CArray::svd casting

henrique-borba · henrique-borba · commit c1212d6b9da3 · 2019-07-08T12:34:14.000-03:00
diff --git a/kernel/linalg.c b/kernel/linalg.c
@@ -614,17 +614,28 @@ CArray_Vdot(CArray * a, CArray * b, MemoryPointer * out)
 CArray **
 CArray_Svd(CArray * a, int full_matrices, int compute_uv, MemoryPointer * out)
 {
-    int m, n;
+    int m, n, casted = 0;
     int lda, ldu, ldvt, info, lwork;
     int * iwork;
     double * s, * u, * vt;
-    CArray * u_ca, * s_ca, *vh_ca, ** rtn;
+    CArray * u_ca, * s_ca, *vh_ca, ** rtn, * target;
 
     if (CArray_NDIM(a) != 2) {
         throw_valueerror_exception("Expected 2D array");
         return NULL;
     }
 
+    if (CArray_DESCR(a)->type_num != TYPE_DOUBLE_INT) {
+        CArrayDescriptor *descr = CArray_DescrFromType(TYPE_DOUBLE_INT);
+        target = CArray_NewLikeArray(a, CARRAY_CORDER, descr, 0);
+        if(CArray_CastTo(target, a) < 0) {
+            return NULL;
+        }
+        casted = 1;
+    } else {
+        target = a;
+    }
+
     m = CArray_DIMS(a)[0];
     n = CArray_DIMS(a)[1];
 
@@ -637,33 +648,33 @@ CArray_Svd(CArray * a, int full_matrices, int compute_uv, MemoryPointer * out)
     u = emalloc(sizeof(double) * (ldu * m));
     vt = emalloc(sizeof(double) * (ldvt * n));
 
-    info = LAPACKE_dgesdd(LAPACK_ROW_MAJOR, 'S', m, n, DDATA(a), lda, s, u, ldu, vt, ldvt);
+    info = LAPACKE_dgesdd(LAPACK_ROW_MAJOR, 'S', m, n, DDATA(target), lda, s, u, ldu, vt, ldvt);
 
     if( info > 0 ) {
         throw_valueerror_exception( "The algorithm computing SVD failed to converge." );
         return NULL;
     }
 
     u_ca = emalloc(sizeof(CArray));
-    u_ca = CArray_NewFromDescr_int(u_ca, CArray_DESCR(a), CArray_NDIM(a), CArray_DIMS(a), CArray_STRIDES(a),
+    u_ca = CArray_NewFromDescr_int(u_ca, CArray_DESCR(target), CArray_NDIM(target), CArray_DIMS(target), CArray_STRIDES(target),
                                    NULL, 0, NULL, 1, 0);
     efree(u_ca->data);
     u_ca->data = (char *)u;
-    CArrayDescriptor_INCREF(CArray_DESCR(a));
+    CArrayDescriptor_INCREF(CArray_DESCR(target));
 
     s_ca = emalloc(sizeof(CArray));
-    s_ca = CArray_NewFromDescr_int(s_ca, CArray_DESCR(a), 1, &n, NULL,
+    s_ca = CArray_NewFromDescr_int(s_ca, CArray_DESCR(target), 1, &n, NULL,
                                    NULL, 0, NULL, 1, 0);
     efree(s_ca->data);
     s_ca->data = (char *)s;
-    CArrayDescriptor_INCREF(CArray_DESCR(a));
+    CArrayDescriptor_INCREF(CArray_DESCR(target));
 
     vh_ca = emalloc(sizeof(CArray));
-    vh_ca = CArray_NewFromDescr_int(vh_ca, CArray_DESCR(a), CArray_NDIM(a), CArray_DIMS(a), CArray_STRIDES(a),
+    vh_ca = CArray_NewFromDescr_int(vh_ca, CArray_DESCR(target), CArray_NDIM(target), CArray_DIMS(target), CArray_STRIDES(target),
                                     NULL, 0, NULL, 1, 0);
     efree(vh_ca->data);
     vh_ca->data = (char *)vt;
-    CArrayDescriptor_INCREF(CArray_DESCR(a));
+    CArrayDescriptor_INCREF(CArray_DESCR(target));
 
     if (out != NULL) {
         add_to_buffer(&(out[0]), u_ca, sizeof(CArray));
@@ -676,5 +687,10 @@ CArray_Svd(CArray * a, int full_matrices, int compute_uv, MemoryPointer * out)
     rtn[1] = s_ca;
     rtn[2] = vh_ca;
 
+    if (casted) {
+        CArrayDescriptor_DECREF(CArray_DESCR(target));
+        CArray_Free(target);
+    }
+
     return rtn;
 }
