Add grad for lodtensor array ops (#5461)

* Add LoDRankTable

LoD Rank Table stores the `level` of `lod` which is ordered by sequence
length in descending order. It is useful when implement dynamic RNN and
is shared by dynamic RNN memory, dynamic RNN slice input and dynamic
RNN slice output operators.

* Add skeleton for array_to_lod_tensor and lod_tensor_to_array

* Add VarType::LoDTensorArray

* Add PyBind of LoDTensorArray

* Add InferVarType

* Add first unittest

* Add ut

* Add unittest

* Add unittest

* Add unittests

* update

* init

* add infershape for lod_tensor_to_array_op

* compelete array_to_lod_tensor_op

* copy data

* clean code

* clean code

* Fix unittest data

* fix bugs

* fix compile error

* Refine TensorToArrayOp

* refactor array_to_lod_tensor

* Unittest

* fix bugs

* Fix unittest

* Fix unittest

* debug

* Debug

* Fix unittest

* Add grad for ops

* Debug

* Fix a bug

* fix a bug

* fix a bug

paddle/operators/array_to_lod_tensor_op.cc

@@ -140,6 +140,23 @@ class ArrayToLoDTensorInferShape : public framework::InferShapeBase {
                    "ArrayToLoDTensorOp must has input X.");
     PADDLE_ENFORCE(context->HasInput("RankTable"),
                    "ArrayToLoDTensorOp must has input RankTable.");
+    context->SetOutputDim("Out", context->GetInputDim("X"));
+  }
+};
+
+class ArrayToLoDTensorGradMaker : public framework::SingleGradOpDescMaker {
+ public:
+  using framework::SingleGradOpDescMaker::SingleGradOpDescMaker;
+
+ protected:
+  std::unique_ptr<framework::OpDescBind> Apply() const override {
+    auto *grad_op = new framework::OpDescBind();
+    grad_op->SetType("lod_tensor_to_array");
+    grad_op->SetInput("X", OutputGrad("Out"));
+    grad_op->SetInput("RankTable", Input("RankTable"));
+    grad_op->SetOutput("Out", InputGrad("X"));
+    grad_op->SetAttrMap(Attrs());
+    return std::unique_ptr<framework::OpDescBind>(grad_op);
   }
 };
 
@@ -149,4 +166,5 @@ class ArrayToLoDTensorInferShape : public framework::InferShapeBase {
 namespace ops = paddle::operators;
 REGISTER_OPERATOR(array_to_lod_tensor, ops::ArrayToLoDTensorOp,
                   ops::ArrayToLoDTensorOpProtoMaker,
-                  ops::ArrayToLoDTensorInferShape);
+                  ops::ArrayToLoDTensorInferShape,
+                  ops::ArrayToLoDTensorGradMaker);

paddle/operators/lod_tensor_to_array_op.cc

@@ -133,6 +133,22 @@ class LoDTensorToArrayInferVarType : public framework::VarTypeInference {
   }
 };
 
+class LoDTensorToArrayGradMaker : public framework::SingleGradOpDescMaker {
+ public:
+  using framework::SingleGradOpDescMaker::SingleGradOpDescMaker;
+
+ protected:
+  std::unique_ptr<framework::OpDescBind> Apply() const override {
+    auto *grad_op = new framework::OpDescBind();
+    grad_op->SetType("array_to_lod_tensor");
+    grad_op->SetInput("X", OutputGrad("Out"));
+    grad_op->SetInput("RankTable", Input("RankTable"));
+    grad_op->SetOutput("Out", InputGrad("X"));
+    grad_op->SetAttrMap(Attrs());
+    return std::unique_ptr<framework::OpDescBind>(grad_op);
+  }
+};
+
 }  // namespace operators
 }  // namespace paddle
 
