[GitHub] [incubator-mxnet] samskalicky commented on a change in pull request #17569: Adding sparse support to MXTensor for custom operators

[GitBox <gi...@apache.org>]

samskalicky commented on a change in pull request #17569: Adding sparse support to MXTensor for custom operators
URL: https://github.com/apache/incubator-mxnet/pull/17569#discussion_r387401876
 
 

 ##########
 File path: include/mxnet/lib_api.h
 ##########
 @@ -1111,16 +1197,49 @@ extern "C" {
 
     // create a vector of tensors for inputs
     std::vector<MXTensor> inputs(num_in);
+    // create a vector for sparse inputs
+    std::vector<MXInSparse> in_sparse(num_in);
+
     for (int i = 0; i < num_in; i++) {
-      inputs[i].setTensor(indata[i], (MXDType)intypes[i], inshapes[i], indims[i],
-                          inIDs[i], {indev_type[i], indev_id[i]});
+      // Dense representation.
+      if(!in_indices_shapes) {
+        inputs[i].setTensor(indata[i], (MXDType)intypes[i], inshapes[i], indims[i],
+                            inIDs[i], {indev_type[i], indev_id[i]}, kDefaultStorage);
+      }
+      // Sparse representation.
+      else {
+        MXStorageType type;
+        if(!in_indptr_shapes) {
+          type = kRowSparseStorage;
+          in_sparse[i].set(indata[i], inshapes[i], indims[i], in_indices[i], in_indices_shapes[i]);
+        }
+        else {
+          type = kCSRStorage;
+          in_sparse[i].set(indata[i], inshapes[i], indims[i], in_indices[i],
+                           in_indices_shapes[i], in_indptr[i], in_indptr_shapes[i]);
+        }
+        inputs[i].setTensor((void*)(&in_sparse[i]), (MXDType)intypes[i], inshapes[i], indims[i],
+                            inIDs[i], {indev_type[i], indev_id[i]}, type);
+      }
     }
 
     // create a vector of tensors for outputs
     std::vector<MXTensor> outputs(num_out);
+    // create a vector for sparse outputs
+    std::vector<MXOutSparse> out_sparse;
 
 Review comment:
   This assumes we should allocate a sparse tensor for every input, before even checking if the input is sparse. Maybe theres a better way to do it 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services