move meshgrid yaml (#41411)

paddle/phi/api/lib/api_custom_impl.cc

@@ -410,5 +410,153 @@ std::vector<Tensor> stack_grad_impl(const std::vector<Tensor>& x,
   return x_grad;
 }
 
+std::vector<Tensor> meshgrid_impl(const std::vector<Tensor>& inputs) {
+  Backend kernel_backend = Backend::UNDEFINED;
+  DataLayout kernel_layout = DataLayout::UNDEFINED;
+  DataType kernel_data_type = DataType::UNDEFINED;
+
+  if (kernel_backend == Backend::UNDEFINED ||
+      kernel_layout == DataLayout::UNDEFINED ||
+      kernel_data_type == DataType::UNDEFINED) {
+    auto kernel_key_set = ParseKernelKeyByInputArgs(inputs);
+    auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey();
+    if (kernel_backend == Backend::UNDEFINED) {
+      kernel_backend = kernel_key.backend();
+    }
+    if (kernel_layout == DataLayout::UNDEFINED) {
+      kernel_layout = kernel_key.layout();
+    }
+    if (kernel_data_type == DataType::UNDEFINED) {
+      kernel_data_type = kernel_key.dtype();
+    }
+  }
+
+  const auto& kernel = phi::KernelFactory::Instance().SelectKernelOrThrowError(
+      "meshgrid", {kernel_backend, kernel_layout, kernel_data_type});
+  VLOG(6) << "meshgrid API kernel key: [" << kernel_backend << ", "
+          << kernel_layout << ", " << kernel_data_type << "]";
+  VLOG(6) << "meshgrid API kernel: " << kernel;
+
+  auto* dev_ctx = GetDeviceContextByBackend(kernel_backend);
+
+  auto input_inputs_vec = PrepareData(inputs, kernel.InputAt(0), {});
+  std::vector<const phi::DenseTensor*> input_inputs(input_inputs_vec->size());
+  for (size_t i = 0; i < input_inputs.size(); ++i) {
+    input_inputs[i] = &input_inputs_vec->at(i);
+  }
+
+  auto x_meta_vec = MakeMetaTensor(input_inputs);
+  std::vector<phi::MetaTensor*> inputs_metas(x_meta_vec.size());
+  for (size_t i = 0; i < x_meta_vec.size(); ++i) {
+    inputs_metas[i] = &x_meta_vec[i];
+  }
+
+  // Calculate the number of out tensors
+  size_t out_number = inputs.size();
+
+  std::vector<Tensor> out;
+  auto dense_outs = SetKernelOutput(out_number, kernel_backend, &out);
+
+  std::vector<phi::MetaTensor> meta_outs;
+  meta_outs.reserve(out_number);
+  std::vector<phi::MetaTensor*> meta_out_ptrs;
+  meta_out_ptrs.reserve(out_number);
+  for (size_t i = 0; i < out_number; ++i) {
+    meta_outs.push_back(dense_outs[i]);
+    meta_out_ptrs.push_back(&meta_outs.back());
+  }
+  phi::MeshgridInferMeta(inputs_metas, meta_out_ptrs);
+
+  using kernel_signature = void (*)(const platform::DeviceContext&,
+                                    const std::vector<const phi::DenseTensor*>&,
+                                    std::vector<phi::DenseTensor*>&);
+  auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
+  (*kernel_fn)(*dev_ctx, input_inputs, dense_outs);
+
+  return out;
+}
+
+std::vector<Tensor> meshgrid_grad_impl(
+    const std::vector<Tensor>& inputs,
+    const std::vector<Tensor>& outputs_grad) {
+  Backend kernel_backend = Backend::UNDEFINED;
+  DataLayout kernel_layout = DataLayout::UNDEFINED;
+  DataType kernel_data_type = DataType::UNDEFINED;
+
+  if (kernel_backend == Backend::UNDEFINED ||
+      kernel_layout == DataLayout::UNDEFINED ||
+      kernel_data_type == DataType::UNDEFINED) {
+    auto kernel_key_set = ParseKernelKeyByInputArgs(inputs, outputs_grad);
+    auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey();
+    if (kernel_backend == Backend::UNDEFINED) {
+      kernel_backend = kernel_key.backend();
+    }
+    if (kernel_layout == DataLayout::UNDEFINED) {
+      kernel_layout = kernel_key.layout();
+    }
+    if (kernel_data_type == DataType::UNDEFINED) {
+      kernel_data_type = kernel_key.dtype();
+    }
+  }
+
+  const auto& kernel = phi::KernelFactory::Instance().SelectKernelOrThrowError(
+      "meshgrid_grad", {kernel_backend, kernel_layout, kernel_data_type});
+  VLOG(6) << "meshgrid_grad API kernel key: [" << kernel_backend << ", "
+          << kernel_layout << ", " << kernel_data_type << "]";
+  VLOG(6) << "meshgrid_grad API kernel: " << kernel;
+
+  auto* dev_ctx = GetDeviceContextByBackend(kernel_backend);
+
+  auto input_inputs_vec = PrepareData(inputs, kernel.InputAt(0), {});
+  std::vector<const phi::DenseTensor*> input_inputs(input_inputs_vec->size());
+  for (size_t i = 0; i < input_inputs.size(); ++i) {
+    input_inputs[i] = &input_inputs_vec->at(i);
+  }
+  auto input_outputs_grad_vec =
+      PrepareData(outputs_grad, kernel.InputAt(1), {});
+  std::vector<const phi::DenseTensor*> input_outputs_grad(
+      input_outputs_grad_vec->size());
+  for (size_t i = 0; i < input_outputs_grad.size(); ++i) {
+    input_outputs_grad[i] = &input_outputs_grad_vec->at(i);
+  }
+
+  size_t out_number = inputs.size();
+  std::vector<Tensor> api_output;
+  auto kernel_out = SetKernelOutput(out_number, kernel_backend, &api_output);
+
+  auto inputs_meta_vec = MakeMetaTensor(input_inputs);
+  std::vector<phi::MetaTensor*> inputs_metas(inputs_meta_vec.size());
+  for (size_t i = 0; i < inputs_meta_vec.size(); ++i) {
+    inputs_metas[i] = &inputs_meta_vec[i];
+  }
+
+  auto outputs_grad_meta_vec = MakeMetaTensor(input_outputs_grad);
+  std::vector<phi::MetaTensor*> outputs_grad_metas(
+      outputs_grad_meta_vec.size());
+  for (size_t i = 0; i < outputs_grad_meta_vec.size(); ++i) {
+    outputs_grad_metas[i] = &outputs_grad_meta_vec[i];
+  }
+
+  std::vector<phi::MetaTensor> meta_outs;
+  meta_outs.reserve(out_number);
+  std::vector<phi::MetaTensor*> meta_out_ptrs;
+  meta_out_ptrs.reserve(out_number);
+  for (size_t i = 0; i < out_number; ++i) {
+    meta_outs.push_back(kernel_out[i]);
+    meta_out_ptrs.push_back(&meta_outs.back());
+  }
+
+  phi::MeshgridGradInferMeta(inputs_metas, outputs_grad_metas, meta_out_ptrs);
+
+  using kernel_signature = void (*)(const platform::DeviceContext&,
+                                    const std::vector<const phi::DenseTensor*>&,
+                                    const std::vector<const phi::DenseTensor*>&,
+                                    std::vector<phi::DenseTensor*>&);
+  auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
+  (*kernel_fn)(*dev_ctx, input_inputs, input_outputs_grad, kernel_out);
+
+  return api_output;
+}
+
 }  // namespace experimental
 }  // namespace paddle

