Bugfix: kAddTo support for swap_axes.

apache · Jan 24, 2018 · 758e4f8 · 758e4f8
1 parent 20253d5
commit 758e4f8
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/src/operator/swapaxis-inl.h b/src/operator/swapaxis-inl.h
@@ -101,15 +101,17 @@ class SwapAxisOp : public Operator {
   }
 
   void SwapAxis(mshadow::Stream<xpu> *s,
-                  const std::vector<TBlob> &in_data,
-                  const std::vector<TBlob> &out_data) {
+                const std::vector<TBlob> &in_data,
+                const std::vector<TBlob> &out_data,
+                const std::vector<OpReqType> &req) {
     using namespace mshadow;
     using namespace mshadow::expr;
     uint32_t dim1 = param_.dim1;
     uint32_t dim2 = param_.dim2;
 
     TBlob data_in = in_data[swapaxisenum::kData];
     TBlob data_out = out_data[swapaxisenum::kData];
+    OpReqType out_req = req[swapaxisenum::kData];
 
     TShape shape_in = data_in.shape_;
     TShape shape_out = data_out.shape_;
@@ -125,7 +127,11 @@ class SwapAxisOp : public Operator {
 
     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);
+    if (out_req == kAddTo) {
+        inter_data_out += swapaxis<3, 1>(inter_data_in);
+    } else {
+        inter_data_out = swapaxis<3, 1>(inter_data_in);
+    }
   }
 
   virtual void Forward(const OpContext &ctx,
@@ -136,7 +142,7 @@ class SwapAxisOp : public Operator {
     using namespace mshadow;
     Stream<xpu> *s = ctx.get_stream<xpu>();
 
-    SwapAxis(s, in_data, out_data);
+    SwapAxis(s, in_data, out_data, req);
   }
 
   virtual void Backward(const OpContext &ctx,
@@ -149,7 +155,7 @@ class SwapAxisOp : public Operator {
     using namespace mshadow;
     Stream<xpu> *s = ctx.get_stream<xpu>();
 
-    SwapAxis(s, out_grad, in_grad);
+    SwapAxis(s, out_grad, in_grad, req);
   }
 
   SwapAxisParam param_;