// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <gtest/gtest.h>

#include "paddle/common/macros.h"
#include "paddle/fluid/framework/new_executor/pir_interpreter.h"
#include "paddle/fluid/framework/new_executor/standalone_executor.h"
#include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
#include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
#include "paddle/fluid/pir/transforms/pd_op_to_kernel_pass.h"
#include "paddle/pir/include/core/block.h"
#include "paddle/pir/include/core/builtin_attribute.h"
#include "paddle/pir/include/core/builtin_op.h"
#include "paddle/pir/include/core/ir_context.h"
#include "paddle/pir/include/core/program.h"
#include "paddle/pir/include/core/utils.h"

DECLARE_FILE_SYMBOLS(kernel_dialect);

PD_DECLARE_KERNEL(full, CPU, ALL_LAYOUT);
PD_DECLARE_KERNEL(tanh, CPU, ALL_LAYOUT);
PD_DECLARE_KERNEL(mean, CPU, ALL_LAYOUT);
PD_DECLARE_KERNEL(tanh_grad, CPU, ALL_LAYOUT);
PD_DECLARE_KERNEL(mean_grad, CPU, ALL_LAYOUT);
PD_DECLARE_KERNEL(add_grad, CPU, ALL_LAYOUT);

namespace paddle {
namespace framework {

pir::Operation* GetOpFromProgram(const std::string& op_name,
                                 const pir::Program& program) {
  for (auto& op : *(program.block())) {
    if (op.name() == op_name) {
      return &op;
    }
  }
  return nullptr;
}

TEST(VJP, TanhBackwardTest) {
  pir::IrContext* ctx = pir::IrContext::Instance();
  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
  pir::Program program((ctx));
  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);

  std::shared_ptr<pir::Builder> builder =
      paddle::dialect::ApiBuilder::Instance().GetBuilder();
  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{1}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
  paddle::dialect::TanhOp op2 =
      builder->Build<paddle::dialect::TanhOp>(op1.out());

  paddle::dialect::FullOp op3 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{1}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());

  std::vector<std::vector<bool>> stop_gradients{{false}};
  std::vector<std::vector<pir::Value>> inputs{{op1.out()}};
  std::vector<std::vector<pir::Value>> outputs{{op2.out()}};
  std::vector<std::vector<pir::Value>> out_grads{{op3.out()}};

  pir::OpInfo op2_info = ctx->GetRegisteredOpInfo("pd_op.tanh");
  auto tanh_vjp_interface_impl =
      op2_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();
  tanh_vjp_interface_impl->vjp_(
      op2.operation(), inputs, outputs, out_grads, stop_gradients);

  builder->Build<pir::ShadowOutputOp>(op2->result(0), "tanh_out");
  builder->Build<pir::ShadowOutputOp>(
      GetOpFromProgram("pd_op.tanh_grad", program)->result(0), "tanh_grad_out");

  auto kernel_program = pir::PdOpLowerToKernelPass(&program);

  auto place = phi::CPUPlace();
  Scope scope;

  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);
  test_core.SetSkipGcVars({"tanh_out", "tanh_grad_out"});
  test_core.Run({});
  auto out_tensor = test_core.local_scope() == nullptr
                        ? scope.FindVar("tanh_out")->Get<phi::DenseTensor>()
                        : test_core.local_scope()
                              ->FindVar("tanh_out")
                              ->Get<phi::DenseTensor>();
  auto grad_out_tensor =
      test_core.local_scope() == nullptr
          ? scope.FindVar("tanh_grad_out")->Get<phi::DenseTensor>()
          : test_core.local_scope()
                ->FindVar("tanh_grad_out")
                ->Get<phi::DenseTensor>();

  ASSERT_NEAR(out_tensor.data<float>()[0], 0.76159, 1e-5);
  ASSERT_NEAR(grad_out_tensor.data<float>()[0], 0.83995, 1e-5);
}

