add treatment of out keyword

v923z · Mar 26, 2024 · 386fb78 · 386fb78
1 parent 43efc91
commit 386fb78
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 6 deletions.
diff --git a/code/numpy/create.c b/code/numpy/create.c
@@ -837,7 +837,7 @@ mp_obj_t create_take(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args)
         }
     }
 
-    int32_t axis_len = (int32_t)a->shape[ULAB_MAX_DIMS - 1 - a->ndim + axis];
+    int32_t axis_len = (int32_t)a->shape[ULAB_MAX_DIMS - a->ndim + axis];
     uint8_t mode;
 
     if(args[4].u_obj == mp_const_none) {
@@ -887,6 +887,22 @@ mp_obj_t create_take(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args)
         indices[i] = (size_t)index;
     }
 
+    size_t *shape = m_new0(size_t, ULAB_MAX_DIMS);
+    for(uint8_t i = 0; i < ULAB_MAX_DIMS; i++) {
+        shape[i] = a->shape[i];
+        if(i == (ULAB_MAX_DIMS - a->ndim + axis)) {
+            shape[i] = indices_len;
+        }
+    }
+
+    ndarray_obj_t *out = NULL;
+    if(args[3].u_obj == mp_const_none) {
+        out = ndarray_new_dense_ndarray(a->ndim, shape, a->ndim);
+    } else {
+        // TODO: deal with last argument being false!
+        out = ulab_tools_inspect_out(args[3].u_obj, a->dtype, a->ndim, shape, true);
+    }
+    (void)out;
     m_del(size_t, indices, indices_len);
     return mp_const_none;
 }

diff --git a/code/numpy/random/random.c b/code/numpy/random/random.c
@@ -175,8 +175,6 @@ static mp_obj_t random_normal(size_t n_args, const mp_obj_t *pos_args, mp_map_t
 
     mp_float_t *array = (mp_float_t *)ndarray->array;
 
-    // numpy's random supports only dense output arrays, so we can simply 
-    // loop through the elements in a linear fashion
     for(size_t i = 0; i < ndarray->len; i = i + 2) {
         #if MICROPY_FLOAT_IMPL == MICROPY_FLOAT_IMPL_FLOAT
         uint32_t x = pcg32_next(&self->state);

diff --git a/code/ulab_tools.c b/code/ulab_tools.c
@@ -274,3 +274,31 @@ bool ulab_tools_mp_obj_is_scalar(mp_obj_t obj) {
     }
     #endif
 }
+
+ndarray_obj_t *ulab_tools_inspect_out(mp_obj_t out, uint8_t dtype, uint8_t ndim, size_t *shape, bool dense_only) {
+    if(!mp_obj_is_type(out, &ulab_ndarray_type)) {
+        mp_raise_TypeError(MP_ERROR_TEXT("out has wrong type"));
+    }
+    ndarray_obj_t *ndarray = MP_OBJ_TO_PTR(out);
+
+    if(ndarray->dtype != dtype) {
+        mp_raise_ValueError(MP_ERROR_TEXT("out array has wrong dtype"));
+    }
+
+    if(ndarray->ndim != ndim) {
+        mp_raise_ValueError(MP_ERROR_TEXT("out array has wrong dimension"));
+    }
+
+    for(uint8_t i = 0; i < ULAB_MAX_DIMS; i++) {
+        if(ndarray->shape[i] != shape[i]) {
+            mp_raise_ValueError(MP_ERROR_TEXT("out array has wrong shape"));
+        }
+    }
+
+    if(dense_only) {
+        if(!ndarray_is_dense(ndarray)) {
+            mp_raise_ValueError(MP_ERROR_TEXT("output array must be contiguous"));
+        }
+    }
+    return ndarray;
+}
diff --git a/code/ulab_tools.h b/code/ulab_tools.h
@@ -44,7 +44,6 @@ void ulab_rescale_float_strides(int32_t *);
 
 bool ulab_tools_mp_obj_is_scalar(mp_obj_t );
 
-#if ULAB_NUMPY_HAS_RANDOM_MODULE
-ndarray_obj_t *ulab_tools_create_out(mp_obj_tuple_t , mp_obj_t , uint8_t , bool );
-#endif
+ndarray_obj_t *ulab_tools_inspect_out(mp_obj_t , uint8_t , uint8_t , size_t *, bool );
+
 #endif