[MXNET-382] Shape and Size Operator

docs/api/python/ndarray/ndarray.md

@@ -124,6 +124,8 @@ The `ndarray` package provides several classes:
     :nosignatures:
 
     NDArray.T
+    NDArray.shape_nd
+    NDArray.size_nd
     NDArray.reshape
     NDArray.reshape_like
     NDArray.flatten
@@ -375,6 +377,8 @@ The `ndarray` package provides several classes:
     :nosignatures:
 
     cast
+    shape_nd
+    size_nd
     reshape
     reshape_like
     flatten

docs/api/python/symbol/symbol.md

@@ -191,6 +191,8 @@ Composite multiple symbols into a new one by an operator.
     :nosignatures:
 
     Symbol.astype
+    Symbol.shape_nd
+    Symbol.size_nd
     Symbol.reshape
     Symbol.reshape_like
     Symbol.flatten
@@ -371,6 +373,8 @@ Composite multiple symbols into a new one by an operator.
     :nosignatures:
 
     cast
+    shape_nd
+    size_nd
     reshape
     reshape_like
     flatten

python/mxnet/ndarray/ndarray.py

@@ -1254,6 +1254,22 @@ def flatten(self, *args, **kwargs):
         """
         return op.flatten(self, *args, **kwargs)
 
+    def shape_nd(self, *args, **kwargs):
+        """Convenience fluent method for :py:func:`shape_op`.
+
+        The arguments are the same as for :py:func:`shape_op`, with
+        this array as data.
+        """
+        return op.shape_nd(self, *args, **kwargs)
+
+    def size_nd(self, *args, **kwargs):
+        """Convenience fluent method for :py:func:`size_op`.
+
+        The arguments are the same as for :py:func:`size_op`, with
+        this array as data.
+        """
+        return op.size_nd(self, *args, **kwargs)
+
     def expand_dims(self, *args, **kwargs):
         """Convenience fluent method for :py:func:`expand_dims`.
 

python/mxnet/symbol/symbol.py

@@ -1978,6 +1978,22 @@ def flatten(self, *args, **kwargs):
         """
         return op.flatten(self, *args, **kwargs)
 
+    def shape_nd(self, *args, **kwargs):
+        """Convenience fluent method for :py:func:`shape_op`.
+
+        The arguments are the same as for :py:func:`shape_op`, with
+        this array as data.
+        """
+        return op.shape_nd(self, *args, **kwargs)
+
+    def size_nd(self, *args, **kwargs):
+        """Convenience fluent method for :py:func:`size_op`.
+
+        The arguments are the same as for :py:func:`size_op`, with
+        this array as data.
+        """
+        return op.size_nd(self, *args, **kwargs)
+
     def expand_dims(self, *args, **kwargs):
         """Convenience fluent method for :py:func:`expand_dims`.
 

src/operator/tensor/elemwise_unary_op.h

@@ -388,6 +388,39 @@ void CastCompute(const nnvm::NodeAttrs& attrs,
   });
 }
 
