Modify ndarray slice to have numpy compatbile behaviou

Author: Mike Mao

src/operator/tensor/matrix_op-inl.h

@@ -685,13 +685,13 @@ inline void GetIndexRange(const mxnet::TShape& dshape,
     << "Static array size=" << ndim
     << " is not equal to data shape ndim=" << dshape.ndim();
 
-  if (param_step.ndim() != 0) {
+  if (param_step.ndim() > 0) {
     CHECK_EQ(param_step.ndim(), param_begin.ndim())
       << "step and begin must have the same length";
   }
 
   for (int i = 0; i < param_begin.ndim(); ++i) {
-    index_t s = param_step.ndim() != 0U && param_step[i].has_value() ? param_step[i].value() : 1;
+    index_t s = param_step.ndim() > 0 && param_step[i].has_value() ? param_step[i].value() : 1;
     CHECK_NE(s, 0) << "slice op step[" << i << "] cannot be 0";
 
     index_t b = 0, e = 0;
@@ -703,29 +703,44 @@ inline void GetIndexRange(const mxnet::TShape& dshape,
       // checking upper and lower bounds for begin
       if (b < 0) {
         b += len;
-        CHECK_GE(b, 0) << "slicing with begin[" << i << "]=" << b - len
-                       << " exceeds limit of input dimension[" << i << "]=" << len;
+        if (!Imperative::Get()->is_np_shape()) {
+          CHECK_GE(b, 0) << "slicing with begin[" << i << "]=" << b - len
+                         << " exceeds limit of input dimension[" << i << "]=" << len;
+        }
+      }
+      if (!Imperative::Get()->is_np_shape()) {
+        CHECK_LT(b, len) << "slicing with begin[" << i << "]=" << b
+                         << " exceeds limit of input dimension[" << i << "]=" << len;
       }
-      CHECK_LT(b, len) << "slicing with begin[" << i << "]=" << b
-                       << " exceeds limit of input dimension[" << i << "]=" << len;
-
       // checking upper and lower bounds for end
       if (e < 0 && param_end[i].has_value()) {
-        if (!(s < 0 && e == -1)) {
-          // Keep end=-1 as one-beyond-limits index for negative stride
-          e += len;
+        e += len;
+        if (!Imperative::Get()->is_np_shape()) {
+          CHECK_GE(e, 0) << "slicing with end[" << i << "]=" << e - len
+                         << " exceeds limit of input dimension[" << i << "]=" << len;
         }
-        CHECK_GE(e, 0) << "slicing with end[" << i << "]=" << e - len
-                       << " exceeds limit of input dimension[" << i << "]=" << len;
       }
-      CHECK_LE(e, len) << "slicing with end[" << i << "]=" << e
-                       << " exceeds limit of input dimension[" << i << "]=" << len;
+      if (!Imperative::Get()->is_np_shape()) {
+        CHECK_LE(e, len) << "slicing with end[" << i << "]=" << e
+                         << " exceeds limit of input dimension[" << i << "]=" << len;
+      }
 
       // checking begin==end case which is not supported
-      CHECK_NE(b, e) << "slicing with begin[" << i << "]=end[" << i << "]="
-                     << e << " results in an empty tensor and is not supported";
+      if (!Imperative::Get()->is_np_shape()) {
+        CHECK_NE(b, e) << "slicing with begin[" << i << "]=end[" << i << "]="
+                       << e << " results in an empty tensor and is not supported";
+      }
     }
 
+    if (Imperative::Get()->is_np_shape()) {
+      // move the begin and end to correct position for calculating dim size
+      b = b < 0 && s > 0 ? 0 : b;
+      b = b > len-1 && s < 0 ? len-1 : b;
+      // if the start value lead to empty tensor under step s, use -1 for indication
+      b = b < 0 || b > len-1 ? -1 : b;
+      e = e > -1 ? e : -1;
+      e = e > len ? len : e;
+    }
     (*begin)[i] = b;
     (*end)[i] = e;
     (*step)[i] = s;
@@ -741,17 +756,29 @@ inline void GetIndexRange(const mxnet::TShape& dshape,
 inline void SetSliceOpOutputDimSize(const index_t i, const int b,
                                     const int e, const int s,
                                     mxnet::TShape* oshape) {
-  if (e != b) {
-    if (s > 0) {
-      CHECK_LT(b, e) << "slicing with begin[" << i << "]=" << b << ", end[" << i << "]="
-                     << e << ", and step[" << i << "]=" << s << " is invalid";
-      (*oshape)[i] = (e - b - 1) / s + 1;
+  if (!Imperative::Get()->is_np_shape()) { //handle as ndarray
+    if (e != b) {
+      if (s > 0) {
+        CHECK_LT(b, e) << "slicing with begin=[" << i << "]=" << b << ", end[" << i << "]="
+                      << e << ", and step[" << i << "]=" << s << " is invalid";
+        (*oshape)[i] = (e - b - 1) / s + 1;
+      } else {
+        CHECK_LT(e, b) << "slicing with begin=[" << i << "]=" << b << ", end[" << i << "]="
+                      << e << ", and step[" << i << "]=" << s << " is invalid";
+        (*oshape)[i] = (b - e - 1) / (-s) + 1;
+      }
+    }  // else leave oshape[i] as 0 for partial infer
+  } else { //handle as numpy compatible array
+    if (e != b && b >= 0) {
+      if (s > 0) {
+        (*oshape)[i] = e > b ? (e - b - 1) / s + 1 : 0;
+      } else {
+        (*oshape)[i] = e < b ? (b - e - 1) / (-s) + 1 : 0;
+      }
     } else {
-      CHECK_LT(e, b) << "slicing with begin[" << i << "]=" << b << ", end[" << i << "]="
-                     << e << ", and step[" << i << "]=" << s << " is invalid";
-      (*oshape)[i] = (b - e - 1) / (-s) + 1;
+        (*oshape)[i] = 0;
     }
-  }  // else leave oshape[i] as 0 for partial infer
+  }
 }
 
 inline bool SliceOpShape(const nnvm::NodeAttrs& attrs,
@@ -852,6 +879,7 @@ void SliceOpForward(const nnvm::NodeAttrs& attrs,
   Stream<xpu>* s = ctx.get_stream<xpu>();
   const TBlob& data = inputs[0];
   const TBlob& out = outputs[0];
+  if (Imperative::Get()->is_np_shape() && out.Size() == 0) return;
   const SliceParam& param = nnvm::get<SliceParam>(attrs.parsed);
   MXNET_NDIM_SWITCH(data.ndim(), ndim, {
     common::StaticArray<index_t, ndim> begin, end, step;
@@ -951,6 +979,7 @@ void SliceOpBackward(const nnvm::NodeAttrs& attrs,
   } else if (req[0] == kWriteInplace) {
     LOG(FATAL) << "_slice_backward does not support kWriteInplace";
   }
+  if (Imperative::Get()->is_np_shape() && ograd.Size() == 0) return;
   MXNET_NDIM_SWITCH(ograd.ndim(), ndim, {
     common::StaticArray<index_t, ndim> begin, end, step;
     GetIndexRange(igrad.shape_, param.begin, param.end, param.step, &begin, &end, &step);

src/operator/tensor/matrix_op.cc

@@ -506,6 +506,7 @@ Example::
                                                             [5.,  7.],
                                                             [1.,  3.]]
 )code" ADD_FILELINE)
+.add_alias("_npx_slice")
 .set_attr_parser(ParamParser<SliceParam>)
 .set_attr<mxnet::FInferShape>("FInferShape", SliceOpShape)
 .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>)

tests/python/unittest/test_numpy_op.py

@@ -92,6 +92,78 @@ def is_int(dtype):
                         assert_almost_equal(mx_out.asnumpy(), np_out, rtol=1e-3, atol=1e-5)
 
 
+@with_seed()
+@use_np
+def test_npx_slice():
+    class TestSlice(HybridBlock):
+        def __init__(self, begin, end, step):
+            super(TestSlice, self).__init__()
+            self._begin = begin
+            self._end = end
+            self._step = step
+
+        def hybrid_forward(self, F, a, *args, **kwargs):
+            return F.npx.slice(a, begin=self._begin, end=self._end, step=self._step)
+
+    def get_start_end_step(shape):
+        start = []
+        end = []
+        step_switch = random.randint(-1,1)
+        step = None if step_switch == 0 else []
+        for i in range(len(shape)):
+            s = random.randint(0, shape[i]-1)
+            e = random.randint(s+1, shape[i])
+            if step_switch == 1:
+                step.append(1)
+                start.append(s)
+                end.append(e)
+            elif step_switch == -1:
+                step.append(-1)
+                if e == shape[i]:
+                    e -= 1
+                    s -= 1
+                    if s == -1:
+                        s = None
+                start.append(e)
+                end.append(s)
+            else:
+                start.append(s)
+                end.append(e)
+        return start, end, step
+
+    for hybridize in [True, False]:
+        for i in range(10):
+            dim = random.randint(1,4)
+            shape = [random.randint(1,5) for i in range(dim)]
+
+            # test gluon
+            start, end, step = get_start_end_step(shape)
+            test_slice = TestSlice(begin=start, end=end, step=step)
+            if hybridize:
+                test_slice.hybridize()
+
+            a = mx.nd.random.uniform(shape=shape).as_np_ndarray()
+            a.attach_grad()
+            if step is not None:
+                expected_ret = a.as_nd_ndarray().slice(start, end, step)
+            else:
+                expected_ret = a.as_nd_ndarray().slice(start, end)
+            with mx.autograd.record():
+                y = test_slice(a)
+
+            assert_almost_equal(y.asnumpy(), expected_ret.asnumpy(), rtol=1e-3, atol=1e-5)
+
+            # test backward
+            mx.autograd.backward(y)
+            expected_grad = _np.zeros(shape)
+            basic_index = tuple([ 
+                slice(start[i], end[i], step[i]) if step is not None else slice(start[i], end[i])
+                for i in range(len(start))
+                ])
+            expected_grad[basic_index] = 1
+            assert_almost_equal(a.grad.asnumpy(), expected_grad, rtol=1e-3, atol=1e-5)
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()