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: us...@infra.apache.org With regards, Apache Git Services