enable other DTypes in BlockGrad

Author: vchuravy

src/operator/block_grad-inl.h

@@ -24,7 +24,7 @@ enum BlockGradientOpInputs {kData};
 enum BlockGradientOpOutputs {kOut};
 }  // namespace blockgrad
 
-template<typename xpu>
+template<typename xpu, typename DType>
 class BlockGradientOp : public Operator {
  public:
   virtual void Forward(const OpContext &ctx,
@@ -37,8 +37,8 @@ class BlockGradientOp : public Operator {
     CHECK_EQ(in_data.size(), 1);
     CHECK_EQ(out_data.size(), 1);
     Stream<xpu> *s = ctx.get_stream<xpu>();
-    Tensor<xpu, 2> data = in_data[blockgrad::kData].FlatTo2D<xpu, real_t>(s);
-    Tensor<xpu, 2> out = out_data[blockgrad::kOut].FlatTo2D<xpu, real_t>(s);
+    Tensor<xpu, 2, DType> data = in_data[blockgrad::kData].FlatTo2D<xpu, DType>(s);
+    Tensor<xpu, 2, DType> out = out_data[blockgrad::kOut].FlatTo2D<xpu, DType>(s);
     out = F<mshadow_op::identity>(data);
   }
 
@@ -52,13 +52,13 @@ class BlockGradientOp : public Operator {
     using namespace mshadow;
     using namespace mshadow::expr;
     Stream<xpu> *s = ctx.get_stream<xpu>();
-    Tensor<xpu, 2> grad = in_grad[blockgrad::kData].FlatTo2D<xpu, real_t>(s);
+    Tensor<xpu, 2, DType> grad = in_grad[blockgrad::kData].FlatTo2D<xpu, DType>(s);
     grad = 0.f;
   }
 };  // class BlockGradientOp
 
 template<typename xpu>
-Operator *CreateOp();
+Operator *CreateOp(int dtype);
 
 #if DMLC_USE_CXX11
 class BlockGradientProp : public OperatorProperty {
@@ -81,6 +81,18 @@ class BlockGradientProp : public OperatorProperty {
     return true;
   }
 
+  bool InferType(std::vector<int> *in_type,
+                 std::vector<int> *out_type,
+                 std::vector<int> *aux_type) const override {
+
+    CHECK_EQ(in_type->size(), 1);
+    int dtype = (*in_type)[0];
+    CHECK_NE(dtype, -1) << "Input must have specified type";
+    out_type->clear();
+    out_type->push_back(dtype);
+    return true;
+  }
+
   OperatorProperty* Copy() const override {
     return new BlockGradientProp();
   }
@@ -102,7 +114,13 @@ class BlockGradientProp : public OperatorProperty {
     return {{in_data[blockgrad::kData], out_data[blockgrad::kOut]}};
   }
 
-  Operator* CreateOperator(Context ctx) const override;
+  Operator* CreateOperator(Context ctx) const override {
+    LOG(FATAL) << "Not Implemented";
+    return NULL;
+  }
+
+  Operator* CreateOperatorEx(Context ctx, std::vector<TShape> *in_shape,
+                             std::vector<int> *in_type) const override;
 };  // class BlockGradientProperty
 
 #endif  // DMLC_USE_CXX11

src/operator/block_grad.cc

@@ -9,12 +9,21 @@
 namespace mxnet {
 namespace op {
 template<>
-Operator *CreateOp<cpu>() {
-  return new BlockGradientOp<cpu>();
+Operator *CreateOp<cpu>(int dtype) {
+  Operator *op = NULL;
+  MSHADOW_REAL_TYPE_SWITCH(dtype, DType, {
+    op = new BlockGradientOp<cpu, DType>();
+  });
+  return op;
 }
 
-Operator *BlockGradientProp::CreateOperator(Context ctx) const {
-  DO_BIND_DISPATCH(CreateOp);
+Operator *BlockGradientProp::CreateOperatorEx(Context ctx, std::vector<TShape> *in_shape,
+                                              std::vector<int> *in_type) const {
+  std::vector<TShape> out_shape, aux_shape;
+  std::vector<int> out_type, aux_type;
+  CHECK(InferType(in_type, &out_type, &aux_type));
+  CHECK(InferShape(in_shape, &out_shape, &aux_shape));
+  DO_BIND_DISPATCH(CreateOp, in_type->at(0));
 }
 
 MXNET_REGISTER_OP_PROPERTY(BlockGrad, BlockGradientProp)

src/operator/block_grad.cu

@@ -9,8 +9,12 @@
 namespace mxnet {
 namespace op {
 template<>
-Operator *CreateOp<gpu>() {
-  return new BlockGradientOp<gpu>();
+Operator *CreateOp<gpu>(int dtype) {
+  Operator *op = NULL;
+  MSHADOW_REAL_TYPE_SWITCH(dtype, DType, {
+    op = new BlockGradientOp<gpu, DType>();
+  });
+  return op;
 }
 
 }  // namespace op

tests/python/gpu/test_operator_gpu.py

@@ -114,6 +114,15 @@ def test_reshape_with_type():
                 {'ctx': mx.cpu(0), 'reshape_data': (2, 2, 2, 10), 'type_dict': {'reshape_data': np.float32}}]
     check_consistency(sym, ctx_list)
 
+def test_blockgrad_with_type():
+    sym = mx.sym.BlockGrad(name='bg')
+    ctx_list = [{'ctx': mx.gpu(0), 'bg_data': (2, 2, 2, 10), 'type_dict': {'bg_data': np.float64}},
+                {'ctx': mx.gpu(0), 'bg_data': (2, 2, 2, 10), 'type_dict': {'bg_data': np.float32}},
+                {'ctx': mx.gpu(0), 'bg_data': (2, 2, 2, 10), 'type_dict': {'bg_data': np.float16}},
+                {'ctx': mx.cpu(0), 'bg_data': (2, 2, 2, 10), 'type_dict': {'bg_data': np.float64}},
+                {'ctx': mx.cpu(0), 'bg_data': (2, 2, 2, 10), 'type_dict': {'bg_data': np.float32}}]
+    check_consistency(sym, ctx_list)
+
 def test_fullyconnected_with_type():
     sym = mx.sym.FullyConnected(num_hidden=3, name='inner')
     ctx_list = [{'ctx': mx.gpu(0), 'inner_data': (2, 10), 'type_dict': {'inner_data': np.float64}},
@@ -139,6 +148,7 @@ def test_activation_with_type():
     test_upsampling_with_type()
     test_concat_with_type()
     test_reshape_with_type()
+    test_blockgrad_with_type()
     test_fullyconnected_with_type()
     test_activation_with_type()
     #test_softmax_with_shape((3,4), mx.gpu())