TEST(VJP, Tanh_BackwardTest) {
  pir::IrContext* ctx = pir::IrContext::Instance();
  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
  pir::Program program((ctx));
  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);

  std::shared_ptr<pir::Builder> builder =
      paddle::dialect::ApiBuilder::Instance().GetBuilder();
  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{1}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
  paddle::dialect::Tanh_Op op2 =
      builder->Build<paddle::dialect::Tanh_Op>(op1.out());

  paddle::dialect::FullOp op3 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{1}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());

  std::vector<std::vector<bool>> stop_gradients{{false}};
  std::vector<std::vector<pir::Value>> inputs{{op1.out()}};
  std::vector<std::vector<pir::Value>> outputs{{op2.out()}};
  std::vector<std::vector<pir::Value>> out_grads{{op3.out()}};

  pir::OpInfo op2_info = ctx->GetRegisteredOpInfo("pd_op.tanh_");
  auto tanh_vjp_interface_impl =
      op2_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();
  tanh_vjp_interface_impl->vjp_(
      op2.operation(), inputs, outputs, out_grads, stop_gradients);

  std::string tanh_out = "tanh_out";
  std::string tanh_grad_out = "tanh_grad_out";
  builder->Build<pir::ShadowOutputOp>(op2->result(0), tanh_out);
  builder->Build<pir::ShadowOutputOp>(
      GetOpFromProgram("pd_op.tanh_grad", program)->result(0), tanh_grad_out);

  auto kernel_program = pir::PdOpLowerToKernelPass(&program);

  auto place = phi::CPUPlace();
  Scope scope;

  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);
  test_core.SetSkipGcVars({tanh_out, tanh_grad_out});
  test_core.Run({});
  auto out_tensor =
      test_core.local_scope() == nullptr
          ? scope.FindVar(tanh_out)->Get<phi::DenseTensor>()
          : test_core.local_scope()->FindVar(tanh_out)->Get<phi::DenseTensor>();
  auto grad_out_tensor =
      test_core.local_scope() == nullptr
          ? scope.FindVar(tanh_grad_out)->Get<phi::DenseTensor>()
          : test_core.local_scope()
                ->FindVar(tanh_grad_out)
                ->Get<phi::DenseTensor>();

  ASSERT_NEAR(out_tensor.data<float>()[0], 0.76159, 1e-5);
  ASSERT_NEAR(grad_out_tensor.data<float>()[0], 0.83995, 1e-5);
}

TEST(VJP, MeanBackwardTest) {
  pir::IrContext* ctx = pir::IrContext::Instance();
  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
  pir::Program program((ctx));
  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);

  std::shared_ptr<pir::Builder> builder =
      paddle::dialect::ApiBuilder::Instance().GetBuilder();
  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{2, 2}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
  paddle::dialect::MeanOp op2 =
      builder->Build<paddle::dialect::MeanOp>(op1.out());

  paddle::dialect::FullOp op3 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());

  std::vector<std::vector<bool>> stop_gradients{{false}};
  std::vector<std::vector<pir::Value>> inputs{{op1.out()},
                                              {op2.operand_source(1)}};
  std::vector<std::vector<pir::Value>> outputs{{op2.out()}};
  std::vector<std::vector<pir::Value>> out_grads{{op3.out()}};

  pir::OpInfo op2_info = ctx->GetRegisteredOpInfo("pd_op.mean");
  auto mean_vjp_interface_impl =
      op2_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();
  mean_vjp_interface_impl->vjp_(
      op2.operation(), inputs, outputs, out_grads, stop_gradients);

  builder->Build<pir::ShadowOutputOp>(op2->result(0), "mean_out");
  builder->Build<pir::ShadowOutputOp>(
      GetOpFromProgram("pd_op.mean_grad", program)->result(0), "mean_grad_out");

  auto kernel_program = pir::PdOpLowerToKernelPass(&program);

  auto place = phi::CPUPlace();
  Scope scope;

  ProgramDesc prog_desc;
  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);
  test_core.SetSkipGcVars({"mean_out", "mean_grad_out"});
  test_core.Run({});
  auto out_tensor = test_core.local_scope() == nullptr
                        ? scope.FindVar("mean_out")->Get<phi::DenseTensor>()
                        : test_core.local_scope()
                              ->FindVar("mean_out")
                              ->Get<phi::DenseTensor>();
  auto grad_out_tensor =
      test_core.local_scope() == nullptr
          ? scope.FindVar("mean_grad_out")->Get<phi::DenseTensor>()
          : test_core.local_scope()
                ->FindVar("mean_grad_out")
                ->Get<phi::DenseTensor>();
  ASSERT_EQ(out_tensor.data<float>()[0], 2.0);
  ASSERT_EQ(grad_out_tensor.data<float>()[0], 0.25);
  ASSERT_EQ(grad_out_tensor.data<float>()[1], 0.25);
  ASSERT_EQ(grad_out_tensor.data<float>()[2], 0.25);
  ASSERT_EQ(grad_out_tensor.data<float>()[3], 0.25);
}

