fix np.diag

Author: v923z

code/numpy/create.c

@@ -367,52 +367,49 @@ mp_obj_t create_diag(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args)
     mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
     mp_arg_parse_all(n_args, pos_args, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
 
-    if(!mp_obj_is_type(args[0].u_obj, &ulab_ndarray_type)) {
-        mp_raise_TypeError(translate("input must be an ndarray"));
-    }
-    ndarray_obj_t *source = MP_OBJ_TO_PTR(args[0].u_obj);
-    if(source->ndim == 1) { // return a rank-2 tensor with the prescribed diagonal
-        ndarray_obj_t *target = ndarray_new_dense_ndarray(2, ndarray_shape_vector(0, 0, source->len, source->len), source->dtype);
+    ndarray_obj_t *source = ndarray_from_iterable(args[0].u_obj, NDARRAY_FLOAT);
+    ndarray_obj_t *target = NULL;
+
+    if(source->ndim == 2) { // return the diagonal
+        size_t len = MIN(source->shape[ULAB_MAX_DIMS - 1], source->shape[ULAB_MAX_DIMS - 2]);
+        target = ndarray_new_linear_array(len, source->dtype);
         uint8_t *sarray = (uint8_t *)source->array;
         uint8_t *tarray = (uint8_t *)target->array;
-        for(size_t i=0; i < source->len; i++) {
+        for(size_t i=0; i < len; i++) {
             memcpy(tarray, sarray, source->itemsize);
-            sarray += source->strides[ULAB_MAX_DIMS - 1];
-            tarray += (source->len + 1) * target->itemsize;
+            sarray += (source->strides[ULAB_MAX_DIMS - 1] + source->strides[ULAB_MAX_DIMS - 2]);
+            tarray += target->itemsize;
         }
-        return MP_OBJ_FROM_PTR(target);
-    }
-    if(source->ndim > 2) {
-        mp_raise_TypeError(translate("input must be a tensor of rank 2"));
-    }
-    int32_t k = args[1].u_int;
-    size_t len = 0;
-    uint8_t *sarray = (uint8_t *)source->array;
-    if(k < 0) { // move the pointer "vertically"
-        if(-k < (int32_t)source->shape[ULAB_MAX_DIMS - 2]) {
-            sarray -= k * source->strides[ULAB_MAX_DIMS - 2];
-            len = MIN(source->shape[ULAB_MAX_DIMS - 2] + k, source->shape[ULAB_MAX_DIMS - 1]);
+    } else if(source->ndim == 1) { // return a rank-2 tensor with the prescribed diagonal
+        int32_t k = args[1].u_int;
+        size_t len = source->len;
+        if(k < 0) {
+            len -= k;
+        } else {
+            len += k;
         }
-    } else { // move the pointer "horizontally"
-        if(k < (int32_t)source->shape[ULAB_MAX_DIMS - 1]) {
-            sarray += k * source->strides[ULAB_MAX_DIMS - 1];
-            len = MIN(source->shape[ULAB_MAX_DIMS - 1] - k, source->shape[ULAB_MAX_DIMS - 2]);
+        target = ndarray_new_dense_ndarray(2, ndarray_shape_vector(0, 0, len, len), source->dtype);
+        uint8_t *sarray = (uint8_t *)source->array;
+        uint8_t *tarray = (uint8_t *)target->array;
+
+        if(k < 0) {
+            k = -k;
+            tarray += len * k * target->itemsize;
+        } else {
+            tarray += k * target->itemsize;
+        }
+        for(size_t i = 0; i < source->len; i++) {
+            memcpy(tarray, sarray, source->itemsize);
+            sarray += source->strides[ULAB_MAX_DIMS - 1];
+            tarray += (len + 1) * target->itemsize;
         }
     }
-
-    if(len == 0) {
-        mp_raise_ValueError(translate("offset is too large"));
+    #if ULAB_MAX_DIMS > 2
+    else {
+        mp_raise_ValueError(translate("input must be 1- or 2-d"));
     }
+    #endif
 
-    ndarray_obj_t *target = ndarray_new_linear_array(len, source->dtype);
-    uint8_t *tarray = (uint8_t *)target->array;
-
-    for(size_t i=0; i < len; i++) {
-        memcpy(tarray, sarray, source->itemsize);
-        sarray += source->strides[ULAB_MAX_DIMS - 2];
-        sarray += source->strides[ULAB_MAX_DIMS - 1];
-        tarray += source->itemsize;
-    }
     return MP_OBJ_FROM_PTR(target);
 }
 

code/ulab.c

@@ -33,7 +33,7 @@
 #include "user/user.h"
 #include "utils/utils.h"
 
-#define ULAB_VERSION 5.0.1
+#define ULAB_VERSION 5.0.2
 #define xstr(s) str(s)
 #define str(s) #s
 

docs/ulab-change-log.md

@@ -1,3 +1,9 @@
+Tue, 8 Feb 2022
+
+version 5.0.2
+
+    fix np.diag
+
 Thu, 3 Feb 2022
 
 version 5.0.1