enable other DTypes in SwapAxis

Author: vchuravy

src/operator/swapaxis-inl.h

@@ -40,7 +40,7 @@ struct SwapAxisParam : public dmlc::Parameter<SwapAxisParam> {
 };
 
 
-template<typename xpu>
+template<typename xpu, typename DType>
 class SwapAxisOp : public Operator {
  public:
   explicit SwapAxisOp(SwapAxisParam p) {
@@ -99,12 +99,12 @@ class SwapAxisOp : public Operator {
 
     Reshape2Five(&inter_shape, shape_in, dim1, dim2);
 
-    Tensor<xpu, 5> inter_data_in = data_in.get_with_shape<xpu, 5, real_t>(inter_shape, s);
+    Tensor<xpu, 5, DType> inter_data_in = data_in.get_with_shape<xpu, 5, DType>(inter_shape, s);
 
     Shape<5> inter_shape2 = inter_shape;
     std::swap(inter_shape2[1], inter_shape2[3]);
 
-    Tensor<xpu, 5> inter_data_out = data_out.get_with_shape<xpu, 5, real_t>(inter_shape2, s);
+    Tensor<xpu, 5, DType> inter_data_out = data_out.get_with_shape<xpu, 5, DType>(inter_shape2, s);
 
     inter_data_out = swapaxis<3, 1>(inter_data_in);
   }
@@ -138,7 +138,7 @@ class SwapAxisOp : public Operator {
 
 
 template<typename xpu>
-Operator* CreateOp(SwapAxisParam param);
+Operator* CreateOp(SwapAxisParam param, int dtype);
 
 
 #if DMLC_USE_CXX11
@@ -171,6 +171,17 @@ class SwapAxisProp : 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 {
     auto ptr = new SwapAxisProp();
     ptr->param_ = param_;
@@ -188,7 +199,13 @@ class SwapAxisProp : public OperatorProperty {
     return {out_grad[swapaxisenum::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;
 
  private:
   SwapAxisParam param_;

src/operator/swapaxis.cc

@@ -11,12 +11,21 @@ namespace mxnet {
 namespace op {
 
 template<>
-Operator* CreateOp<cpu>(SwapAxisParam param) {
-  return new SwapAxisOp<cpu>(param);
+Operator* CreateOp<cpu>(SwapAxisParam param, int dtype) {
+  Operator *op = NULL;
+  MSHADOW_REAL_TYPE_SWITCH(dtype, DType, {
+    op = new SwapAxisOp<cpu, DType>(param);
+  });
+  return op;
 }
 
-Operator* SwapAxisProp::CreateOperator(Context ctx) const {
-  DO_BIND_DISPATCH(CreateOp, param_);
+Operator* SwapAxisProp::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(InferShape(in_shape, &out_shape, &aux_shape));
+  CHECK(InferType(in_type, &out_type, &aux_type));
+  DO_BIND_DISPATCH(CreateOp, param_, in_type->at(0));
 }
 
 

src/operator/swapaxis.cu

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

tests/python/gpu/test_operator_gpu.py

@@ -123,6 +123,15 @@ def test_blockgrad_with_type():
                 {'ctx': mx.cpu(0), 'bg_data': (2, 2, 2, 10), 'type_dict': {'bg_data': np.float32}}]
     check_consistency(sym, ctx_list)
 
+def test_swapaxis_with_type():
+    sym = mx.sym.SwapAxis(name='swap', dim1=1)
+    ctx_list = [{'ctx': mx.gpu(0), 'swap_data': (2, 2, 2, 10), 'type_dict': {'swap_data': np.float64}},
+                {'ctx': mx.gpu(0), 'swap_data': (2, 2, 2, 10), 'type_dict': {'swap_data': np.float32}},
+                {'ctx': mx.gpu(0), 'swap_data': (2, 2, 2, 10), 'type_dict': {'swap_data': np.float16}},
+                {'ctx': mx.cpu(0), 'swap_data': (2, 2, 2, 10), 'type_dict': {'swap_data': np.float64}},
+                {'ctx': mx.cpu(0), 'swap_data': (2, 2, 2, 10), 'type_dict': {'swap_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}},
@@ -149,6 +158,7 @@ def test_activation_with_type():
     test_concat_with_type()
     test_reshape_with_type()
     test_blockgrad_with_type()
+    test_swapaxis_with_type()
     test_fullyconnected_with_type()
     test_activation_with_type()
     #test_softmax_with_shape((3,4), mx.gpu())