TEST(VJP, ConcatBackwardTest) {
  pir::IrContext* ctx = pir::IrContext::Instance();
  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
  pir::Program program((ctx));
  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);

  std::shared_ptr<pir::Builder> builder =
      paddle::dialect::ApiBuilder::Instance().GetBuilder();
  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{1, 2}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
  std::vector<pir::Value> combine_input{{op1.out(), op1.out()}};
  pir::CombineOp op2 = builder->Build<pir::CombineOp>(combine_input);
  paddle::dialect::ConcatOp op3 =
      builder->Build<paddle::dialect::ConcatOp>(op2.out(), 0);

  paddle::dialect::FullOp op4 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{2, 2}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
  std::vector<std::vector<bool>> stop_gradients{{false, false}};
  std::vector<std::vector<pir::Value>> inputs{{op1.out(), op1.out()},
                                              {op3.axis()}};
  std::vector<std::vector<pir::Value>> outputs{{op3.out()}};
  std::vector<std::vector<pir::Value>> out_grads{{op4.out()}};
  pir::OpInfo op2_info = ctx->GetRegisteredOpInfo("pd_op.concat");
  auto concat_vjp_interface_impl =
      op2_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();
  concat_vjp_interface_impl->vjp_(
      op3.operation(), inputs, outputs, out_grads, stop_gradients);

  builder->Build<pir::ShadowOutputOp>(op3->result(0), "concat_out");
  builder->Build<pir::ShadowOutputOp>(
      GetOpFromProgram("builtin.split", program)->result(0), "split_out_0");
  builder->Build<pir::ShadowOutputOp>(
      GetOpFromProgram("builtin.split", program)->result(1), "split_out_1");

  auto kernel_program = pir::PdOpLowerToKernelPass(&program);

  auto place = phi::CPUPlace();
  Scope scope;

  ProgramDesc prog_desc;
  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);
  test_core.SetSkipGcVars({"concat_out", "split_out_0", "split_out_1"});
  test_core.Run({});
  auto out_tensor = test_core.local_scope() == nullptr
                        ? scope.FindVar("concat_out")->Get<phi::DenseTensor>()
                        : test_core.local_scope()
                              ->FindVar("concat_out")
                              ->Get<phi::DenseTensor>();
  auto grad_out_tensor_0 =
      test_core.local_scope() == nullptr
          ? scope.FindVar("split_out_0")->Get<phi::DenseTensor>()
          : test_core.local_scope()
                ->FindVar("split_out_0")
                ->Get<phi::DenseTensor>();
  auto grad_out_tensor_1 =
      test_core.local_scope() == nullptr
          ? scope.FindVar("split_out_1")->Get<phi::DenseTensor>()
          : test_core.local_scope()
                ->FindVar("split_out_1")
                ->Get<phi::DenseTensor>();
  ASSERT_EQ(out_tensor.data<float>()[0], 2.0);
  ASSERT_EQ(grad_out_tensor_0.data<float>()[0], 1.0);
  ASSERT_EQ(grad_out_tensor_0.data<float>()[1], 1.0);
  ASSERT_EQ(grad_out_tensor_1.data<float>()[0], 1.0);
  ASSERT_EQ(grad_out_tensor_1.data<float>()[1], 1.0);
}