+template<typename xpu>
+void ShapeCompute(const nnvm::NodeAttrs& attrs,
+                  const OpContext& ctx,
+                  const std::vector<TBlob>& inputs,
+                  const std::vector<OpReqType>& req,
+                  const std::vector<TBlob>& outputs) {
+  CHECK_EQ(inputs.size(), 1U);
+  CHECK_EQ(outputs.size(), 1U);
+  CHECK_EQ(req.size(), 1U);
+  const TBlob& in_data = inputs[0];
+  const TBlob& out_data = outputs[0];
+  mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
+  const TShape& in_shape = in_data.shape_;
+  MSHADOW_TYPE_SWITCH(out_data.type_flag_, DType, {
+    mxnet_op::Kernel<mshadow_op::identity_with_cast, xpu>::Launch(
+      s, in_data.ndim(), out_data.dptr<int64_t>(), in_shape.data());
+  });
+}
+
+template<typename xpu>
+void SizeCompute(const nnvm::NodeAttrs& attrs,
+                 const OpContext& ctx,
+                 const std::vector<TBlob>& inputs,
+                 const std::vector<OpReqType>& req,
+                 const std::vector<TBlob>& outputs) {
+  CHECK_EQ(inputs.size(), 1U);
+  CHECK_EQ(outputs.size(), 1U);
+  CHECK_EQ(req.size(), 1U);
+  const TBlob& in_data = inputs[0];
+  const TBlob& out_data = outputs[0];
+  out_data.dptr<int64_t>()[0] = in_data.Size();
+}
+
 struct HardSigmoidParam : public dmlc::Parameter<HardSigmoidParam> {
   real_t alpha;
   real_t beta;

src/operator/tensor/elemwise_unary_op_basic.cc

@@ -399,6 +399,71 @@ NNVM_REGISTER_OP(reshape_like)
 .add_argument("lhs", "NDArray-or-Symbol", "First input.")
 .add_argument("rhs", "NDArray-or-Symbol", "Second input.");
 
+NNVM_REGISTER_OP(shape_nd)
+.describe(R"code(Returns a 1D int64 array containing the shape of data.
+
+Example::
+
+  shape_nd([[1,2,3,4], [5,6,7,8]]) = [2,4]
+
+)code" ADD_FILELINE)
+.set_num_inputs(1)
+.set_num_outputs(1)
+.set_attr<nnvm::FIgnoreInputs>("FIgnoreInputs",
+    [](const NodeAttrs& attrs) { return std::vector<uint32_t>(1, 1); })
+.set_attr<FCompute>("FCompute<cpu>", ShapeCompute<cpu>)
+.set_attr<nnvm::FInferShape>("FInferShape",
+   [](const nnvm::NodeAttrs& attrs,
+      std::vector<TShape> *in_attrs,
+      std::vector<TShape> *out_attrs) {
+       CHECK_EQ(in_attrs->size(), 1U);
+       CHECK_EQ(out_attrs->size(), 1U);
+       TShape target_shape(1);
+       target_shape[0] = in_attrs->at(0).ndim();
+       SHAPE_ASSIGN_CHECK(*out_attrs, 0, target_shape);
+       return !shape_is_none(out_attrs->at(0));
+     })
+.set_attr<nnvm::FInferType>("FInferType",
+   [](const nnvm::NodeAttrs& attrs,
+      std::vector<int>* in_attrs,
+      std::vector<int>* out_attrs) {
+       CHECK_EQ(in_attrs->size(), 1U);
+       CHECK_EQ(out_attrs->size(), 1U);
+       TYPE_ASSIGN_CHECK(*out_attrs, 0, mshadow::kInt64);
+       return out_attrs->at(0) != -1;
+     })
+.add_argument("data", "NDArray-or-Symbol", "Input Array.");
+
+NNVM_REGISTER_OP(size_nd)
+.describe(R"code(Returns a 1D int64 array containing the size of data.
+
+Example::
+
+  size_nd([[1,2,3,4], [5,6,7,8]]) = [8]
+
+)code" ADD_FILELINE)
+.set_num_inputs(1)
+.set_num_outputs(1)
+.set_attr<FCompute>("FCompute<cpu>", SizeCompute<cpu>)
+.set_attr<nnvm::FInferShape>("FInferShape",
+   [](const nnvm::NodeAttrs& attrs,
+      std::vector<TShape> *in_attrs,
+      std::vector<TShape> *out_attrs) {
+       CHECK_EQ(in_attrs->size(), 1U);
+       CHECK_EQ(out_attrs->size(), 1U);
+       SHAPE_ASSIGN_CHECK(*out_attrs, 0, 1U);
+       return !shape_is_none(out_attrs->at(0));
+      })
+.set_attr<nnvm::FInferType>("FInferType",
+   [](const nnvm::NodeAttrs& attrs,
+      std::vector<int>* in_attrs,
+      std::vector<int>* out_attrs) {
+       CHECK_EQ(in_attrs->size(), 1U);
+       CHECK_EQ(out_attrs->size(), 1U);
+       TYPE_ASSIGN_CHECK(*out_attrs, 0, mshadow::kInt64);
+       return out_attrs->at(0) != -1;
+      })
+.add_argument("data", "NDArray-or-Symbol", "Input Array.");
 
 DMLC_REGISTER_PARAMETER(CastParam);
 NNVM_REGISTER_OP(Cast)

src/operator/tensor/elemwise_unary_op_basic.cu

@@ -77,6 +77,12 @@ NNVM_REGISTER_OP(_identity_with_attr_like_rhs)
 NNVM_REGISTER_OP(reshape_like)
 .set_attr<FCompute>("FCompute<gpu>", UnaryOp::IdentityCompute<gpu>);
 
+NNVM_REGISTER_OP(shape_nd)
+.set_attr<FCompute>("FCompute<gpu>", ShapeCompute<gpu>);
+
+NNVM_REGISTER_OP(size_nd)
+.set_attr<FCompute>("FCompute<gpu>", SizeCompute<gpu>);
+
 NNVM_REGISTER_OP(Cast)
 .set_attr<FCompute>("FCompute<gpu>", CastCompute<gpu>);
 

tests/python/unittest/test_operator.py

@@ -584,6 +584,24 @@ def fsigmoid(a):
     check_symbolic_forward(y, [xa], [ya])
     check_symbolic_backward(y, [xa], [np.ones(shape)], [ya * (1 - ya)])
 
+@with_seed()
+def test_shape():
+    shape = [3, 4]
+    x = mx.symbol.Variable("x")
+    y = mx.sym.shape_nd(x)
+    xa = np.random.uniform(low=-1.0,high=1.0,size=shape)
+    ya = np.shape(xa)
+    check_symbolic_forward(y, [xa], [ya])
+
+@with_seed()
+def test_size():
+    shape = [3, 4]
+    x = mx.symbol.Variable("x")
+    y = mx.sym.size_nd(x)
+    xa = np.random.uniform(low=-1.0,high=1.0,size=shape)
+    ya = np.array(xa.size).astype(np.int64)
+    check_symbolic_forward(y, [xa], [ya])
+
 @with_seed()
 def test_hard_sigmoid():
     def fhardsigmoid(a, alpha=0.2, beta=0.5):