paddle/phi/api/lib/api_custom_impl.h

@@ -59,6 +59,9 @@ std::vector<Tensor> concat_grad_impl(const std::vector<Tensor>& x,
 std::vector<Tensor> stack_grad_impl(const std::vector<Tensor>& x,
                                     const Tensor& out_grad,
                                     int axis);
+std::vector<Tensor> meshgrid_impl(const std::vector<Tensor>& inputs);
+std::vector<Tensor> meshgrid_grad_impl(const std::vector<Tensor>& inputs,
+                                       const std::vector<Tensor>& outputs_grad);
 
 }  // namespace experimental
 }  // namespace paddle

paddle/phi/infermeta/backward.cc

@@ -245,6 +245,20 @@ void MaxPoolWithIndexGradInferMeta(const MetaTensor& x,
   dx->share_meta(x);
 }
 
+void MeshgridGradInferMeta(const std::vector<MetaTensor*>& inputs,
+                           const std::vector<MetaTensor*>& outputs_grad,
+                           std::vector<MetaTensor*> inputs_grad) {
+  PADDLE_ENFORCE_GT(outputs_grad.size(),
+                    1,
+                    errors::InvalidArgument(
+                        "Number of Inputs(Out@Grad) should be larger than 1."
+                        "But received Inputs(Out@Grad)' size = %d .",
+                        outputs_grad.size()));
+  for (size_t i = 0; i < inputs.size(); i++) {
+    inputs_grad[i]->share_meta(*inputs[i]);
+  }
+}
+
 void NllLossGradInferMeta(const MetaTensor& x,
                           const MetaTensor& label,
                           paddle::optional<const MetaTensor&> weight,

paddle/phi/infermeta/backward.h

@@ -115,6 +115,10 @@ void MaxPoolWithIndexGradInferMeta(const MetaTensor& x,
                                    bool adaptive,
                                    MetaTensor* dx);
 
+void MeshgridGradInferMeta(const std::vector<MetaTensor*>& inputs,
+                           const std::vector<MetaTensor*>& outputs_grad,
+                           std::vector<MetaTensor*> inputs_grad);
+
 void NllLossGradInferMeta(const MetaTensor& input,
                           const MetaTensor& label,
                           paddle::optional<const MetaTensor&> weight,

python/paddle/fluid/tests/unittests/test_meshgrid_op.py

@@ -20,6 +20,7 @@ from op_test import OpTest, skip_check_grad_ci
 import paddle.fluid as fluid
 import paddle
 from paddle.fluid import compiler, Program, program_guard, core
+from paddle.fluid.framework import _test_eager_guard
 
 
 class TestMeshgridOp(OpTest):
@@ -149,6 +150,10 @@ class TestMeshgridOp6(unittest.TestCase):
             assert np.array_equal(res_3.shape, [100, 200])
             assert np.array_equal(res_4.shape, [100, 200])
 
+    def test_api_eager_dygraph(self):
+        with _test_eager_guard():
+            self.test_api_with_dygraph()
+
 
 class TestMeshgridOp7(unittest.TestCase):
     def test_api_with_dygraph_list_input(self):
@@ -163,6 +168,10 @@ class TestMeshgridOp7(unittest.TestCase):
             assert np.array_equal(res_3.shape, [100, 200])
             assert np.array_equal(res_4.shape, [100, 200])
 
+    def test_api_eager_dygraph(self):
+        with _test_eager_guard():
+            self.test_api_with_dygraph_list_input()
+
 
 class TestMeshgridOp8(unittest.TestCase):
     def test_api_with_dygraph_tuple_input(self):
@@ -177,6 +186,40 @@ class TestMeshgridOp8(unittest.TestCase):
             assert np.array_equal(res_3.shape, [100, 200])
             assert np.array_equal(res_4.shape, [100, 200])
 
+    def test_api_eager_dygraph(self):
+        with _test_eager_guard():
+            self.test_api_with_dygraph_tuple_input()
+
+
+class TestMeshgridEager(unittest.TestCase):
+    def test_dygraph_final_state_api(self):
+        input_1 = np.random.randint(0, 100, [100, ]).astype('int32')
+        input_2 = np.random.randint(0, 100, [200, ]).astype('int32')
+
+        with fluid.dygraph.guard():
+            tensor_1 = fluid.dygraph.to_variable(input_1)
+            tensor_2 = fluid.dygraph.to_variable(input_2)
+            tensor_1.stop_gradient = False
+            tensor_2.stop_gradient = False
+            res_1, res_2 = paddle.tensor.meshgrid((tensor_1, tensor_2))
+            sum = paddle.add_n([res_1, res_2])
+            sum.backward()
+            with _test_eager_guard():
+                tensor_eager_1 = fluid.dygraph.to_variable(input_1)
+                tensor_eager_2 = fluid.dygraph.to_variable(input_2)
+                tensor_eager_1.stop_gradient = False
+                tensor_eager_2.stop_gradient = False
+                res_eager_1, res_eager_2 = paddle.tensor.meshgrid(
+                    (tensor_eager_1, tensor_eager_2))
+                sum_eager = paddle.add_n([res_eager_1, res_eager_2])
+                sum_eager.backward()
+                self.assertEqual((
+                    tensor_1.grad.numpy() == tensor_eager_1.grad.numpy()).all(),
+                                 True)
+                self.assertEqual((
+                    tensor_2.grad.numpy() == tensor_eager_2.grad.numpy()).all(),
+                                 True)
+
 
 if __name__ == '__main__':
     paddle.enable_static()

python/paddle/tensor/creation.py

@@ -776,10 +776,12 @@ def meshgrid(*args, **kwargs):
 
     if len(args) == 1 and isinstance(args[0], (list, tuple)):
         args = args[0]
-    if paddle.in_dynamic_mode():
+    if _in_legacy_dygraph():
         num = len(args)
         out = _C_ops.meshgrid(list(args), num)
         return out
+    if in_dygraph_mode():
+        return _C_ops.final_state_meshgrid(list(args))
 
     name = kwargs.get("name", None)
     helper = LayerHelper('meshgrid', **locals())

python/paddle/utils/code_gen/api.yaml

@@ -1120,6 +1120,12 @@
     func : mean
   backward : mean_grad
 
+- api : meshgrid
+  args : (Tensor[] inputs)
+  output : Tensor[]
+  invoke : meshgrid_impl(inputs)
+  backward : meshgrid_grad
+
 - api : min
   args : (Tensor x, int64_t[] dims={}, bool keep_dim=false)
   output : Tensor(out)

python/paddle/utils/code_gen/backward.yaml

@@ -777,6 +777,12 @@
   kernel :
     func : mean_grad
 
+- backward_api : meshgrid_grad
+  forward : meshgrid (Tensor[] inputs) -> Tensor[](outputs)
+  args : (Tensor[] inputs, Tensor[] outputs_grad)
+  output : Tensor[](inputs_grad)
+  invoke : meshgrid_grad_impl(inputs, outputs_grad)
+
 - backward_api : min_grad
   forward: min (Tensor x,  int64_t[] dims={},  bool keep_dim=false) -> Tensor(out)
   args : (Tensor x, Tensor out, Tensor out_grad, int64_t[] dims={},  bool keep_dim=false, bool reduce_all=false)