TEST(VJP, AddBackwardTest) {
  pir::IrContext* ctx = pir::IrContext::Instance();
  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
  pir::Program program((ctx));
  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);

  std::shared_ptr<pir::Builder> builder =
      paddle::dialect::ApiBuilder::Instance().GetBuilder();
  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{1}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
  paddle::dialect::FullOp op2 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{1}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
  paddle::dialect::AddOp op3 =
      builder->Build<paddle::dialect::AddOp>(op1.out(), op2.out());

  paddle::dialect::FullOp op4 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{1}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());

  std::vector<std::vector<bool>> stop_gradients{{false}, {false}};
  std::vector<std::vector<pir::Value>> inputs{{op1.out()}, {op2.out()}};
  std::vector<std::vector<pir::Value>> outputs{{op3.out()}};
  std::vector<std::vector<pir::Value>> out_grads{{op4.out()}};

  pir::OpInfo op3_info = ctx->GetRegisteredOpInfo("pd_op.add");
  auto add_vjp_interface_impl =
      op3_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();
  add_vjp_interface_impl->vjp_(
      op3.operation(), inputs, outputs, out_grads, stop_gradients);

  builder->Build<pir::ShadowOutputOp>(op3->result(0), "add_out");
  builder->Build<pir::ShadowOutputOp>(
      GetOpFromProgram("pd_op.add_grad", program)->result(0), "add_grad_out_0");
  builder->Build<pir::ShadowOutputOp>(
      GetOpFromProgram("pd_op.add_grad", program)->result(1), "add_grad_out_1");

  auto kernel_program = pir::PdOpLowerToKernelPass(&program);

  auto place = phi::CPUPlace();
  Scope scope;

  ProgramDesc prog_desc;
  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);
  test_core.SetSkipGcVars({"add_out", "add_grad_out_0", "add_grad_out_1"});
  test_core.Run({});
  auto out_tensor = test_core.local_scope() == nullptr
                        ? scope.FindVar("add_out")->Get<phi::DenseTensor>()
                        : test_core.local_scope()
                              ->FindVar("add_out")
                              ->Get<phi::DenseTensor>();
  auto dx = test_core.local_scope() == nullptr
                ? scope.FindVar("add_grad_out_0")->Get<phi::DenseTensor>()
                : test_core.local_scope()
                      ->FindVar("add_grad_out_0")
                      ->Get<phi::DenseTensor>();

  auto dy = test_core.local_scope() == nullptr
                ? scope.FindVar("add_grad_out_1")->Get<phi::DenseTensor>()
                : test_core.local_scope()
                      ->FindVar("add_grad_out_1")
                      ->Get<phi::DenseTensor>();
  ASSERT_EQ(out_tensor.data<float>()[0], 4.0);
  ASSERT_EQ(dx.data<float>()[0], 1.0);
  ASSERT_EQ(dy.data<float>()[0], 1.0);
}

TEST(VJP, Add_BackwardTest) {
  pir::IrContext* ctx = pir::IrContext::Instance();
  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
  pir::Program program((ctx));
  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);

  std::shared_ptr<pir::Builder> builder =
      paddle::dialect::ApiBuilder::Instance().GetBuilder();
  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{1}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
  paddle::dialect::FullOp op2 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{1}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
  paddle::dialect::Add_Op op3 =
      builder->Build<paddle::dialect::Add_Op>(op1.out(), op2.out());

  paddle::dialect::FullOp op4 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{1}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());

  std::vector<std::vector<bool>> stop_gradients{{false}, {false}};
  std::vector<std::vector<pir::Value>> inputs{{op1.out()}, {op2.out()}};
  std::vector<std::vector<pir::Value>> outputs{{op3.out()}};
  std::vector<std::vector<pir::Value>> out_grads{{op4.out()}};

  pir::OpInfo op3_info = ctx->GetRegisteredOpInfo("pd_op.add_");
  auto add_inplace_vjp_interface_impl =
      op3_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();
  add_inplace_vjp_interface_impl->vjp_(
      op3.operation(), inputs, outputs, out_grads, stop_gradients);

  builder->Build<pir::ShadowOutputOp>(op1->result(0), "full_op1_out");
  builder->Build<pir::ShadowOutputOp>(
      GetOpFromProgram("pd_op.add_grad", program)->result(0), "add_grad_out_0");
  builder->Build<pir::ShadowOutputOp>(
      GetOpFromProgram("pd_op.add_grad", program)->result(1), "add_grad_out_1");

  auto kernel_program = pir::PdOpLowerToKernelPass(&program);

  auto place = phi::CPUPlace();
  Scope scope;

  ProgramDesc prog_desc;
  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);

  test_core.SetSkipGcVars({"full_op1_out", "add_grad_out_0", "add_grad_out_1"});
  test_core.Run({});
  auto out_tensor = test_core.local_scope() == nullptr
                        ? scope.FindVar("full_op1_out")->Get<phi::DenseTensor>()
                        : test_core.local_scope()
                              ->FindVar("full_op1_out")
                              ->Get<phi::DenseTensor>();
  auto dx = test_core.local_scope() == nullptr
                ? scope.FindVar("add_grad_out_0")->Get<phi::DenseTensor>()
                : test_core.local_scope()
                      ->FindVar("add_grad_out_0")
                      ->Get<phi::DenseTensor>();

  auto dy = test_core.local_scope() == nullptr
                ? scope.FindVar("add_grad_out_1")->Get<phi::DenseTensor>()
                : test_core.local_scope()
                      ->FindVar("add_grad_out_1")
                      ->Get<phi::DenseTensor>();
  ASSERT_EQ(out_tensor.data<float>()[0], 4.0);
  ASSERT_EQ(dx.data<float>()[0], 1.0);
  ASSERT_EQ(dy.data<float>()[0], 1.0);
}

