refactor conversion api (apache#25)

* add multiple stage test

* refactor GetFCompute

* adjust jenkins python version. pass all unit test

* register cast storage type op

* refactor conversion interface

Jenkinsfile

@@ -201,17 +201,17 @@ del /Q *.7z
 // Python unittest for CPU
 def python_ut(docker_type) {
   timeout(time: max_time, unit: 'MINUTES') {
-    sh "${docker_run} ${docker_type} PYTHONPATH=./python/ nosetests --with-timer --verbose tests/python/unittest"
+    sh "${docker_run} ${docker_type} PYTHONPATH=./python/ nosetests-2.7 --with-timer --verbose tests/python/unittest"
     // sh "${docker_run} ${docker_type} PYTHONPATH=./python/ nosetests-3.4 --with-timer --verbose tests/python/unittest"
-    sh "${docker_run} ${docker_type} PYTHONPATH=./python/ nosetests --with-timer --verbose tests/python/train"
+    sh "${docker_run} ${docker_type} PYTHONPATH=./python/ nosetests-2.7 --with-timer --verbose tests/python/train"
   }
 }
 
 // GPU test has two parts. 1) run unittest on GPU, 2) compare the results on
 // both CPU and GPU
 def python_gpu_ut(docker_type) {
   timeout(time: max_time, unit: 'MINUTES') {
-    sh "${docker_run} ${docker_type} PYTHONPATH=./python/ nosetests --with-timer --verbose tests/python/gpu"
+    sh "${docker_run} ${docker_type} PYTHONPATH=./python/ nosetests-2.7 --with-timer --verbose tests/python/gpu"
     // sh "${docker_run} ${docker_type} PYTHONPATH=./python/ nosetests-3.4 --with-timer --verbose tests/python/gpu"
   }
 }

include/mxnet/c_api.h

@@ -272,11 +272,6 @@ MXNET_DLL int MXNDArrayCreateSparseEx(int storage_type,
                     int *aux_types,
                     NDArrayHandle *out);
 
-// TEMPORARY API FOR TESTING PURPOSE. Conversion should be an op instead
-MXNET_DLL int MXNDArrayConvert(NDArrayHandle in,
-                               int storage_type,
-                               NDArrayHandle *out);
-
 /*!
  * \brief create a NDArray handle that is loaded from raw bytes.
  * \param buf the head of the raw bytes

include/mxnet/ndarray.h

@@ -98,7 +98,7 @@ class NDArray {
   }
   /*! \brief constructor for NDArray with chunk type
    */
-  NDArray(NDArrayStorageType storage_type, const TShape &shape, Context ctx,
+  NDArray(const NDArrayStorageType storage_type, const TShape &shape, Context ctx,
           bool delay_alloc = true, int dtype = mshadow::default_type_flag,
           std::vector<int> aux_types = {})
       : shape_(shape), offset_(0), dtype_(dtype), entry_({nullptr, 0, 0}) {
@@ -127,6 +127,7 @@ class NDArray {
       Mkl_mem_ = std::make_shared<MKLMemHolder>();
 #endif
   }
+  // TODO this constructor should be removed
   NDArray(NDArray data, const std::vector<NDArray> aux_data, Context ctx,
           NDArrayStorageType storage_type, const TShape &shape)
       : ptr_(std::make_shared<Chunk>(data, aux_data, ctx, storage_type)), shape_(shape),
@@ -137,11 +138,6 @@ class NDArray {
       CHECK(aux_data.size() == 1) << "Multiple aux_data not supported yet";
   }
 
-  template<typename xpu>
-  NDArray ConvertTo(NDArrayStorageType storage_type, mshadow::Stream<xpu> *s) const {
-    CHECK_EQ(storage_type, kDefaultStorage) << "other storage type not supported yet";
-    return ToDefault<xpu>(s);
-  }
   /*!
    * \return the shape of current NDArray.
    */
@@ -487,44 +483,10 @@ class NDArray {
 
  private:
   friend class autograd::AutogradRuntime;
-  // Make a copy of the ndarray in dense format
-  template<typename xpu>
-  NDArray ToDefault(mshadow::Stream<xpu>* s) const {
-    NDArray result(shape_, ptr_->ctx, false, dtype());
-    this->WaitToRead();
-    if (storage_type() == kDefaultStorage) {
-      MSHADOW_TYPE_SWITCH(dtype(), DType, {
-        mshadow::Copy(result.data().FlatTo1D<xpu, DType>(), data().FlatTo1D<xpu, DType>());
-      });
-      return result;
-    }
-    CHECK(storage_type() == kRowSparseStorage);
-    MSHADOW_TYPE_SWITCH(dtype(), DType, {
-      MSHADOW_TYPE_SWITCH(aux_type(rowsparse::kIdx), AuxType, {
-        // Fill in zeros
-        result.data().FlatTo1D<xpu, DType>(s) = 0;
-        result.data().shape_ = shape_;
-        // data() is not empty
-        if (storage_shape().ndim() != 0) {
-          // Copy over
-          auto in_data = data().FlatTo2D<xpu, DType>(s);
-          auto out_data = result.data().FlatTo2D<xpu, DType>(s);
-          auto num_rows = aux_shape(rowsparse::kIdx)[0];
-          auto in_idx = aux_data(rowsparse::kIdx).FlatTo1D<xpu, AuxType>(s);
-          for (size_t i = 0; i < num_rows; i += 1) {
-            mshadow::Copy(out_data[in_idx[i]], in_data[i], s);
-          }
-        }
-      });
-    });
-    return result;
-  }
-
   /*! \brief the real data chunk that backs NDArray */
   // shandle is used to store the actual values in the NDArray
   // aux_handles store the aux data(such as indices) if it's needed by non-default storage.
   struct Chunk {
-    // TODO(haibin) Also specify the capacity & size of the chunk, we don't want to resize it
     // every time a new element is added to a non default storage
     /*! \brief storage handle from storage engine.
                for non-default storage, shandle stores the data(value) array.
@@ -551,7 +513,7 @@ class NDArray {
     // context of data
     Context ctx;
     // The shape of the chunk data.
-    // This might not be the same shape as the NDArray, since the chunk may be sparse.
+    // This might not be the same shape as the NDArray, since the storage may be sparse.
     TShape storage_shape;
     // The shape of aux data. The default value for the shape is 0.
     std::vector<TShape> aux_shapes;
@@ -660,7 +622,7 @@ class NDArray {
       CHECK_EQ(storage_type, kRowSparseStorage) << "Not yet implemented";
       // calculate size, perform allocation
       if (delay_alloc) {
-        // For row sparse chunk, aux_shape indicates the number of rows to allocate
+        // For row sparse storage, aux_shape indicates the number of rows to allocate
         auto aux_shape = aux_shapes[0];
         CHECK_EQ(aux_shape.ndim(), 1);
         auto num_rows = aux_shape[0];
@@ -670,7 +632,7 @@ class NDArray {
         shandle = Storage::Get()->Alloc(dbytes, ctx);
         aux_handles.push_back(Storage::Get()->Alloc(aux_bytes, ctx));
         delay_alloc = false;
-        // Initialize aux_shape and shape
+        // Initialize shapes
         this->aux_shapes = aux_shapes;
         storage_shape = shape;
         storage_shape[0] = num_rows;

include/mxnet/op_attr_types.h

@@ -66,8 +66,6 @@ using FCompute = std::function<void (const nnvm::NodeAttrs& attrs,
  *
  * \note Register under "FComputeEx<cpu, `storage_type`>" and "FComputeEx<gpu, `storage_type`>" 
  * e.g FComputeEx<cpu, row_sparse>
- * TODO should probably change const std::vector<NDArray>& outputs to
-                               std::vector<NDArray> *outputs
  */
 using FComputeEx = std::function<void (const nnvm::NodeAttrs& attrs,
                                      const OpContext& ctx,

python/mxnet/sparse_ndarray.py

@@ -92,11 +92,7 @@ class SparseNDArray(NDArray):
     ''' sparse ndarray '''
     __slots__ = []
 
-    def __repr__(self):
-        """Return a string representation of the array"""
-        # TODO(haibin) also print shape info
-        #shape_info = 'x'.join(['%d' % x for x in self.shape])
-        return '<%s>' % (self.__class__.__name__)
+    #def __repr__(self):
     def __reduce__(self):
         return (SparseNDArray, (None,), self.__getstate__())
     def __add__(self, other):
@@ -160,26 +156,22 @@ def _sync_copyfrom(self, source_array):
     def _slice(self, start, stop):
         raise Exception('Not implemented for SparseND yet!')
     def _at(self, idx):
-        raise Exception('Not implemented for SparseND yet!')
+        raise Exception('at operator for SparseND is not supported.')
     def reshape(self, shape):
         raise Exception('Not implemented for SparseND yet!')
     def broadcast_to(self, shape):
         raise Exception('Not implemented for SparseND yet!')
     #def wait_to_read(self):
-    #inherited from parent
     #@property
     #def shape(self):
-    #inherited from parent
 
     @property
     def size(self):
         raise Exception('Not implemented for SparseND yet!')
-    @property
-    def context(self):
-        raise Exception('Not implemented for SparseND yet!')
-    @property
-    def dtype(self):
-        raise Exception('Not implemented for SparseND yet!')
+    #@property
+    #def context(self):
+    #@property
+    #def dtype(self):
     @property
     # pylint: disable= invalid-name, undefined-variable
     def T(self):
@@ -202,16 +194,16 @@ def as_in_context(self, context):
     def to_dense(self):
         return to_dense(self)
 
-#TODO(haibin) also add aux_types. Not tested yet.
-#We need a to_dense method to test it
-def csr(values, indptr, idx, shape, ctx=Context.default_ctx, dtype=mx_real_t):
+#TODO We need a to_dense method to test it
+def csr(values, indptr, idx, shape, ctx=Context.default_ctx, dtype=mx_real_t, aux_types=None):
     ''' constructor '''
     hdl = NDArrayHandle()
     #TODO currently only supports NDArray input
     assert(isinstance(values, NDArray))
     assert(isinstance(index, NDArray))
     indices = c_array(NDArrayHandle, [idx.handle, indptr.handle])
     num_aux = mx_uint(2)
+    # TODO create an empty handle with specified types, then assign values
     check_call(_LIB.MXNDArrayCreateSparse(
         values.handle, num_aux, indices,
         c_array(mx_uint, shape),
@@ -226,13 +218,14 @@ def csr(values, indptr, idx, shape, ctx=Context.default_ctx, dtype=mx_real_t):
 
 # pylint: enable= no-member
 #TODO(haibin) also specify aux_types
-def row_sparse(values, index, shape, ctx=Context.default_ctx, dtype=mx_real_t):
+def row_sparse(values, index, shape, ctx=Context.default_ctx, dtype=mx_real_t, aux_types=None):
     ''' constructor '''
     hdl = NDArrayHandle()
     assert(isinstance(values, NDArray))
     assert(isinstance(index, NDArray))
     indices = c_array(NDArrayHandle, [index.handle])
     num_aux = mx_uint(1)
+    # TODO create an empty handle with specified types, then assign values
     check_call(_LIB.MXNDArrayCreateSparse(
         values.handle, num_aux, indices,
         c_array(mx_uint, shape),
@@ -245,26 +238,27 @@ def row_sparse(values, index, shape, ctx=Context.default_ctx, dtype=mx_real_t):
         ctypes.byref(hdl)))
     return SparseNDArray(hdl)
 
-def array(values, index_list, storage_type, shape, ctx=None, dtype=mx_real_t):
+def array(values, index_list, storage_type, shape, ctx=None, dtype=mx_real_t, aux_types=None):
     # TODO check input array types. Assume NDArray class for now
     # TODO support other types
     assert(storage_type == 'row_sparse')
     if isinstance(shape, int):
         shape = (shape, )
     if ctx is None:
         ctx = Context.default_ctx
-    arr = row_sparse(values, index_list[0], shape, ctx=ctx, dtype=dtype)
+    arr = row_sparse(values, index_list[0], shape, ctx=ctx, dtype=dtype, aux_types=aux_types)
     return arr
 
 # Temporary function for testing purpose
 def to_dense(source):
-    hdl = NDArrayHandle()
+    return ndarray.cast_storage(source, storage_type=1)
+    '''hdl = NDArrayHandle()
     check_call(_LIB.MXNDArrayConvert(
         source.handle, _STORAGE_TYPE_STR_TO_ID['default'],
         ctypes.byref(hdl)))
-    return ndarray.NDArray(handle=hdl, writable=True)
+    return ndarray.NDArray(handle=hdl, writable=True)'''
 
-def zeros(shape, storage_type, ctx=None, dtype=mx_real_t):
+def zeros(shape, storage_type, ctx=None, dtype=mx_real_t, aux_types=None):
     """Return a new array of given shape and type, filled with zeros.
 
     Parameters
@@ -294,12 +288,13 @@ def zeros(shape, storage_type, ctx=None, dtype=mx_real_t):
     """
     if ctx is None:
         ctx = Context.default_ctx
-    if storage_type == 'row_sparse':
-        # pylint: disable= no-member, protected-access
-        out = SparseNDArray(_new_alloc_handle(storage_type, shape, ctx,
-                                              aux_types=_STORAGE_AUX_TYPES['row_sparse']))
-        return _internal._zeros(shape=shape, ctx=ctx, dtype=dtype, out=out)
-    return _internal._zeros(shape=shape, ctx=ctx, dtype=dtype)
+    assert(storage_type == 'row_sparse')
+    if aux_types == None:
+        aux_types = _STORAGE_AUX_TYPES['row_sparse']
+    # pylint: disable= no-member, protected-access
+    out = SparseNDArray(_new_alloc_handle(storage_type, shape, ctx,
+                                          aux_types=aux_types))
+    return _internal._zeros(shape=shape, ctx=ctx, dtype=dtype, out=out)
     # pylint: enable= no-member, protected-access
 
 _STORAGE_TYPE_TO_ND_CLASS = {

src/c_api/c_api.cc

@@ -131,6 +131,7 @@ int MXNDArrayCreate(const mx_uint *shape,
   API_END();
 }
 
+// TODO remove this API
 int MXNDArrayCreateSparse(NDArrayHandle data,
                     mx_uint num_aux,
                     NDArrayHandle *aux_vec,
@@ -155,16 +156,6 @@ int MXNDArrayCreateSparse(NDArrayHandle data,
   API_END();
 }
 
-// TODO(haibin) Should also consider context
-int MXNDArrayConvert(NDArrayHandle in,
-                     int storage_type,
-                     NDArrayHandle *out) {
-  API_BEGIN();
-  NDArray* nd = reinterpret_cast<NDArray*>(in);
-  *out = new NDArray(nd->ConvertTo<cpu>(static_cast<NDArrayStorageType>(storage_type), nullptr));
-  API_END();
-}
-
 int MXNDArrayCreateEx(const mx_uint *shape,
                     mx_uint ndim,
                     int dev_type,
@@ -363,7 +354,6 @@ int MXNDArrayGetStorageType(NDArrayHandle handle,
                      int *out_storage_type) {
   API_BEGIN();
   NDArray *arr = static_cast<NDArray*>(handle);
-  // Check is_none?
   if (!arr->is_none()) {
     *out_storage_type = arr->storage_type();
   } else {

src/c_api/c_api_common.h

@@ -58,7 +58,7 @@ struct MXAPIThreadLocalEntry {
   std::vector<TShape> arg_shapes, out_shapes, aux_shapes;
   /*! \brief result holder for returning type flags */
   std::vector<int> arg_types, out_types, aux_types;
-  /*! \brief result holder for returning chunk types */
+  /*! \brief result holder for returning storage types */
   std::vector<int> arg_storage_types, out_storage_types, aux_storage_types;
   /*! \brief result holder for returning shape dimensions */
   std::vector<mx_uint> arg_shape_ndim, out_shape_ndim, aux_shape_ndim;