[MXNET-382] Shape and Size Operator

src/operator/tensor/elemwise_unary_op.h

@@ -388,6 +388,26 @@ 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>();
+  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<DType>(), in_shape.data());
+  });
+}
+
+
 struct HardSigmoidParam : public dmlc::Parameter<HardSigmoidParam> {
   real_t alpha;
   real_t beta;

src/operator/tensor/elemwise_unary_op_basic.cc

@@ -399,6 +399,37 @@ 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)
+.describe("Returns a 1D int64 array containing the shape of data.")
+.set_num_inputs(1)
+.set_num_outputs(1)
+.set_attr<nnvm::FInplaceIdentity>("FInplaceIdentity",
+    [](const NodeAttrs& attrs){ return std::vector<bool>{true}; })
+.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 out_attrs->at(0).ndim() != 0U && out_attrs->at(0).Size() != 0U;
+	    })
+.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, 0U);
+	    return out_attrs->at(0) != -1;
+		})
+.set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
+.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,9 @@ 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)
+.set_attr<FCompute>("FCompute<gpu>", ShapeCompute<gpu>);
+
 NNVM_REGISTER_OP(Cast)
 .set_attr<FCompute>("FCompute<gpu>", CastCompute<gpu>);
 

tests/python/unittest/test_operator.py

@@ -584,6 +584,15 @@ 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(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_hard_sigmoid():
     def fhardsigmoid(a, alpha=0.2, beta=0.5):