support leading dimension of -1 in ravel/unravel

src/operator/tensor/ravel.cc

@@ -31,12 +31,13 @@ DMLC_REGISTER_PARAMETER(RavelParam);
 
 NNVM_REGISTER_OP(_ravel_multi_index)
 .add_alias("ravel_multi_index")
-.describe(R"code(Converts a batch of index arrays into an array of flat indices. The operator follows numpy conventions so a single multi index is given by a column of the input matrix. 
+.describe(R"code(Converts a batch of index arrays into an array of flat indices. The operator follows numpy conventions so a single multi index is given by a column of the input matrix. The leading dimension may be left unspecified by using -1 as placeholder.  
 
 Examples::
 
    A = [[3,6,6],[4,5,1]]
    ravel(A, shape=(7,6)) = [22,41,37]
+   ravel(A, shape=(-1,6)) = [22,41,37]
 
 )code" ADD_FILELINE)
 .set_num_inputs(1)
@@ -55,12 +56,13 @@ Examples::
 
 NNVM_REGISTER_OP(_unravel_index)
 .add_alias("unravel_index")
-.describe(R"code(Converts an array of flat indices into a batch of index arrays. The operator follows numpy conventions so a single multi index is given by a column of the output matrix.
+.describe(R"code(Converts an array of flat indices into a batch of index arrays. The operator follows numpy conventions so a single multi index is given by a column of the output matrix. The leading dimension may be left unspecified by using -1 as placeholder.  
 
 Examples::
 
    A = [22,41,37]
    unravel(A, shape=(7,6)) = [[3,6,6],[4,5,1]]
+   unravel(A, shape=(-1,6)) = [[3,6,6],[4,5,1]]
 
 )code" ADD_FILELINE)
 .set_num_inputs(1)

src/operator/tensor/ravel.h

@@ -110,11 +110,12 @@ struct unravel_index {
                                   DType *unravelled, DType *ravelled) {
     index_t idx(ravelled[i]);
     #pragma unroll
-    for (int j = ndim; j--; ) {
+    for (int j = ndim-1; j > 0; --j) {
       index_t tmp = idx / shape[j];
       unravelled[i+j*N] = idx - tmp*shape[j];
       idx = tmp;
     }
+    unravelled[i] = idx;
   }
 };
 

tests/python/unittest/test_operator.py

@@ -7106,6 +7106,13 @@ def test_ravel():
       check_symbolic_forward(b, location={'a': data}, expected=[ravel_npy])
       c = mx.sym.unravel_index(a, shape=shape)
       check_symbolic_forward(c, location={'a': ravel_npy}, expected=[data])
+      # Test with leading dimension set to -1.
+      shape2 = shape
+      shape2 = (-1,)+shape[1:]
+      b = mx.sym.ravel_multi_index(a, shape=shape2)
+      check_symbolic_forward(b, location={'a': data}, expected=[ravel_npy])
+      c = mx.sym.unravel_index(a, shape=shape2)
+      check_symbolic_forward(c, location={'a': ravel_npy}, expected=[data])
 
 def test_context_num_gpus():
     try: