Code changes based on reviews (apache#144)

* code changes according to review comments

remove executor debug. add doc to optimizer

update sparse sgd test

add dtype option to rand_sparse_ndarray

* overhauled reqs for sparse operators

* patch FCompExFallback with mutable inputs. update test_optimizer with more fallback cases

* change executor debug macro to env var

* add comment

* update doc

* change ndarray.aux_shape() to return const reference

* remove todense to_rsp to_csr. replace with tostype

* replace manual calls to cast_storage with tostype

* disable gpu fallback test for optimizer

* fix lint

* add backward pass for cast_storage. refactor cast_storage test

* rand_sparse_ndarray bug fix

* fix cast_storage for gpu

* disable csr test for fp16

* update row sparse ndarray doc

* update doc

benchmark/python/cast_storage.py

@@ -26,7 +26,7 @@ def run_cast_storage_synthetic():
     def dense_to_sparse(m, n, density, ctx, repeat, stype):
         set_default_context(ctx)
         data_shape = (m, n)
-        dns_data = rand_ndarray(data_shape, stype, density).todense()
+        dns_data = rand_ndarray(data_shape, stype, density).tostype('default')
         dns_data.wait_to_read()
 
         # do one warm up run, verify correctness

benchmark/python/sparse_op.py

@@ -92,7 +92,7 @@ def get_iter(path, data_shape, batch_size):
         for batch in train_iter:
             data = train_iter.getdata()
             csr_data.append(data)
-            dns_data.append(data.todense())
+            dns_data.append(data.tostype('default'))
             num_batch += 1
         bag_of_data = [csr_data, dns_data]
         num_repeat = 5
@@ -140,10 +140,10 @@ def bench_dot_forward(m, k, n, density, ctx, repeat):
         dns = mx.nd.random_uniform(shape=(k, n)).copyto(ctx)
         data_shape = (m, k)
         csr_data = rand_ndarray(data_shape, 'csr', density)
-        dns_data = csr_data.todense()
+        dns_data = csr_data.tostype('default')
         rhs_dns_np = dns.asnumpy()
         lhs_csr_sp = sp.csr_matrix(dns_data.asnumpy())  # csr in scipy
-        lhs_dns_np = lhs_csr_sp.todense()
+        lhs_dns_np = lhs_csr_sp.tostype('default')
 
         data = [dns_data, csr_data]
         costs = []
@@ -169,10 +169,10 @@ def bench_dot_backward(m, k, n, density, ctx, repeat):
         dns = mx.nd.random_uniform(shape=(m, n)).copyto(ctx)
         data_shape = (m, k)
         csr_data = rand_ndarray(data_shape, 'csr', density)
-        dns_data = csr_data.todense()
+        dns_data = csr_data.tostype('default')
         rhs_dns_np = dns.asnumpy()
         lhs_csr_sp = sp.csr_matrix(dns_data.asnumpy())
-        lhs_dns_np = lhs_csr_sp.todense()
+        lhs_dns_np = lhs_csr_sp.tostype('default')
 
         data = [dns_data, csr_data]
         costs = []

docs/api/python/ndarray.md

@@ -107,6 +107,7 @@ We summarize the interface for each class in the following sections.
     NDArray.asnumpy
     NDArray.asscalar
     NDArray.astype
+    NDArray.tostype
 ```
 
 ### Array change shape
@@ -191,6 +192,7 @@ We summarize the interface for each class in the following sections.
     :nosignatures:
 
     RowSparseNDArray.copyto
+    RowSparseNDArray.tostype
     RowSparseNDArray.__setitem__
     RowSparseNDArray.__getitem__
     RowSparseNDArray.data
@@ -204,6 +206,7 @@ We summarize the interface for each class in the following sections.
     :nosignatures:
 
     CSRNDArray.copyto
+    CSRNDArray.tostype
     CSRNDArray.__setitem__
     CSRNDArray.__getitem__
     CSRNDArray.data

include/mxnet/c_api.h

@@ -719,6 +719,30 @@ MXNET_DLL int MXCreateCachedOp(SymbolHandle handle,
  * \brief free cached operator
  */
 MXNET_DLL int MXFreeCachedOp(CachedOpHandle handle);
+/*!
+ * \brief invoke cached operator
+ */
+MXNET_DLL int MXInvokeCachedOp(CachedOpHandle handle,
+                               int num_inputs,
+                               NDArrayHandle *inputs,
+                               int *num_outputs,
+                               NDArrayHandle **outputs);
+/*!
+ * \brief invoke a cached op
+ * \param handle the handle to the cached op
+ * \param num_inputs number of input NDArrays
+ * \param inputs input NDArrays
+ * \param num_outputs number of output NDArrays
+ * \param outputs output NDArrays
+ * \param out_stypes output ndarrays' stypes
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXInvokeCachedOpEx(CachedOpHandle handle,
+                                 int num_inputs,
+                                 NDArrayHandle *inputs,
+                                 int *num_outputs,
+                                 NDArrayHandle **outputs,
+                                 const int** out_stypes);
 /*!
  * \brief invoke cached operator
  */

include/mxnet/ndarray.h

@@ -193,13 +193,16 @@ class NDArray {
   }
 
   /*!
-   * \return the shape of aux data at ith index. If it doesn't exist, return an empty one.
+   * \brief get the shape of aux_data(index)
+   * \param index the index of the aux data
+   * \return the shape of aux data at given index
    */
-  inline const TShape aux_shape(size_t i) const {
+  inline const TShape& aux_shape(size_t index) const {
     CHECK(storage_type() != kDefaultStorage);
-    return ptr_->aux_shapes[i];
+    return ptr_->aux_shapes[index];
   }
 
+  /* \return the shapes of all aux data */
   const std::vector<TShape>& aux_shapes() const {
     CHECK(storage_type() != kDefaultStorage);
     return ptr_->aux_shapes;
@@ -212,8 +215,8 @@ class NDArray {
    * for the final result. After the operation is done, the exact size of
    * the shape is known and need to be reset using this function.
    */
-  inline void set_aux_shape(size_t i, const TShape& shape) const {
-    ptr_->set_aux_shape(i, shape);
+  inline void set_aux_shape(size_t index, const TShape& shape) const {
+    ptr_->set_aux_shape(index, shape);
   }
 
   /*!
@@ -851,7 +854,6 @@ size_t num_aux_data(NDArrayStorageType stype);
  * \param from the ndarray we want to copy data from
  * \param to the target ndarray
  * \param priority Priority of the action.
- * \param alloc_output whether to allocate memory for the output ndarray
  * \note The function name explicitly marks the order of from and to
  *     due to different possible convention carried by copy function.
  */

python/mxnet/_ctypes/ndarray.py

@@ -127,17 +127,24 @@ def __call__(self, *args, **kwargs):
                 "CachedOp.__call__ got unexpected keyword argument(s): " + \
                 ', '.join(kwargs.keys()))
 
-        check_call(_LIB.MXInvokeCachedOp(
+        # return output stypes to avoid the c_api call for checking
+        # a handle's stype in _ndarray_cls
+        out_stypes = ctypes.POINTER(ctypes.c_int)()
+
+        check_call(_LIB.MXInvokeCachedOpEx(
             self.handle,
             ctypes.c_int(len(args)),
             c_array(NDArrayHandle, [arr.handle for arr in args]),
             ctypes.byref(num_output),
-            ctypes.byref(output_vars)))
+            ctypes.byref(output_vars),
+            ctypes.byref(out_stypes)))
 
         if original_output is not None:
             return original_output
         if num_output.value == 1:
-            return _ndarray_cls(ctypes.cast(output_vars[0], NDArrayHandle))
+            return _ndarray_cls(ctypes.cast(output_vars[0], NDArrayHandle),
+                                stype=_STORAGE_TYPE_ID_TO_STR[out_stypes[0]])
         else:
-            return [_ndarray_cls(ctypes.cast(output_vars[i], NDArrayHandle))
+            return [_ndarray_cls(ctypes.cast(output_vars[i], NDArrayHandle),
+                                 stype=_STORAGE_TYPE_ID_TO_STR[out_stypes[i]])
                     for i in range(num_output.value)]

python/mxnet/kvstore.py

@@ -255,11 +255,12 @@ def row_sparse_pull(self, key, out=None, priority=0, row_ids=None):
         row_ids : NDArray or list of NDArray
             The row_ids for which to pull for each value. The row_ids doesn't have to be unique
             or sorted.
+
         Examples
         --------
         >>> shape = (3, 3)
-        >>> kv.init('3', mx.nd.ones(shape)._to_rsp())
-        >>> a = mx.nd.zeros(shape)
+        >>> kv.init('3', mx.nd.ones(shape).tostype('row_sparse'))
+        >>> a = mx.nd.zeros(shape, stype='row_sparse')
         >>> row_ids = mx.nd.array([0, 2], dtype='int64')
         >>> kv.row_sparse_pull('3', out=a, row_ids=row_ids)
         >>> print a.asnumpy()

python/mxnet/ndarray/__init__.py

@@ -5,6 +5,6 @@
 from .op import CachedOp
 # pylint: disable=wildcard-import
 from .ndarray import *
-from .ndarray_utils import load, save, zeros, empty, array
-from .sparse_ndarray import _ndarray_cls, todense
-from .sparse_ndarray import csr, row_sparse, BaseSparseNDArray, RowSparseNDArray, CSRNDArray
+from .utils import load, save, zeros, empty, array
+from .sparse_ndarray import _ndarray_cls, csr_matrix, row_sparse_array
+from .sparse_ndarray import BaseSparseNDArray, RowSparseNDArray, CSRNDArray

python/mxnet/ndarray/ndarray.py

@@ -1089,13 +1089,15 @@ def backward(self, out_grad=None, retain_graph=False, is_train=True):
             ctypes.c_int(retain_graph),
             ctypes.c_int(is_train)))
 
-    def _to_csr(self):
-        # pylint: disable=undefined-variable
-        return cast_storage(self, stype='csr')
+    def tostype(self, stype):
+        """Return a copy of the array with chosen storage type.
 
-    def _to_rsp(self):
-        # pylint: disable=undefined-variable
-        return cast_storage(self, stype='row_sparse')
+        Returns
+        -------
+        NDArray, CSRNDArray or RowSparseNDArray
+            A copy of the array with the chosen storage stype
+        """
+        return cast_storage(self, stype=stype)
 
 def onehot_encode(indices, out):
     """One-hot encoding indices into matrix out.