@@ -140,4 +156,5 @@ namespace ops = paddle::operators;
 REGISTER_OPERATOR(lod_tensor_to_array, ops::LoDTensorToArrayOp,
                   ops::LoDTensorToArrayOpProtoMaker,
                   ops::LoDTensorToArrayInferShape,
-                  ops::LoDTensorToArrayInferVarType);
+                  ops::LoDTensorToArrayInferVarType,
+                  ops::LoDTensorToArrayGradMaker);

paddle/operators/mean_op.cc

@@ -51,6 +51,7 @@ class MeanGradOp : public framework::OperatorWithKernel {
 
   void InferShape(framework::InferShapeContext* ctx) const override {
     ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
+    ctx->ShareLoD("X", framework::GradVarName("X"));
   }
 };
 

python/paddle/v2/framework/layers.py

@@ -87,7 +87,8 @@ def data(name,
          type=core.VarDesc.VarType.LOD_TENSOR,
          append_batch_size=True,
          main_program=None,
-         startup_program=None):
+         startup_program=None,
+         stop_gradient=True):
     helper = LayerHelper('data', **locals())
     shape = list(shape)
     for i in xrange(len(shape)):
@@ -101,7 +102,11 @@ def data(name,
         shape = [-1] + shape  # append batch size as -1
 
     return helper.create_global_variable(
-        name=name, shape=shape, dtype=data_type, type=type, stop_gradient=True)
+        name=name,
+        shape=shape,
+        dtype=data_type,
+        type=type,
+        stop_gradient=stop_gradient)
 
 
 def _convert_(name):
@@ -845,7 +850,8 @@ def lod_tensor_to_array(x, table, main_program=None):
     helper = LayerHelper("lod_tensor_to_array", **locals())
     array = helper.create_variable(
         name=unique_name("lod_tensor_to_array"),
-        type=core.VarDesc.VarType.LOD_TENSOR_ARRAY)
+        type=core.VarDesc.VarType.LOD_TENSOR_ARRAY,
+        dtype=x.data_type)
     helper.append_op(
         type='lod_tensor_to_array',
         inputs={'X': x,

python/paddle/v2/framework/tests/test_lod_tensor_array_ops.py

@@ -4,6 +4,7 @@ import numpy
 import paddle.v2.framework.layers as layers
 from paddle.v2.framework.framework import Program
 from paddle.v2.framework.executor import Executor
+from paddle.v2.framework.backward import append_backward_ops
 
 
 class TestCPULoDTensorArrayOps(unittest.TestCase):
@@ -123,5 +124,42 @@ class TestCPULoDTensorArrayOps(unittest.TestCase):
         self.assertEqual(actual.lod(), expect.lod())
 
 
+class TestCPULoDTensorArrayOpGrad(unittest.TestCase):
+    def test_grad(self):
+        place = core.CPUPlace()
+        program = Program()
+
+        x = layers.data(
+            name='x',
+            shape=[1],
+            data_type='float32',
+            main_program=program,
+            stop_gradient=False)
+        table = layers.lod_rank_table(x, level=0, main_program=program)
+        array = layers.lod_tensor_to_array(x, table, main_program=program)
+        result = layers.array_to_lod_tensor(array, table, main_program=program)
+
+        mean = layers.mean(x=result, main_program=program)
+
+        append_backward_ops(mean)
+
+        tensor = core.LoDTensor()
+        tensor.set(numpy.arange(10).reshape(10, 1).astype('float32'), place)
+        tensor.set_lod([[0, 3, 9, 10]])
+
+        g_vars = program.global_block().var(x.name + "@GRAD")
+
+        exe = Executor(place)
+        g_out = [
+            item.sum()
+            for item in map(
+                numpy.array,
+                exe.run(program, feed={'x': tensor}, fetch_list=[g_vars]))
+        ]
+        g_out_sum = numpy.array(g_out).sum()
+
+        self.assertAlmostEqual(1.0, g_out_sum, delta=0.1)
+
+
 if __name__ == '__main__':
     unittest.main()