TEST(VJP, SplitBackwardTest) {
  pir::IrContext* ctx = pir::IrContext::Instance();
  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
  pir::Program program((ctx));
  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);

  std::shared_ptr<pir::Builder> builder =
      paddle::dialect::ApiBuilder::Instance().GetBuilder();
  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{2, 2}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());

  paddle::dialect::SplitOp op2 = builder->Build<paddle::dialect::SplitOp>(
      op1.out(), std::vector<int64_t>{1, 1}, 0);

  pir::SplitOp op3 = builder->Build<pir::SplitOp>(op2.out());

  paddle::dialect::FullOp op4 = builder->Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{1, 2}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());

  std::vector<std::vector<bool>> stop_gradients{{false}};
  std::vector<std::vector<pir::Value>> inputs{
      {op2.x()}, {op2.sections()}, {op2.axis()}};
  std::vector<std::vector<pir::Value>> outputs(1);
  std::vector<pir::Value> res;
  for (uint32_t i = 0; i < op3.outputs().size(); i++) {
    res.push_back(op3.outputs()[i]);
  }
  outputs[0] = res;
  std::vector<std::vector<pir::Value>> out_grads{{op3.result(0), op4.out()}};
  pir::OpInfo op2_info = ctx->GetRegisteredOpInfo("pd_op.split");

  auto concat_vjp_interface_impl =
      op2_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();

  concat_vjp_interface_impl->vjp_(
      op2.operation(), inputs, outputs, out_grads, stop_gradients);

  std::string split_out1 = "split_out1";
  std::string split_out2 = "split_out2";
  std::string concat_out = "concat_out";

  builder->Build<pir::ShadowOutputOp>(op3->result(0), split_out1);
  builder->Build<pir::ShadowOutputOp>(op3->result(1), split_out2);
  builder->Build<pir::ShadowOutputOp>(
      GetOpFromProgram("pd_op.concat", program)->result(0), concat_out);

  auto kernel_program = pir::PdOpLowerToKernelPass(&program);

  auto place = phi::CPUPlace();
  Scope scope;

  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);

  test_core.SetSkipGcVars({split_out1, split_out2, concat_out});
  test_core.Run({});
  auto out_tensor_0 = test_core.local_scope() == nullptr
                          ? scope.FindVar(split_out1)->Get<phi::DenseTensor>()
                          : test_core.local_scope()
                                ->FindVar(split_out1)
                                ->Get<phi::DenseTensor>();
  auto out_tensor_1 = test_core.local_scope() == nullptr
                          ? scope.FindVar(split_out2)->Get<phi::DenseTensor>()
                          : test_core.local_scope()
                                ->FindVar(split_out2)
                                ->Get<phi::DenseTensor>();
  auto grad_out_tensor_0 =
      test_core.local_scope() == nullptr
          ? scope.FindVar(concat_out)->Get<phi::DenseTensor>()
          : test_core.local_scope()
                ->FindVar(concat_out)
                ->Get<phi::DenseTensor>();
  ASSERT_EQ(out_tensor_0.data<float>()[0], 2.0);
  ASSERT_EQ(out_tensor_0.data<float>()[1], 2.0);
  ASSERT_EQ(out_tensor_1.data<float>()[0], 2.0);
  ASSERT_EQ(out_tensor_1.data<float>()[1], 2.0);
  ASSERT_EQ(grad_out_tensor_0.data<float>()[0], 2.0);
  ASSERT_EQ(grad_out_tensor_0.data<float>()[1], 2.0);
  ASSERT_EQ(grad_out_tensor_0.data<float>()[2], 1.0);
  ASSERT_EQ(grad_out_tensor_0.data<float>()[3], 1.0);
}

}  // namespace framework
}  // namespace paddle