/* Copyright (c) 2022 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 "paddle/phi/common/scalar.h"
#include "test/cpp/auto_parallel/spmd_rule_test_util.h"

namespace paddle {
namespace distributed {
namespace auto_parallel {

TEST(MatmulSPMDRule, Ctor) {
  // build input data class
  std::vector<int64_t> x_shape = {64, 32};
  std::vector<int64_t> y_shape = {32, 48};

  std::vector<int64_t> mesh_shape = {2, 4};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({1, -1}));
  x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  TensorDistAttr y_dist_attr = TensorDistAttr();
  y_dist_attr.set_process_mesh(process_mesh);
  y_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1}));
  y_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  size_t input_size = 2;
  size_t output_size = 1;

  phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor y(common::make_ddim(y_shape), y_dist_attr);

  auto matmul_spmd_rule =
      phi::distributed::SpmdRuleFactory::Instance().GetSpmdRule("matmul");

  // mk[1, -1],kn[-1, -1] --> mk[1, -1],kn[-1, -1] = nm[1, -1] partial[]
  phi::distributed::InferSpmdContext ctx(
      {x, y}, {/*trans_x=*/false, /*trans_y=*/false});
  auto inferred_dist_attrs = matmul_spmd_rule.InferForward(ctx);

  EXPECT_EQ(inferred_dist_attrs.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs.second.size(), output_size);
  check_dim_mapping(inferred_dist_attrs.first[0], {1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {-1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), false);
  VLOG(4) << "test1 done." << std::endl << std::endl << std::endl;

  // mk[-1,-1],kn[-1,0] --> mk[-1,-1],kn[-1,0] = nm[-1,0] partial[]
  x_dist_attr.set_dims_mapping({-1, -1});
  y_dist_attr.set_dims_mapping({-1, 0});
  x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
  y = phi::distributed::DistMetaTensor(common::make_ddim(y_shape), y_dist_attr);
  ctx = phi::distributed::InferSpmdContext(
      {x, y}, {/*trans_x=*/false, /*trans_x=*/false});
  inferred_dist_attrs = matmul_spmd_rule.InferForward(ctx);
  check_dim_mapping(inferred_dist_attrs.first[0], {-1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {-1, 0});
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, 0});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), false);
  VLOG(4) << "test2 done." << std::endl << std::endl << std::endl;
  // mk[1, 0],kn[-1,-1] --> mk[1, 0],kn[0, -1] = nm[1, -1] partial[0]: done
  x_dist_attr.set_dims_mapping({1, 0});
  y_dist_attr.set_dims_mapping({-1, -1});
  x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
  y = phi::distributed::DistMetaTensor(common::make_ddim(y_shape), y_dist_attr);
  ctx = phi::distributed::InferSpmdContext(
      {x, y}, {/*trans_x=*/false, /*trans_x=*/false});
  inferred_dist_attrs = matmul_spmd_rule.InferForward(ctx);
  check_dim_mapping(inferred_dist_attrs.first[0], {1, 0});
  check_dim_mapping(inferred_dist_attrs.first[1], {0, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), true);
  check_partial_dims(inferred_dist_attrs.second[0], {0});
  VLOG(4) << "test3 done." << std::endl << std::endl << std::endl;

  // mk[-1,-1],kn[1,0] --> mk[-1, 1],kn[1, 0] = nm[-1, 0] partial[1]: done
  x_dist_attr.set_dims_mapping({-1, -1});
  y_dist_attr.set_dims_mapping({1, 0});
  x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
  y = phi::distributed::DistMetaTensor(common::make_ddim(y_shape), y_dist_attr);
  ctx = phi::distributed::InferSpmdContext(
      {x, y}, {/*trans_x=*/false, /*trans_x=*/false});
  inferred_dist_attrs = matmul_spmd_rule.InferForward(ctx);
  check_dim_mapping(inferred_dist_attrs.first[0], {-1, 1});
  check_dim_mapping(inferred_dist_attrs.first[1], {1, 0});
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, 0});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), true);
  check_partial_dims(inferred_dist_attrs.second[0], {1});
  VLOG(4) << "test4 done." << std::endl << std::endl << std::endl;

  // abcmk[1, 0, -1, -1],kn[-1, -1] --> abcmk[1, 0, -1, -1],kn[-1, -1] =
  // abcmn[1, 0, -1, -1] partial[]: done
  x_shape = {512, 48, 64, 32};
  x_dist_attr.set_dims_mapping({0, 1, -1, -1});
  y_dist_attr.set_dims_mapping({-1, -1});
  x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
  y = phi::distributed::DistMetaTensor(common::make_ddim(y_shape), y_dist_attr);
  ctx = phi::distributed::InferSpmdContext(
      {x, y}, {/*trans_x=*/false, /*trans_x=*/false});
  inferred_dist_attrs = matmul_spmd_rule.InferForward(ctx);
  check_dim_mapping(inferred_dist_attrs.first[0], {0, 1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {-1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {0, 1, -1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), false);
  VLOG(4) << "test5 done." << std::endl << std::endl << std::endl;

  // abcmk[1, -1, -1, 0],kn[-1, -1] --> abcmk[1, -1, -1, 0],kn[0, -1] = abcmn[1,
  // -1, -1, -1] partial[0]: done
  x_dist_attr.set_dims_mapping({1, -1, -1, 0});
  y_dist_attr.set_dims_mapping({-1, -1});
  x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
  y = phi::distributed::DistMetaTensor(common::make_ddim(y_shape), y_dist_attr);
  ctx = phi::distributed::InferSpmdContext(
      {x, y}, {/*trans_x=*/false, /*trans_x=*/false});
  inferred_dist_attrs = matmul_spmd_rule.InferForward(ctx);
  check_dim_mapping(inferred_dist_attrs.first[0], {1, -1, -1, 0});
  check_dim_mapping(inferred_dist_attrs.first[1], {0, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {1, -1, -1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), true);
  check_partial_dims(inferred_dist_attrs.second[0], {0});
  VLOG(4) << "test6 done." << std::endl << std::endl << std::endl;

  // abcmk[1, -1, -1, 0], kn[-1, -1] --> abcmk[1, -1, -1, 0],kn[-1, -1] =
  // abcmn[1, -1, 0, -1] partial[]: done
  x_dist_attr.set_dims_mapping({1, -1, -1, 0});
  y_dist_attr.set_dims_mapping({-1, -1});
  y_shape = {64, 48};
  x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
  y = phi::distributed::DistMetaTensor(common::make_ddim(y_shape), y_dist_attr);
  ctx = phi::distributed::InferSpmdContext(
      {x, y}, {/*trans_x=*/true, /*trans_x=*/false});
  inferred_dist_attrs = matmul_spmd_rule.InferForward(ctx);

  check_dim_mapping(inferred_dist_attrs.first[0], {1, -1, -1, 0});
  check_dim_mapping(inferred_dist_attrs.first[1], {-1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {1, -1, 0, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), false);

  VLOG(4) << "test7 done." << std::endl << std::endl << std::endl;

  // abcmk[-1, -1, -1, -1], kn[1, 0] --> abcmk[-1, -1, -1, 0],kn[1, 0] =
  // abcmn[-1, -1, -1, 1] partial[0]: done
  x_dist_attr.set_dims_mapping({-1, -1, -1, -1});
  y_dist_attr.set_dims_mapping({1, 0});
  y_shape = {48, 32};
  x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
  y = phi::distributed::DistMetaTensor(common::make_ddim(y_shape), y_dist_attr);
  ctx = phi::distributed::InferSpmdContext(
      {x, y}, {/*trans_x=*/false, /*trans_x=*/true});
  inferred_dist_attrs = matmul_spmd_rule.InferForward(ctx);
  check_dim_mapping(inferred_dist_attrs.first[0], {-1, -1, -1, 0});
  check_dim_mapping(inferred_dist_attrs.first[1], {1, 0});
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, -1, -1, 1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), true);
  check_partial_dims(inferred_dist_attrs.second[0], {0});
  clean_partial_dims(&inferred_dist_attrs.second[0], {0});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), false);
  VLOG(4) << "test8 done." << std::endl << std::endl << std::endl;

  // abcmk[-1, -1, 0, 1]+trans_x=true, kn[1, 0]+trans_y=true --> abcmk[-1, -1,
  // 0, -1],kn[-1, 0] = abcmn[-1, -1, 1, -1] partial[0]: done
  x_dist_attr.set_dims_mapping({-1, -1, 0, 1});
  y_dist_attr.set_dims_mapping({1, 0});
  y_shape = {48, 64};
  x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
  y = phi::distributed::DistMetaTensor(common::make_ddim(y_shape), y_dist_attr);
  ctx = phi::distributed::InferSpmdContext(
      {x, y}, {/*trans_x=*/true, /*trans_x=*/true});
  inferred_dist_attrs = matmul_spmd_rule.InferForward(ctx);

  check_dim_mapping(inferred_dist_attrs.first[0], {-1, -1, 0, -1});
  check_dim_mapping(inferred_dist_attrs.first[1],
                    {1, 0});  // conflict and should be changed to [1, 0]
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, -1, -1, 1});
  check_partial_dims(inferred_dist_attrs.second[0], {0});

  clean_partial_status(&inferred_dist_attrs.second[0]);
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), false);
  EXPECT_ANY_THROW(set_partial_status(&inferred_dist_attrs.second[0], {1}));
  VLOG(4) << "test9 done." << std::endl << std::endl << std::endl;

  // abcmk[-1, -1, 1, 0], kn[1, 0] --> abcmk[-1, -1, -1, 0],kn[1, 0] =
  // abcmn[-1, -1, -1, 1] partial[0]: done
  x_dist_attr.set_dims_mapping({-1, -1, 1, 0});
  y_dist_attr.set_dims_mapping({1, 0});
  x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
  y = phi::distributed::DistMetaTensor(common::make_ddim(y_shape), y_dist_attr);
  ctx = phi::distributed::InferSpmdContext(
      {x, y}, {/*trans_x=*/true, /*trans_x=*/true});
  inferred_dist_attrs = matmul_spmd_rule.InferForward(ctx);
  check_dim_mapping(inferred_dist_attrs.first[0], {-1, -1, -1, 0});
  check_dim_mapping(inferred_dist_attrs.first[1],
                    {1, -1});  // conflict and should be changed to [1, -1]
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, -1, 0, 1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), false);
  VLOG(4) << "test10 done." << std::endl << std::endl << std::endl;

  // abcmk[-1, -1, 1, 0], kn[0, 1] --> abcmk[-1, -1, 1, 0],kn[0, 1] =
  // abcmn[-1, -1, 1, -1] partial[0]:
  x_dist_attr.set_dims_mapping({-1, -1, 0, 1});
  y_dist_attr.set_dims_mapping({1, 0});
  x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
  y = phi::distributed::DistMetaTensor(common::make_ddim(y_shape), y_dist_attr);
  ctx = phi::distributed::InferSpmdContext(
      {x, y}, {/*trans_x=*/true, /*trans_x=*/true});
  inferred_dist_attrs = matmul_spmd_rule.InferForward(ctx);
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, -1, -1, 1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), true);
  check_partial_dims(inferred_dist_attrs.second[0], {0});

  // try to clean partial on a dim which is not partial
  EXPECT_ANY_THROW(clean_partial_dims(&inferred_dist_attrs.second[0], {1}));
  // try to clean partial on a dims which is sharded
  EXPECT_ANY_THROW(set_partial_status(&inferred_dist_attrs.second[0], {1}));

  // clean partial and then re-set again
  clean_partial_dims(&inferred_dist_attrs.second[0], {0});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), false);
  set_partial_status(&inferred_dist_attrs.second[0], {0});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), true);
  check_partial_dims(inferred_dist_attrs.second[0], {0});
  VLOG(4) << "test11 done." << std::endl << std::endl << std::endl;
}

TEST(IndexPut, Ctor) {
  // build input data class
  std::vector<int64_t> x_shape = {64, 64, 64};
  std::vector<int64_t> indice_shape = {32};
  std::vector<int64_t> value_shape = {32, 64};
  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 0, 1}));
  x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));

  TensorDistAttr value_dist_attr = TensorDistAttr();
  value_dist_attr.set_process_mesh(process_mesh);
  value_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1}));
  value_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
  TensorDistAttr indice_dist_attr = TensorDistAttr();
  indice_dist_attr.set_process_mesh(process_mesh);
  indice_dist_attr.set_dims_mapping(std::vector<int64_t>({-1}));
  indice_dist_attr.set_dynamic_dims(std::vector<bool>({false}));

  // Test forward.
  // [-1,0, 1], [[-1],[-1]], [-1,-1] --> [-1,-1, 1]
  // infer input:[-1,-1, 1], [[-1],[-1]], [-1,1]
  phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor value(common::make_ddim(value_shape),
                                         value_dist_attr);
  std::vector<phi::distributed::DistMetaTensor> indices;
  for (int i = 0; i < 2; ++i) {
    phi::distributed::DistMetaTensor indice(common::make_ddim(indice_shape),
                                            indice_dist_attr);
    indices.push_back(indice);
  }
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::IndexPutInferSpmd(x, indices, value);
  size_t input_size = 3;
  size_t output_size = 1;
  EXPECT_EQ(forward_info.first.size(), input_size);
  EXPECT_EQ(forward_info.second.size(), output_size);
  check_dim_mapping(forward_info.first[0], {-1, -1, 1});
  std::vector<TensorDistAttr> indices_dist_attr =
      paddle::get<1>(forward_info.first[1]);
  for (auto& attr : indices_dist_attr) {
    check_dim_mapping(attr, {-1});
  }

  check_dim_mapping(forward_info.first[2], {-1, 1});
  check_dim_mapping(forward_info.second[0], {-1, -1, 1});
  VLOG(4) << "test forward done.";

  // Test backward.
  // [-1,0, 1], [[-1],[-1]], [-1,-1],[-1,0, 1] --> [-1,-1, 1], [-1,1]
  // infer input:[-1,-1, 1], [[-1],[-1]], [-1,1],[-1,-1, 1]
  phi::distributed::DistMetaTensor out_grad(common::make_ddim(x_shape),
                                            x_dist_attr);

  phi::distributed::SpmdInfo backward_info =
      phi::distributed::IndexPutGradInferSpmd(x, indices, value, out_grad);
  input_size = 4;
  output_size = 2;
  EXPECT_EQ(backward_info.first.size(), input_size);
  EXPECT_EQ(backward_info.second.size(), output_size);
  check_dim_mapping(backward_info.first[0], {-1, -1, 1});
  indices_dist_attr = paddle::get<1>(backward_info.first[1]);
  for (auto& attr : indices_dist_attr) {
    check_dim_mapping(attr, {-1});
  }

  check_dim_mapping(backward_info.first[2], {-1, 1});
  check_dim_mapping(backward_info.first[3], {-1, -1, 1});
  check_dim_mapping(backward_info.second[0], {-1, -1, 1});
  check_dim_mapping(backward_info.second[1], {-1, 1});
  VLOG(4) << "test backward done.";
}
TEST(InstanceNorm, Ctor) {
  // build input data class
  std::vector<int64_t> x_shape = {64, 64, 64, 64};  // N,C,H,W
  std::vector<int64_t> scale_shape = {64};
  std::vector<int64_t> bias_shape = {64};
  std::vector<int64_t> mean_and_variance_shape = {64, 64};
  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1, -1}));
  x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false, false}));

  TensorDistAttr scale_dist_attr = TensorDistAttr();
  scale_dist_attr.set_process_mesh(process_mesh);
  scale_dist_attr.set_dims_mapping(std::vector<int64_t>({-1}));
  scale_dist_attr.set_dynamic_dims(std::vector<bool>({false}));

  TensorDistAttr bias_dist_attr = TensorDistAttr();
  bias_dist_attr.set_process_mesh(process_mesh);
  bias_dist_attr.set_dims_mapping(std::vector<int64_t>({-1}));
  bias_dist_attr.set_dynamic_dims(std::vector<bool>({false}));

  float epsilon = 1e-5;
  // Test forward.
  // [-1,0, 1, -1], [-1], [-1] --> [-1,0, -1, -1], [-1,0], [-1,0]

  x_dist_attr.set_dims_mapping({-1, 0, 1, -1});
  scale_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
  bias_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
  phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor scale(common::make_ddim(scale_shape),
                                         scale_dist_attr);
  phi::distributed::DistMetaTensor bias(common::make_ddim(bias_shape),
                                        bias_dist_attr);
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::InstanceNormInferSpmd(x, scale, bias, epsilon);
  size_t input_size = 3;
  size_t output_size = 3;
  EXPECT_EQ(forward_info.first.size(), input_size);
  EXPECT_EQ(forward_info.second.size(), output_size);
  check_dim_mapping(forward_info.first[0], {-1, 0, -1, -1});
  check_dim_mapping(forward_info.first[1], {-1});
  check_dim_mapping(forward_info.first[2], {-1});
  check_dim_mapping(forward_info.second[0], {-1, 0, -1, -1});
  check_dim_mapping(forward_info.second[1], {-1, 0});
  check_dim_mapping(forward_info.second[2], {-1, 0});
  VLOG(4) << "test forward done.";

  // Test backward.
  // [-1,0, 1, -1], [-1], [-1,-1], [-1,-1], [-1,0, 1,
  // -1]-->[-1,0,-1,-1],[-1],[-1]

  TensorDistAttr mean_and_variance_dist_attr = TensorDistAttr();
  mean_and_variance_dist_attr.set_process_mesh(process_mesh);
  mean_and_variance_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1}));
  mean_and_variance_dist_attr.set_dynamic_dims(
      std::vector<bool>({false, false}));
  phi::distributed::DistMetaTensor y_grad(common::make_ddim(x_shape),
                                          x_dist_attr);
  phi::distributed::DistMetaTensor saved_mean(
      common::make_ddim(mean_and_variance_shape), mean_and_variance_dist_attr);
  phi::distributed::DistMetaTensor saved_variance(
      common::make_ddim(mean_and_variance_shape), mean_and_variance_dist_attr);
  phi::distributed::SpmdInfo backward_info =
      phi::distributed::InstanceNormGradInferSpmd(
          x, scale, bias, saved_mean, saved_variance, y_grad, epsilon);
  input_size = 5;
  output_size = 3;
  EXPECT_EQ(backward_info.first.size(), input_size);
  EXPECT_EQ(backward_info.second.size(), output_size);
  check_dim_mapping(backward_info.first[0], {-1, 0, -1, -1});
  check_dim_mapping(backward_info.first[1], {-1});
  check_dim_mapping(backward_info.first[2], {-1, 0});
  check_dim_mapping(backward_info.first[3], {-1, 0});
  check_dim_mapping(backward_info.first[4], {-1, 0, -1, -1});
  check_dim_mapping(backward_info.second[0], {-1, 0, -1, -1});
  check_dim_mapping(backward_info.second[1], {-1});
  check_dim_mapping(backward_info.second[2], {-1});
  VLOG(4) << "test backward done.";
}
TEST(LayerNormSPMDRule, Ctor) {
  // build input data class
  std::vector<int64_t> x_shape = {64, 32, 1024};
  std::vector<int64_t> scale_shape = {1024};
  std::vector<int64_t> bias_shape = {1024};

  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({1, -1, -1}));
  x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));

  TensorDistAttr scale_dist_attr = TensorDistAttr();
  scale_dist_attr.set_process_mesh(process_mesh);
  scale_dist_attr.set_dims_mapping(std::vector<int64_t>({-1}));
  scale_dist_attr.set_dynamic_dims(std::vector<bool>({false}));

  TensorDistAttr bias_dist_attr = TensorDistAttr();
  bias_dist_attr.set_process_mesh(process_mesh);
  bias_dist_attr.set_dims_mapping(std::vector<int64_t>({-1}));
  bias_dist_attr.set_dynamic_dims(std::vector<bool>({false}));

  double epsilon = 1e-5;
  int begin_norm_axis = 2;

  auto layer_norm_rule =
      phi::distributed::SpmdRuleFactory::Instance().GetSpmdRule("layer_norm");

  // ijk[1, -1, -1], k[-1], k[-1] --> ijk[1, -1, -1], z[1], z[1], z=ij,
  // begin_norm_axis=2
  begin_norm_axis = 2;
  x_dist_attr.set_dims_mapping({1, -1, -1});
  scale_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
  bias_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
  phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor scale(common::make_ddim(scale_shape),
                                         scale_dist_attr);
  phi::distributed::DistMetaTensor bias(common::make_ddim(bias_shape),
                                        bias_dist_attr);
  phi::distributed::InferSpmdContext ctx({x, scale, bias},
                                         {epsilon, begin_norm_axis});
  auto inferred_dist_attrs = layer_norm_rule.InferForward(ctx);

  size_t input_size = 3;
  size_t output_size = 3;
  EXPECT_EQ(inferred_dist_attrs.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs.second.size(), output_size);
  check_dim_mapping(inferred_dist_attrs.first[0], {1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {-1});
  check_dim_mapping(inferred_dist_attrs.first[2], {-1});
  check_dim_mapping(inferred_dist_attrs.second[0], {1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[1], {1, -1});
  check_dim_mapping(inferred_dist_attrs.second[2], {1, -1});
  VLOG(4) << "test1 done.";

  // ijk[1, 0, -1],k[0],k[0] --> ijk[1, -1, -1],z[1, 0],z[1, 0],
  // begin_norm_axis=2
  begin_norm_axis = 2;
  x_dist_attr.set_dims_mapping({1, 0, -1});
  scale_dist_attr.set_dims_mapping(std::vector<int64_t>{0});
  bias_dist_attr.set_dims_mapping(std::vector<int64_t>{0});
  x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
  scale = phi::distributed::DistMetaTensor(common::make_ddim(scale_shape),
                                           scale_dist_attr);
  bias = phi::distributed::DistMetaTensor(common::make_ddim(bias_shape),
                                          bias_dist_attr);
  ctx = phi::distributed::InferSpmdContext({x, scale, bias},
                                           {epsilon, begin_norm_axis});
  inferred_dist_attrs = layer_norm_rule.InferForward(ctx);

  check_dim_mapping(inferred_dist_attrs.first[0], {1, 0, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {-1});
  check_dim_mapping(inferred_dist_attrs.first[2], {-1});
  check_dim_mapping(inferred_dist_attrs.second[0], {1, 0, -1});
  check_dim_mapping(inferred_dist_attrs.second[1], {1, 0});
  check_dim_mapping(inferred_dist_attrs.second[2], {1, 0});
  VLOG(4) << "test2 done.";

  // ijk[0, -1, -1],y[-1],y[1] --> ijk[0, -1, -1], i[0], i[0], y=jk,
  // begin_norm_axis=1
  begin_norm_axis = 1;
  x_dist_attr.set_dims_mapping({0, -1, -1});
  scale_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
  bias_dist_attr.set_dims_mapping(std::vector<int64_t>{1});
  x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
  scale = phi::distributed::DistMetaTensor(common::make_ddim(scale_shape),
                                           scale_dist_attr);
  bias = phi::distributed::DistMetaTensor(common::make_ddim(bias_shape),
                                          bias_dist_attr);
  ctx = phi::distributed::InferSpmdContext({x, scale, bias},
                                           {epsilon, begin_norm_axis});
  inferred_dist_attrs = layer_norm_rule.InferForward(ctx);

  check_dim_mapping(inferred_dist_attrs.first[0], {0, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {-1});
  check_dim_mapping(inferred_dist_attrs.first[2], {-1});
  check_dim_mapping(inferred_dist_attrs.second[0], {0, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[1], {0});
  check_dim_mapping(inferred_dist_attrs.second[2], {0});
  VLOG(4) << "test3 done.";
}

TEST(MatmulSPMDRuleInferBackward, Ctor) {
  // build input data class
  std::vector<int64_t> x_shape = {512, 1024, 64, 32};
  std::vector<int64_t> y_shape = {512, 1, 32, 48};
  std::vector<int64_t> out_shape = {512, 1024, 64, 48};

  std::vector<int64_t> mesh_shape = {2, 4};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping(
      std::vector<int64_t>({-1, 1, 0, -1}));  // no affect
  x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  TensorDistAttr y_dist_attr = TensorDistAttr();
  y_dist_attr.set_process_mesh(process_mesh);
  y_dist_attr.set_dims_mapping(
      std::vector<int64_t>({0, 1, -1, -1}));  // no affect
  y_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  TensorDistAttr out_dist_attr = TensorDistAttr();
  out_dist_attr.set_process_mesh(process_mesh);
  out_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, 1, -1}));
  out_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
  out_dist_attr.set_partial_status(std::vector<int64_t>({0}));

  phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor y(common::make_ddim(y_shape), y_dist_attr);
  phi::distributed::DistMetaTensor out(common::make_ddim(out_shape),
                                       out_dist_attr);

  auto matmul_spmd_rule =
      phi::distributed::SpmdRuleFactory::Instance().GetSpmdRule("matmul");

  // TODO(zyc) update in future: propagate the partial in inferbackward
  // abmn[-1, -1, 1, -1] + partial[0] --> abmk[-1, -1, 1, -1], a1kn[-1, -1, -1,
  // -1]
  phi::distributed::InferSpmdContext ctx(
      {x, y, out}, {/*trans_x=*/false, /*trans_x=*/false});
  auto inferred_dist_attrs = matmul_spmd_rule.InferBackward(ctx);

  size_t input_size = 2;
  size_t output_size = 1;
  EXPECT_EQ(inferred_dist_attrs.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs.second.size(), output_size);

  check_dim_mapping(inferred_dist_attrs.first[0], {-1, -1, 1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, -1, 1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.first[0]), false);
  EXPECT_EQ(is_partial(inferred_dist_attrs.first[1]), false);
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), true);
  VLOG(4) << "test1 done." << std::endl << std::endl << std::endl;
}
TEST(GroupNorm, Ctor) {
  // build input data class
  std::vector<int64_t> x_shape = {64, 64, 64, 64};  // N,C,H,W
  std::vector<int64_t> scale_shape = {64};
  std::vector<int64_t> bias_shape = {64};
  std::vector<int64_t> mean_and_variance_shape = {64, 64};
  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1, -1}));
  x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false, false}));

  TensorDistAttr scale_dist_attr = TensorDistAttr();
  scale_dist_attr.set_process_mesh(process_mesh);
  scale_dist_attr.set_dims_mapping(std::vector<int64_t>({-1}));
  scale_dist_attr.set_dynamic_dims(std::vector<bool>({false}));

  TensorDistAttr bias_dist_attr = TensorDistAttr();
  bias_dist_attr.set_process_mesh(process_mesh);
  bias_dist_attr.set_dims_mapping(std::vector<int64_t>({-1}));
  bias_dist_attr.set_dynamic_dims(std::vector<bool>({false}));

  float epsilon = 1e-5;
  int groups = -1;
  std::string data_format = "NCHW";
  // Test forward.
  // [0, 1, -1, -1], [-1], [-1] --> input: [0, -1, -1, -1], [-1], [-1]
  // output:[0,-1, -1, -1], [0], [0]

  x_dist_attr.set_dims_mapping({0, 1, -1, -1});
  scale_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
  bias_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
  phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor scale(common::make_ddim(scale_shape),
                                         scale_dist_attr);
  phi::distributed::DistMetaTensor bias(common::make_ddim(bias_shape),
                                        bias_dist_attr);
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::GroupNormInferSpmd(
          x, scale, bias, epsilon, groups, data_format);
  size_t input_size = 3;
  size_t output_size = 3;
  EXPECT_EQ(forward_info.first.size(), input_size);
  EXPECT_EQ(forward_info.second.size(), output_size);
  check_dim_mapping(forward_info.first[0], {0, -1, -1, -1});
  check_dim_mapping(forward_info.first[1], {-1});
  check_dim_mapping(forward_info.first[2], {-1});
  check_dim_mapping(forward_info.second[0], {0, -1, -1, -1});
  check_dim_mapping(forward_info.second[1], {0, -1});
  check_dim_mapping(forward_info.second[2], {0, -1});
  VLOG(4) << "test forward done.";

  // Test backward.
  // [0, 1, -1, -1], [-1],[-1], [0, 1, -1, -1],[-1], [-1], [0, 1, -1, -1]
  // infer input:[0, -1, -1, -1], [-1],[-1], [0, -1, -1, -1],[0], [0], [0, -1,
  // -1, -1]
  // infer output:[0, 1, -1, -1], [-1],[-1]
  TensorDistAttr mean_and_variance_dist_attr = TensorDistAttr();
  mean_and_variance_dist_attr.set_process_mesh(process_mesh);
  mean_and_variance_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1}));
  mean_and_variance_dist_attr.set_dynamic_dims(std::vector<bool>({false}));
  phi::distributed::DistMetaTensor y(common::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor y_grad(common::make_ddim(x_shape),
                                          x_dist_attr);
  phi::distributed::DistMetaTensor mean(
      common::make_ddim(mean_and_variance_shape), mean_and_variance_dist_attr);
  phi::distributed::DistMetaTensor variance(
      common::make_ddim(mean_and_variance_shape), mean_and_variance_dist_attr);
  phi::distributed::SpmdInfo backward_info =
      phi::distributed::GroupNormGradInferSpmd(x,
                                               scale,
                                               bias,
                                               y,
                                               mean,
                                               variance,
                                               y_grad,
                                               epsilon,
                                               groups,
                                               data_format);
  input_size = 7;
  output_size = 3;
  EXPECT_EQ(backward_info.first.size(), input_size);
  EXPECT_EQ(backward_info.second.size(), output_size);
  check_dim_mapping(backward_info.first[0], {0, -1, -1, -1});
  check_dim_mapping(backward_info.first[1], {-1});
  check_dim_mapping(backward_info.first[2], {-1});
  check_dim_mapping(backward_info.first[3], {0, -1, -1, -1});
  check_dim_mapping(backward_info.first[4], {0, -1});
  check_dim_mapping(backward_info.first[5], {0, -1});
  check_dim_mapping(backward_info.first[6], {0, -1, -1, -1});
  check_dim_mapping(backward_info.second[0], {0, -1, -1, -1});
  check_dim_mapping(backward_info.second[1], {-1});
  check_dim_mapping(backward_info.second[2], {-1});
  VLOG(4) << "test backward done.";
}
TEST(ReplicatedSPMDRule, Ctor) {
  // build input data class
  std::vector<int64_t> x_shape = {10, 10, 32, 48};
  std::vector<int64_t> y_shape = {32, 48};
  std::vector<int64_t> out1_shape = {10, 10, 32, 48};
  std::vector<int64_t> out2_shape = {10, 32, 48};

  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 1, -1, -1}));
  x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  TensorDistAttr y_dist_attr = TensorDistAttr();
  y_dist_attr.set_process_mesh(process_mesh);
  y_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1}));  // no affect
  y_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  TensorDistAttr out1_dist_attr = TensorDistAttr();
  out1_dist_attr.set_process_mesh(process_mesh);
  out1_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, 1, -1}));
  out1_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  TensorDistAttr out2_dist_attr = TensorDistAttr();
  out2_dist_attr.set_process_mesh(process_mesh);
  out2_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 1, -1}));
  out2_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor y(common::make_ddim(y_shape), y_dist_attr);
  phi::distributed::DistMetaTensor out1(common::make_ddim(out1_shape),
                                        out1_dist_attr);
  phi::distributed::DistMetaTensor out2(common::make_ddim(out2_shape),
                                        out2_dist_attr);

  // 2 inputs 2 outputs
  // call in vector arguments format
  auto inferred_dist_attrs_st =
      phi::distributed::ReplicatedInferSpmd({&x, &y}, {&out1, &out2});
  // call in variadic arguments format
  auto inferred_dist_attrs_dy =
      phi::distributed::VariadicReplicatedInferSpmd(x, y, &out1, &out2);

  size_t input_size = 2;
  size_t output_size = 2;
  EXPECT_EQ(inferred_dist_attrs_st.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs_st.second.size(), output_size);
  EXPECT_EQ(inferred_dist_attrs_dy.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs_dy.second.size(), output_size);

  check_dim_mapping(inferred_dist_attrs_st.first[0], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs_st.first[1], {-1, -1});
  check_dim_mapping(inferred_dist_attrs_st.second[0], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs_st.second[1], {-1, -1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs_st.first[0]), false);
  EXPECT_EQ(is_partial(inferred_dist_attrs_st.first[1]), false);
  EXPECT_EQ(is_partial(inferred_dist_attrs_st.second[0]), false);
  EXPECT_EQ(is_partial(inferred_dist_attrs_st.second[1]), false);
  EXPECT_EQ(inferred_dist_attrs_st.first, inferred_dist_attrs_dy.first);
  EXPECT_EQ(inferred_dist_attrs_st.second, inferred_dist_attrs_dy.second);
  VLOG(4) << "test1 done." << std::endl << std::endl << std::endl;

  // 3 inputs 1 outputs
  // call in vector arguments format
  inferred_dist_attrs_st =
      phi::distributed::ReplicatedInferSpmd({&x, &y, &out1}, {&out2});
  // call in variadic arguments format
  inferred_dist_attrs_dy =
      phi::distributed::VariadicReplicatedInferSpmd(x, y, out1, &out2);

  input_size = 3;
  output_size = 1;
  EXPECT_EQ(inferred_dist_attrs_st.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs_st.second.size(), output_size);
  EXPECT_EQ(inferred_dist_attrs_dy.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs_dy.second.size(), output_size);
  check_dim_mapping(inferred_dist_attrs_dy.first[0], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs_dy.first[1], {-1, -1});
  check_dim_mapping(inferred_dist_attrs_dy.first[2], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs_dy.second[0], {-1, -1, -1});
  EXPECT_EQ(inferred_dist_attrs_st.first, inferred_dist_attrs_dy.first);
  EXPECT_EQ(inferred_dist_attrs_st.second, inferred_dist_attrs_dy.second);
  VLOG(4) << "test2 done." << std::endl << std::endl << std::endl;

  // 1 inputs 3 outputs backward
  // call in vector arguments format
  inferred_dist_attrs_st =
      phi::distributed::ReplicatedInferSpmdReverse({&x}, {&y, &out1, &out2});
  // call in variadic arguments format
  inferred_dist_attrs_dy =
      phi::distributed::VariadicReplicatedInferSpmdReverse(x, &y, &out1, &out2);

  input_size = 1;
  output_size = 3;
  EXPECT_EQ(inferred_dist_attrs_st.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs_st.second.size(), output_size);
  EXPECT_EQ(inferred_dist_attrs_dy.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs_dy.second.size(), output_size);

  check_dim_mapping(inferred_dist_attrs_dy.first[0], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs_dy.second[0], {-1, -1});
  check_dim_mapping(inferred_dist_attrs_dy.second[1], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs_dy.second[2], {-1, -1, -1});
  EXPECT_EQ(inferred_dist_attrs_st.first, inferred_dist_attrs_dy.first);
  EXPECT_EQ(inferred_dist_attrs_st.second, inferred_dist_attrs_dy.second);
  VLOG(4) << "test3 done." << std::endl << std::endl << std::endl;
}

TEST(DefaultDataParallelSPMDRule, Ctor) {
  // build input data class
  std::vector<int64_t> x_shape = {10, 10, 32, 48};
  std::vector<int64_t> y_shape = {32, 48};
  std::vector<int64_t> out1_shape = {10, 10, 32, 48};
  std::vector<int64_t> out2_shape = {10, 32, 48};

  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, -1, 1}));
  x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  TensorDistAttr y_dist_attr = TensorDistAttr();
  y_dist_attr.set_process_mesh(process_mesh);
  y_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1}));
  y_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  TensorDistAttr out1_dist_attr = TensorDistAttr();
  out1_dist_attr.set_process_mesh(process_mesh);
  out1_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, 1, -1}));
  out1_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  TensorDistAttr out2_dist_attr = TensorDistAttr();
  out2_dist_attr.set_process_mesh(process_mesh);
  out2_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 1, -1}));
  out2_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor y(common::make_ddim(y_shape), y_dist_attr);
  phi::distributed::DistMetaTensor out1(common::make_ddim(out1_shape),
                                        out1_dist_attr);
  phi::distributed::DistMetaTensor out2(common::make_ddim(out2_shape),
                                        out2_dist_attr);

  // 2 inputs 2 outputs, batch axis sharding is propagated while other axes are
  // replicated call in vector arguments format
  auto inferred_dist_attrs_st =
      phi::distributed::DefaultDataParallelInferSpmd({&x, &y}, {&out1, &out2});
  // call in variadic arguments format
  auto inferred_dist_attrs_dy =
      phi::distributed::VariadicDefaultDataParallelInferSpmd(
          x, y, &out1, &out2);

  size_t input_size = 2;
  size_t output_size = 2;
  EXPECT_EQ(inferred_dist_attrs_st.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs_st.second.size(), output_size);
  EXPECT_EQ(inferred_dist_attrs_dy.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs_dy.second.size(), output_size);
  check_dim_mapping(inferred_dist_attrs_st.first[0], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs_st.first[1], {0, -1});
  check_dim_mapping(inferred_dist_attrs_st.second[0], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs_st.second[1], {0, -1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs_st.first[0]), false);
  EXPECT_EQ(is_partial(inferred_dist_attrs_st.first[1]), false);
  EXPECT_EQ(is_partial(inferred_dist_attrs_st.second[0]), false);
  EXPECT_EQ(is_partial(inferred_dist_attrs_st.second[1]), false);

  EXPECT_EQ(inferred_dist_attrs_st.first, inferred_dist_attrs_dy.first);
  EXPECT_EQ(inferred_dist_attrs_st.second, inferred_dist_attrs_dy.second);
  VLOG(4) << "test1 done." << std::endl << std::endl << std::endl;

  // 1 inputs 3 outputs, batch axis is un-sharded
  // call in vector arguments format
  inferred_dist_attrs_st =
      phi::distributed::DefaultDataParallelInferSpmd({&x}, {&y, &out1, &out2});
  // call in variadic arguments format
  inferred_dist_attrs_dy =
      phi::distributed::VariadicDefaultDataParallelInferSpmd(
          x, &y, &out1, &out2);

  input_size = 1;
  output_size = 3;
  EXPECT_EQ(inferred_dist_attrs_st.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs_st.second.size(), output_size);
  EXPECT_EQ(inferred_dist_attrs_dy.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs_dy.second.size(), output_size);

  check_dim_mapping(inferred_dist_attrs_dy.first[0], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs_dy.second[0], {-1, -1});
  check_dim_mapping(inferred_dist_attrs_dy.second[1], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs_dy.second[2], {-1, -1, -1});

  EXPECT_EQ(inferred_dist_attrs_st.first, inferred_dist_attrs_dy.first);
  EXPECT_EQ(inferred_dist_attrs_st.second, inferred_dist_attrs_dy.second);
  VLOG(4) << "test2 done." << std::endl << std::endl << std::endl;

  // conflict on batch axis
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1, -1}));
  y_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1}));
  out1_dist_attr.set_dims_mapping(std::vector<int64_t>({1, -1, -1, -1}));
  x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
  y = phi::distributed::DistMetaTensor(common::make_ddim(y_shape), y_dist_attr);
  out1 = phi::distributed::DistMetaTensor(common::make_ddim(out1_shape),
                                          out1_dist_attr);

  EXPECT_ANY_THROW(inferred_dist_attrs_st =
                       phi::distributed::DefaultDataParallelInferSpmd(
                           {&x, &y, &out1}, {&out2}));
  // call in variadic arguments format
  EXPECT_ANY_THROW(inferred_dist_attrs_dy =
                       phi::distributed::VariadicDefaultDataParallelInferSpmd(
                           x, y, out1, &out2));

  VLOG(4) << "test3 done." << std::endl << std::endl << std::endl;

  // 2 inputs 2 outputs, backward
  // call in vector arguments format
  out1_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 0, 1, -1}));
  out2_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1}));
  out1 = phi::distributed::DistMetaTensor(common::make_ddim(out1_shape),
                                          out1_dist_attr);
  out2 = phi::distributed::DistMetaTensor(common::make_ddim(out2_shape),
                                          out2_dist_attr);

  inferred_dist_attrs_st =
      phi::distributed::DefaultDataParallelInferSpmdReverse({&x, &y},
                                                            {&out1, &out2});
  // call in variadic arguments format
  inferred_dist_attrs_dy =
      phi::distributed::VariadicDefaultDataParallelInferSpmdReverse(
          x, y, &out1, &out2);

  input_size = 2;
  output_size = 2;
  EXPECT_EQ(inferred_dist_attrs_st.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs_st.second.size(), output_size);
  EXPECT_EQ(inferred_dist_attrs_dy.first.size(), input_size);
  EXPECT_EQ(inferred_dist_attrs_dy.second.size(), output_size);
  check_dim_mapping(inferred_dist_attrs_dy.first[0], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs_dy.first[1], {0, -1});
  check_dim_mapping(inferred_dist_attrs_dy.second[0], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs_dy.second[1], {0, -1, -1});
  EXPECT_EQ(inferred_dist_attrs_st.first, inferred_dist_attrs_dy.first);
  EXPECT_EQ(inferred_dist_attrs_st.second, inferred_dist_attrs_dy.second);
  VLOG(4) << "test4 done." << std::endl << std::endl << std::endl;
}
TEST(ConcatRule, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  std::vector<std::vector<int64_t>> shapes = {
      {16, 16, 16}, {4, 16, 16}, {2, 16, 16}};
  std::vector<std::vector<int64_t>> dim_mappings = {
      {-1, 0, 1}, {-1, 1, 0}, {-1, -1, 0}};
  std::vector<std::vector<int64_t>> partial_status = {{}, {}, {1}};

  auto build_inputs = [&] {
    std::vector<phi::distributed::DistMetaTensor> inputs;
    for (int i = 0; i < 3; i++) {
      auto t_dist_attr = TensorDistAttr();
      t_dist_attr.set_process_mesh(process_mesh);
      t_dist_attr.set_dims_mapping(dim_mappings[i]);
      t_dist_attr.set_dynamic_dims({false, false, false});
      auto input = phi::distributed::DistMetaTensor(
          common::make_ddim(shapes[i]), t_dist_attr);
      inputs.push_back(input);
    }
    return inputs;
  };

  // test 1, inputs are aligned according to cost, and partial status is cleared
  auto inputs = build_inputs();
  auto inferred_dist_attrs = phi::distributed::ConcatInferSpmd(inputs, 0);
  // list of tensor => single tensor
  EXPECT_EQ(inferred_dist_attrs.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(inferred_dist_attrs.second.size(), static_cast<size_t>(1));
  EXPECT_TRUE(
      paddle::holds_alternative<std::vector<phi::distributed::TensorDistAttr>>(
          inferred_dist_attrs.first[0]));
  EXPECT_TRUE(paddle::holds_alternative<phi::distributed::TensorDistAttr>(
      inferred_dist_attrs.second[0]));
  auto& inputs_infer1 = PADDLE_GET_CONST(std::vector<TensorDistAttr>,
                                         inferred_dist_attrs.first[0]);
  for (auto e : inputs_infer1) {
    check_dim_mapping(e, {-1, 1, 0});
    check_partial_dims(e, {});
  }
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, 1, 0});
  check_partial_dims(inferred_dist_attrs.second[0], {});

  auto build_output = [&](const TensorDistAttr& t_dist_attr,
                          const std::vector<int64_t>& shape) {
    return phi::distributed::DistMetaTensor(common::make_ddim(shape),
                                            t_dist_attr);
  };

  auto& output_dist_attr =
      PADDLE_GET_CONST(TensorDistAttr, inferred_dist_attrs.second[0]);
  auto output = build_output(output_dist_attr, {22, 16, 16});
  // test reverse
  auto inferred_reverse_attrs =
      phi::distributed::ConcatInferSpmdReverse(inputs, output, 0);
  auto& inputs_infer1_reverse = PADDLE_GET_CONST(
      std::vector<TensorDistAttr>, inferred_reverse_attrs.first[0]);
  for (auto e : inputs_infer1_reverse) {
    check_dim_mapping(e, {-1, 1, 0});
    check_partial_dims(e, {});
  }
  check_dim_mapping(inferred_reverse_attrs.second[0],
                    output_dist_attr.dims_mapping());
  // test grad
  auto inferred_grad_attrs =
      phi::distributed::ConcatGradInferSpmdDynamic(inputs, output, 0);
  auto& inputs_infer1_grad = PADDLE_GET_CONST(std::vector<TensorDistAttr>,
                                              inferred_grad_attrs.first[0]);
  for (auto e : inputs_infer1_grad) {
    check_dim_mapping(e, {-1, 1, 0});
    check_partial_dims(e, {});
  }
  check_dim_mapping(inferred_grad_attrs.first[1],
                    output_dist_attr.dims_mapping());
  auto& inferred_grad = PADDLE_GET_CONST(std::vector<TensorDistAttr>,
                                         inferred_grad_attrs.second[0]);
  for (auto e : inferred_grad) {
    check_dim_mapping(e, {-1, 1, 0});
    check_partial_dims(e, {});
  }

  // test 2，force replicate along concat axis
  inputs = build_inputs();
  inferred_dist_attrs = phi::distributed::ConcatInferSpmd(inputs, 1);
  // list of tensor => single tensor
  EXPECT_EQ(inferred_dist_attrs.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(inferred_dist_attrs.second.size(), static_cast<size_t>(1));
  EXPECT_TRUE(
      paddle::holds_alternative<std::vector<phi::distributed::TensorDistAttr>>(
          inferred_dist_attrs.first[0]));
  EXPECT_TRUE(paddle::holds_alternative<phi::distributed::TensorDistAttr>(
      inferred_dist_attrs.second[0]));
  auto& inputs_infer2 = PADDLE_GET_CONST(std::vector<TensorDistAttr>,
                                         inferred_dist_attrs.first[0]);
  for (auto e : inputs_infer2) {
    check_dim_mapping(e, {1, -1, 0});
    check_partial_dims(e, {});
  }
  check_dim_mapping(inferred_dist_attrs.second[0], {1, -1, 0});
  check_partial_dims(inferred_dist_attrs.second[0], {});

  // test 3，special case: concat one dimensional tensor
  shapes = {{16}, {32}, {64}};
  dim_mappings = {{0}, {1}, {-1}};
  partial_status = {{}, {}, {1}};
  inputs = build_inputs();
  inferred_dist_attrs = phi::distributed::ConcatInferSpmd(inputs, 0);
  // list of tensor => single tensor
  EXPECT_EQ(inferred_dist_attrs.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(inferred_dist_attrs.second.size(), static_cast<size_t>(1));
  EXPECT_TRUE(
      paddle::holds_alternative<std::vector<phi::distributed::TensorDistAttr>>(
          inferred_dist_attrs.first[0]));
  EXPECT_TRUE(paddle::holds_alternative<phi::distributed::TensorDistAttr>(
      inferred_dist_attrs.second[0]));
  auto& inputs_infer3 = PADDLE_GET_CONST(std::vector<TensorDistAttr>,
                                         inferred_dist_attrs.first[0]);
  for (auto e : inputs_infer3) {
    check_dim_mapping(e, {-1});
    check_partial_dims(e, {});
  }
  check_dim_mapping(inferred_dist_attrs.second[0], {-1});
  check_partial_dims(inferred_dist_attrs.second[0], {});
}

TEST(StackRule, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  int input_size = 3;
  std::vector<int64_t> input_shape = {16, 8, 4};
  std::vector<std::vector<int64_t>> dim_mappings = {
      {-1, 0, 1}, {-1, 1, 0}, {-1, -1, 0}};
  std::vector<std::vector<int64_t>> partial_status = {{}, {}, {1}};

  auto build_inputs = [&] {
    std::vector<phi::distributed::DistMetaTensor> inputs;
    for (int i = 0; i < input_size; i++) {
      auto t_dist_attr = TensorDistAttr();
      t_dist_attr.set_process_mesh(process_mesh);
      t_dist_attr.set_dims_mapping(dim_mappings[i]);
      t_dist_attr.set_dynamic_dims({false, false, false});
      auto input = phi::distributed::DistMetaTensor(
          common::make_ddim(input_shape), t_dist_attr);
      inputs.push_back(input);
    }
    return inputs;
  };

  auto build_output = [&](const TensorDistAttr& t_dist_attr, int stack_dim) {
    std::vector<int64_t> output_shape;
    std::transform(input_shape.begin(),
                   input_shape.begin() + stack_dim,
                   std::back_inserter(output_shape),
                   [](int64_t x) { return x; });
    output_shape.push_back(input_size);
    std::transform(input_shape.begin() + stack_dim,
                   input_shape.end(),
                   std::back_inserter(output_shape),
                   [](int64_t x) { return x; });
    return phi::distributed::DistMetaTensor(common::make_ddim(output_shape),
                                            t_dist_attr);
  };

  // test 1, inputs are aligned according to cost.
  auto inputs = build_inputs();
  auto inferred_dist_attrs = phi::distributed::StackInferSpmd(inputs, 0);
  // list of tensor => single tensor
  EXPECT_EQ(inferred_dist_attrs.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(inferred_dist_attrs.second.size(), static_cast<size_t>(1));
  EXPECT_TRUE(
      paddle::holds_alternative<std::vector<phi::distributed::TensorDistAttr>>(
          inferred_dist_attrs.first[0]));
  EXPECT_TRUE(paddle::holds_alternative<phi::distributed::TensorDistAttr>(
      inferred_dist_attrs.second[0]));
  auto& inputs_infer1 = PADDLE_GET_CONST(std::vector<TensorDistAttr>,
                                         inferred_dist_attrs.first[0]);
  for (auto e : inputs_infer1) {
    check_dim_mapping(e, {-1, 1, 0});
    check_partial_dims(e, {});
  }
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, -1, 1, 0});
  check_partial_dims(inferred_dist_attrs.second[0], {});

  auto output_dist_attr =
      PADDLE_GET_CONST(TensorDistAttr, inferred_dist_attrs.second[0]);
  auto output = build_output(output_dist_attr, 0);
  // test reverse
  auto inferred_reverse_attrs =
      phi::distributed::StackInferSpmdReverse(inputs, output, 0);
  auto& inputs_infer1_reverse = PADDLE_GET_CONST(
      std::vector<TensorDistAttr>, inferred_reverse_attrs.first[0]);
  for (auto e : inputs_infer1_reverse) {
    check_dim_mapping(e, {-1, 1, 0});
    check_partial_dims(e, {});
  }
  check_dim_mapping(inferred_reverse_attrs.second[0],
                    output_dist_attr.dims_mapping());
  // test grad
  auto inferred_grad_attrs = phi::distributed::StackGradInferSpmd(output, 0);
  check_dim_mapping(inferred_grad_attrs.first[0],
                    output_dist_attr.dims_mapping());
  auto& inferred_grad = PADDLE_GET_CONST(std::vector<TensorDistAttr>,
                                         inferred_grad_attrs.second[0]);
  for (auto e : inferred_grad) {
    check_dim_mapping(e, {-1, 1, 0});
    check_partial_dims(e, {});
  }

  // test 2，force replicate along concat axis
  inputs = build_inputs();
  inferred_dist_attrs = phi::distributed::StackInferSpmd(inputs, 1);
  // list of tensor => single tensor
  EXPECT_EQ(inferred_dist_attrs.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(inferred_dist_attrs.second.size(), static_cast<size_t>(1));
  EXPECT_TRUE(
      paddle::holds_alternative<std::vector<phi::distributed::TensorDistAttr>>(
          inferred_dist_attrs.first[0]));
  EXPECT_TRUE(paddle::holds_alternative<phi::distributed::TensorDistAttr>(
      inferred_dist_attrs.second[0]));
  auto& inputs_infer2 = PADDLE_GET_CONST(std::vector<TensorDistAttr>,
                                         inferred_dist_attrs.first[0]);
  for (auto e : inputs_infer2) {
    check_dim_mapping(e, {-1, 1, 0});
    check_partial_dims(e, {});
  }
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, -1, 1, 0});
  check_partial_dims(inferred_dist_attrs.second[0], {});
}

TEST(WhereRule, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  std::vector<std::vector<int64_t>> shapes = {{16, 16, 16}, {16, 16}, {16}};
  std::vector<std::vector<int64_t>> dim_mappings = {{-1, 0, -1}, {-1, 0}, {-1}};

  auto build_inputs = [&] {
    std::vector<phi::distributed::DistMetaTensor> inputs;
    for (int i = 0; i < 3; i++) {
      auto t_dist_attr = TensorDistAttr();
      t_dist_attr.set_process_mesh(process_mesh);
      t_dist_attr.set_dims_mapping(dim_mappings[i]);
      t_dist_attr.set_dynamic_dims({false, false, false});
      auto input = phi::distributed::DistMetaTensor(
          common::make_ddim(shapes[i]), t_dist_attr);
      inputs.push_back(input);
    }
    return inputs;
  };

  auto inputs = build_inputs();
  auto inferred_dist_attrs = phi::distributed::WhereGradInferSpmd(
      inputs[0], inputs[1], inputs[2], inputs[0]);

  EXPECT_EQ(inferred_dist_attrs.first.size(), static_cast<size_t>(4));
  EXPECT_EQ(inferred_dist_attrs.second.size(), static_cast<size_t>(2));

  check_dim_mapping(inferred_dist_attrs.first[0], {-1, 0, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {0, -1});
  check_dim_mapping(inferred_dist_attrs.first[2], {-1});
  check_dim_mapping(inferred_dist_attrs.first[3], {-1, 0, -1});

  check_dim_mapping(inferred_dist_attrs.second[0], {0, -1});
  check_partial_dims(inferred_dist_attrs.second[0], {});
  check_dim_mapping(inferred_dist_attrs.second[1], {-1});
  check_partial_dims(inferred_dist_attrs.second[1], {0});
}
TEST(ArgminRule, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  // test forward
  auto t_dist_attr = TensorDistAttr();
  t_dist_attr.set_process_mesh(process_mesh);
  t_dist_attr.set_dims_mapping({0, 1, -1});
  t_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
      common::make_ddim({4, 6, 8}), t_dist_attr);
  phi::Scalar axis(1);
  bool keep_dim = false;
  bool flatten = false;
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::ArgMinInferSpmdDynamic(
          x, axis, keep_dim, flatten, phi::DataType::FLOAT32);
  check_dim_mapping(forward_info.first[0], {0, -1, -1});
  check_partial_dims(forward_info.first[0], {});
  check_dim_mapping(forward_info.second[0], {0, -1});
  check_partial_dims(forward_info.second[0], {});
}
TEST(ReduceMaxRule, Ctor) {
  std::vector<int64_t> mesh_shape = {2};
  std::vector<int64_t> process_ids = {0, 1};
  std::vector<std::string> dim_names = {"x"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  // test forward
  auto t_dist_attr = TensorDistAttr();
  t_dist_attr.set_process_mesh(process_mesh);
  t_dist_attr.set_dims_mapping({-1, 0, -1});
  t_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
      common::make_ddim({4, 6, 8}), t_dist_attr);
  phi::IntArray axis = {1};
  bool keep_dim = false;
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::ReductionMaxInferSpmdDynamic(x, axis, keep_dim);
  check_dim_mapping(forward_info.second[0], {-1, -1});
  check_partial_dims(forward_info.second[0], {0});
  // test backward
  phi::distributed::DistMetaTensor out = phi::distributed::DistMetaTensor(
      common::make_ddim({4, 8}),
      PADDLE_GET_CONST(TensorDistAttr, forward_info.second[0]));
  phi::distributed::DistMetaTensor out_grad = out;
  phi::distributed::SpmdInfo backward_info =
      phi::distributed::ReductionGradInferSpmd(
          x, out, out_grad, {1}, false, false);
  check_partial_dims(backward_info.first[1], {});
  check_dim_mapping(backward_info.second[0], {-1, -1, -1});
  check_partial_dims(backward_info.second[0], {});
}

TEST(ReduceMinRule, Ctor) {
  std::vector<int64_t> mesh_shape = {2};
  std::vector<int64_t> process_ids = {0, 1};
  std::vector<std::string> dim_names = {"x"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  // test forward
  auto t_dist_attr = TensorDistAttr();
  t_dist_attr.set_process_mesh(process_mesh);
  t_dist_attr.set_dims_mapping({-1, 0, -1});
  t_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
      common::make_ddim({4, 6, 8}), t_dist_attr);
  phi::IntArray axis = {1};
  bool keep_dim = false;
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::ReductionMinInferSpmdDynamic(x, axis, keep_dim);
  check_dim_mapping(forward_info.second[0], {-1, -1});
  check_partial_dims(forward_info.second[0], {0});
}

TEST(ReduceAllRule, Ctor) {
  std::vector<int64_t> mesh_shape = {2};
  std::vector<int64_t> process_ids = {0, 1};
  std::vector<std::string> dim_names = {"x"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  // test forward
  auto t_dist_attr = TensorDistAttr();
  t_dist_attr.set_process_mesh(process_mesh);
  t_dist_attr.set_dims_mapping({-1, 0, -1});
  t_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor x =
      phi::distributed::DistMetaTensor(phi::make_ddim({4, 6, 8}), t_dist_attr);
  phi::IntArray axis = {1};
  bool keep_dim = false;
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::ReductionAllInferSpmdDynamic(x, axis, keep_dim);
  check_dim_mapping(forward_info.second[0], {-1, -1});
  check_partial_dims(forward_info.second[0], {0});
}

TEST(ReduceAnyRule, Ctor) {
  std::vector<int64_t> mesh_shape = {2};
  std::vector<int64_t> process_ids = {0, 1};
  std::vector<std::string> dim_names = {"x"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  // test forward
  auto t_dist_attr = TensorDistAttr();
  t_dist_attr.set_process_mesh(process_mesh);
  t_dist_attr.set_dims_mapping({-1, 0, -1});
  t_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor x =
      phi::distributed::DistMetaTensor(phi::make_ddim({4, 6, 8}), t_dist_attr);
  phi::IntArray axis = {1};
  bool keep_dim = false;
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::ReductionAnyInferSpmdDynamic(x, axis, keep_dim);
  check_dim_mapping(forward_info.second[0], {-1, -1});
  check_partial_dims(forward_info.second[0], {0});
}

TEST(Numel, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  std::vector<int64_t> shape = {16, 16, 16};
  std::vector<int64_t> dims_mapping = {-1, 0, -1};

  auto t_dist_attr = TensorDistAttr();
  t_dist_attr.set_process_mesh(process_mesh);
  t_dist_attr.set_dims_mapping(dims_mapping);
  t_dist_attr.set_dynamic_dims({false, false, false});
  auto input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  auto inferred_dist_attrs = phi::distributed::NumelInferSpmd(input);
  EXPECT_EQ(inferred_dist_attrs.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(inferred_dist_attrs.second.size(), static_cast<size_t>(1));
  check_dim_mapping(inferred_dist_attrs.first[0], dims_mapping);
  check_dim_mapping(inferred_dist_attrs.second[0], {});
  check_partial_dims(inferred_dist_attrs.second[0], {0});
}

TEST(Triu, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  std::vector<int64_t> shape = {16, 16, 16};
  std::vector<int64_t> dims_mapping = {0, -1, 1};

  auto t_dist_attr = TensorDistAttr();
  t_dist_attr.set_process_mesh(process_mesh);
  t_dist_attr.set_dims_mapping(dims_mapping);
  t_dist_attr.set_dynamic_dims({false, false, false});
  auto input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  auto inferred_dist_attrs = phi::distributed::TriuGradInferSpmd(input, 0);
  EXPECT_EQ(inferred_dist_attrs.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(inferred_dist_attrs.second.size(), static_cast<size_t>(1));
  check_dim_mapping(inferred_dist_attrs.first[0], {0, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {0, -1, -1});
  check_partial_dims(inferred_dist_attrs.second[0], {});
}

TEST(LayerNorm, Ctor) {
  using phi::distributed::PartialStatus;
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  std::vector<int64_t> x_shapes = {16, 32, 32};

  auto build_input = [&](const std::vector<int64_t>& shape,
                         const std::vector<int64_t>& dim_mapping) {
    auto t_dist_attr = TensorDistAttr();
    t_dist_attr.set_process_mesh(process_mesh);
    t_dist_attr.set_dims_mapping(dim_mapping);
    t_dist_attr.set_dynamic_dims({false, false, false});
    auto input =
        phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
    return input;
  };
  // test 1
  auto x = build_input(x_shapes, {0, 1, -1});
  auto out_grad = build_input(x_shapes, {0, 1, -1});
  auto mean = build_input({16, 32}, {0, 1});
  auto variance = build_input({16, 32}, {0, 1});
  auto scale = build_input({32}, {0});
  auto bias = build_input({32}, {0});

  auto spmd1 =
      LayerNormGradInferSpmd(x, scale, bias, mean, variance, out_grad, 1.0, 2);

  EXPECT_EQ(spmd1.first.size(), static_cast<size_t>(6));
  EXPECT_EQ(spmd1.second.size(), static_cast<size_t>(3));

  check_dim_mapping(spmd1.first[0], {0, 1, -1});
  check_dim_mapping(spmd1.first[1], {-1});
  check_dim_mapping(spmd1.first[2], {-1});
  check_dim_mapping(spmd1.first[3], {0, 1});
  check_dim_mapping(spmd1.first[4], {0, 1});
  check_dim_mapping(spmd1.first[5], {0, 1, -1});
  check_dim_mapping(spmd1.second[0], {0, 1, -1});
  check_dim_mapping(spmd1.second[1], {-1});
  check_dim_mapping(spmd1.second[2], {-1});
  check_partial_dims(spmd1.second[1], {0, 1});
  check_partial_dims(spmd1.second[2], {0, 1});
  // test 2
  mean = build_input({16}, {0});
  variance = build_input({16}, {0});
  scale = build_input({32, 32}, {0, 1});
  bias = build_input({32, 32}, {0, 1});
  auto spmd2 =
      LayerNormGradInferSpmd(x, scale, bias, mean, variance, out_grad, 1.0, 1);
  EXPECT_EQ(spmd2.first.size(), static_cast<size_t>(6));
  EXPECT_EQ(spmd2.second.size(), static_cast<size_t>(3));
  check_dim_mapping(spmd2.first[0], {0, -1, -1});
  check_dim_mapping(spmd2.first[1], {-1, -1});
  check_dim_mapping(spmd2.first[2], {-1, -1});
  check_dim_mapping(spmd2.first[3], {0});
  check_dim_mapping(spmd2.first[4], {0});
  check_dim_mapping(spmd2.first[5], {0, -1, -1});
  check_dim_mapping(spmd2.second[0], {0, -1, -1});
  check_dim_mapping(spmd2.second[1], {-1, -1});
  check_dim_mapping(spmd2.second[2], {-1, -1});
  check_partial_dims(spmd2.second[1], {0});
  check_partial_dims(spmd2.second[2], {0});
}

TEST(FlashAtt, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  auto build_input = [&](const std::vector<int64_t>& shape,
                         const std::vector<int64_t>& dim_mapping) {
    auto t_dist_attr = TensorDistAttr();
    t_dist_attr.set_process_mesh(process_mesh);
    t_dist_attr.set_dims_mapping(dim_mapping);
    t_dist_attr.set_dynamic_dims(std::vector<bool>(shape.size(), false));
    auto input =
        phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
    return input;
  };

  // b, s, m, h
  std::vector<int64_t> qkv_shape = {2, 256, 2, 128};
  std::vector<int64_t> dim_mapping = {0, 1, -1, -1};

  auto qkv = build_input(qkv_shape, dim_mapping);
  auto mask = build_input({}, {});
  auto seed_offset = build_input({}, {});

  auto spmd1 = FlashAttInferSpmd(
      qkv, qkv, qkv, seed_offset, mask, 0.5, false, false, false, "");

  EXPECT_EQ(spmd1.first.size(), static_cast<size_t>(5));
  EXPECT_EQ(spmd1.second.size(), static_cast<size_t>(4));
  check_dim_mapping(spmd1.first[0], {0, -1, -1, -1});
  check_dim_mapping(spmd1.first[1], {0, -1, -1, -1});
  check_dim_mapping(spmd1.first[2], {0, -1, -1, -1});
  check_dim_mapping(spmd1.first[3], {});
  check_dim_mapping(spmd1.first[4], {});
  check_dim_mapping(spmd1.second[0], {0, -1, -1, -1});
  check_dim_mapping(spmd1.second[1], {0, -1, -1, -1});
  check_dim_mapping(spmd1.second[2], {0, -1, -1});
  check_dim_mapping(spmd1.second[3], {-1});

  auto out = build_input(qkv_shape, {0, -1, 1, -1});
  auto softmax_lse = build_input({2, 2, 256}, {0, 1, -1});
  auto out_grad = build_input(qkv_shape, {-1, -1, -1, -1});

  auto spmd2 = FlashAttGradInferSpmd(
      qkv, qkv, qkv, out, softmax_lse, seed_offset, mask, out_grad, 0.5, false);

  EXPECT_EQ(spmd2.first.size(), static_cast<size_t>(8));
  EXPECT_EQ(spmd2.second.size(), static_cast<size_t>(3));

  check_dim_mapping(spmd2.first[0], {0, -1, 1, -1});
  check_dim_mapping(spmd2.first[1], {0, -1, 1, -1});
  check_dim_mapping(spmd2.first[2], {0, -1, 1, -1});
  check_dim_mapping(spmd2.first[3], {0, -1, 1, -1});
  check_dim_mapping(spmd2.first[4], {0, 1, -1});
  check_dim_mapping(spmd2.first[5], {});
  check_dim_mapping(spmd2.first[6], {});
  check_dim_mapping(spmd2.first[7], {0, -1, 1, -1});
  check_dim_mapping(spmd2.second[0], {0, -1, 1, -1});
  check_dim_mapping(spmd2.second[1], {0, -1, 1, -1});
  check_dim_mapping(spmd2.second[2], {0, -1, 1, -1});
}

TEST(FlashMask, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  auto build_input = [&](const std::vector<int64_t>& shape,
                         const std::vector<int64_t>& dim_mapping) {
    auto t_dist_attr = TensorDistAttr();
    t_dist_attr.set_process_mesh(process_mesh);
    t_dist_attr.set_dims_mapping(dim_mapping);
    t_dist_attr.set_dynamic_dims(std::vector<bool>(shape.size(), false));
    auto input =
        phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
    return input;
  };

  // b, s, m, h
  std::vector<int64_t> qkv_shape = {2, 256, 2, 128};
  std::vector<int64_t> dim_mapping = {0, 1, -1, -1};

  auto qkv = build_input(qkv_shape, dim_mapping);
  auto mask = build_input({}, {});
  auto seed_offset = build_input({}, {});

  auto spmd1 = FlashMaskInferSpmd(
      qkv, qkv, qkv, seed_offset, mask, 0.5, false, false, false, "");

  EXPECT_EQ(spmd1.first.size(), static_cast<size_t>(5));
  EXPECT_EQ(spmd1.second.size(), static_cast<size_t>(4));
  check_dim_mapping(spmd1.first[0], {0, -1, -1, -1});
  check_dim_mapping(spmd1.first[1], {0, -1, -1, -1});
  check_dim_mapping(spmd1.first[2], {0, -1, -1, -1});
  check_dim_mapping(spmd1.first[3], {});
  check_dim_mapping(spmd1.first[4], {});
  check_dim_mapping(spmd1.second[0], {0, -1, -1, -1});
  check_dim_mapping(spmd1.second[1], {0, -1, -1, -1});
  check_dim_mapping(spmd1.second[2], {0, -1, -1});
  check_dim_mapping(spmd1.second[3], {-1});

  auto out = build_input(qkv_shape, {0, -1, 1, -1});
  auto softmax_lse = build_input({2, 2, 256}, {0, 1, -1});
  auto out_grad = build_input(qkv_shape, {-1, -1, -1, -1});

  auto spmd2 = FlashMaskGradInferSpmd(
      qkv, qkv, qkv, mask, out, softmax_lse, seed_offset, out_grad, 0.5, false);

  EXPECT_EQ(spmd2.first.size(), static_cast<size_t>(8));
  EXPECT_EQ(spmd2.second.size(), static_cast<size_t>(3));

  check_dim_mapping(spmd2.first[0], {0, -1, 1, -1});
  check_dim_mapping(spmd2.first[1], {0, -1, 1, -1});
  check_dim_mapping(spmd2.first[2], {0, -1, 1, -1});
  check_dim_mapping(spmd2.first[3], {});
  check_dim_mapping(spmd2.first[4], {0, -1, 1, -1});
  check_dim_mapping(spmd2.first[5], {0, 1, -1});
  check_dim_mapping(spmd2.first[6], {});
  check_dim_mapping(spmd2.first[7], {0, -1, 1, -1});
  check_dim_mapping(spmd2.second[0], {0, -1, 1, -1});
  check_dim_mapping(spmd2.second[1], {0, -1, 1, -1});
  check_dim_mapping(spmd2.second[2], {0, -1, 1, -1});
}

TEST(Util, Ctor) {
  // test equal test not equal
  using phi::distributed::PartialStatus;
  using phi::distributed::PlacementEqual;
  using phi::distributed::ReplicatedStatus;
  using phi::distributed::ShardStatus;
  auto a = std::make_shared<PartialStatus>(phi::ReduceType::kRedSum);
  auto b = std::make_shared<PartialStatus>(phi::ReduceType::kRedMin);
  EXPECT_TRUE(PlacementEqual(a, a));
  EXPECT_TRUE(!PlacementEqual(a, b));
  auto c = std::make_shared<ShardStatus>(0);
  auto d = std::make_shared<ShardStatus>(1);
  EXPECT_TRUE(!PlacementEqual(a, c));
  EXPECT_TRUE(!PlacementEqual(b, c));
  EXPECT_TRUE(PlacementEqual(c, c));
  EXPECT_TRUE(!PlacementEqual(c, d));
  auto e = std::make_shared<ReplicatedStatus>();
  EXPECT_TRUE(PlacementEqual(e, e));
  EXPECT_TRUE(!PlacementEqual(a, e));
  EXPECT_TRUE(!PlacementEqual(b, e));
  EXPECT_TRUE(!PlacementEqual(c, e));
  EXPECT_TRUE(!PlacementEqual(d, e));
}

TEST(Transpose, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  std::vector<int64_t> shape = {6, 8, 10};
  std::vector<int64_t> dims_mapping = {0, -1, 1};

  TensorDistAttr t_dist_attr = TensorDistAttr();
  t_dist_attr.set_process_mesh(process_mesh);
  t_dist_attr.set_dims_mapping(dims_mapping);
  t_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor x =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  std::vector<int> perm = {1, 2, -3};
  // test forward
  phi::distributed::SpmdInfo forward_spmd_info =
      phi::distributed::TransposeInferSpmd(x, perm);
  EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
  check_dim_mapping(forward_spmd_info.first[0], {0, -1, 1});
  check_dim_mapping(forward_spmd_info.second[0], {-1, 1, 0});
  check_partial_dims(forward_spmd_info.second[0], {});
  // test backward
  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
      common::make_ddim({8, 10, 6}),
      PADDLE_GET_CONST(TensorDistAttr, forward_spmd_info.second[0]));
  phi::distributed::SpmdInfo backward_spmd_info =
      TransposeGradInferSpmd(out_grad, perm);
  EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(1));
  check_dim_mapping(backward_spmd_info.first[0], {-1, 1, 0});
  check_dim_mapping(backward_spmd_info.second[0], {0, -1, 1});
  check_partial_dims(backward_spmd_info.second[0], {});
}

TEST(FusedRope, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  auto build_input = [&](const std::vector<int64_t>& shape,
                         const std::vector<int64_t>& dim_mapping) {
    TensorDistAttr t_dist_attr;
    t_dist_attr.set_process_mesh(process_mesh);
    t_dist_attr.set_dims_mapping(dim_mapping);
    t_dist_attr.set_dynamic_dims(std::vector<bool>(shape.size(), false));
    auto input =
        phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
    return input;
  };

  phi::distributed::DistMetaTensor q =
      build_input({16, 2048, 64, 128}, {0, 1, -1, -1});
  phi::distributed::DistMetaTensor none;

  // 1. test forward
  // 1.1 only q input
  phi::distributed::SpmdInfo forward_spmd_info =
      phi::distributed::FusedRopeInferSpmd(
          q, none, none, none, none, none, false, false);
  EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(6));
  EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(3));
  check_dim_mapping(forward_spmd_info.first[0], {0, -1, -1, -1});
  check_dim_mapping(forward_spmd_info.first[1], {});
  check_dim_mapping(forward_spmd_info.first[2], {});
  check_dim_mapping(forward_spmd_info.first[3], {});
  check_dim_mapping(forward_spmd_info.first[4], {});
  check_dim_mapping(forward_spmd_info.first[5], {});
  check_dim_mapping(forward_spmd_info.second[0], {0, -1, -1, -1});
  check_dim_mapping(forward_spmd_info.second[1], {});
  check_dim_mapping(forward_spmd_info.second[2], {});
  check_partial_dims(forward_spmd_info.second[0], {});

  // 1.2 q, k, sin, cos, position_ids
  phi::distributed::DistMetaTensor k =
      build_input({16, 2048, 64, 128}, {-1, 1, -1, 0});
  phi::distributed::DistMetaTensor sin =
      build_input({1, 2048, 1, 128}, {-1, 0, -1, 1});
  phi::distributed::DistMetaTensor cos =
      build_input({1, 2048, 1, 128}, {-1, 1, -1, -1});
  phi::distributed::DistMetaTensor position_ids =
      build_input({16, 2048}, {0, 1});
  forward_spmd_info = phi::distributed::FusedRopeInferSpmd(
      q, k, none, sin, cos, position_ids, false, false);
  EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(6));
  EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(3));
  check_dim_mapping(forward_spmd_info.first[0], {0, -1, -1, -1});
  check_dim_mapping(forward_spmd_info.first[1], {0, -1, -1, -1});
  check_dim_mapping(forward_spmd_info.first[2], {});
  check_dim_mapping(forward_spmd_info.first[3], {-1, -1, -1, -1});
  check_dim_mapping(forward_spmd_info.first[4], {-1, -1, -1, -1});
  check_dim_mapping(forward_spmd_info.first[5], {0, -1});
  check_dim_mapping(forward_spmd_info.second[0], {0, -1, -1, -1});
  check_dim_mapping(forward_spmd_info.second[1], {0, -1, -1, -1});
  check_dim_mapping(forward_spmd_info.second[2], {});
  check_partial_dims(forward_spmd_info.second[0], {});
  check_partial_dims(forward_spmd_info.second[1], {});
  // 2. test backward
  phi::distributed::SpmdInfo backward_spmd_info =
      FusedRopeGradInferSpmd(sin, cos, position_ids, q, k, none, false, false);
  EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(6));
  EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(3));
  check_dim_mapping(backward_spmd_info.first[0], {-1, -1, -1, -1});
  check_dim_mapping(backward_spmd_info.first[1], {-1, -1, -1, -1});
  check_dim_mapping(backward_spmd_info.first[2], {0, -1});
  check_dim_mapping(backward_spmd_info.first[3], {0, -1, -1, -1});
  check_dim_mapping(backward_spmd_info.first[4], {0, -1, -1, -1});
  check_dim_mapping(backward_spmd_info.first[5], {});
  check_dim_mapping(backward_spmd_info.second[0], {0, -1, -1, -1});
  check_dim_mapping(backward_spmd_info.second[1], {0, -1, -1, -1});
  check_dim_mapping(backward_spmd_info.second[2], {});
  check_partial_dims(backward_spmd_info.second[0], {});
  check_partial_dims(backward_spmd_info.second[1], {});

  // 3. test reverse
  phi::distributed::DistMetaTensor out_q =
      build_input({16, 2048, 64, 128}, {0, 1, -1, -1});
  phi::distributed::DistMetaTensor out_k =
      build_input({16, 2048, 64, 128}, {-1, 1, -1, 0});
  phi::distributed::SpmdInfo reverse_spmd_info = FusedRopeInferSpmdReverse(
      q, k, none, sin, cos, position_ids, out_q, out_k, none, false, false);
  EXPECT_EQ(reverse_spmd_info.first.size(), static_cast<size_t>(6));
  EXPECT_EQ(reverse_spmd_info.second.size(), static_cast<size_t>(3));
  check_dim_mapping(reverse_spmd_info.first[0], {0, -1, -1, -1});
  check_dim_mapping(reverse_spmd_info.first[1], {0, -1, -1, -1});
  check_dim_mapping(reverse_spmd_info.first[2], {});
  check_dim_mapping(reverse_spmd_info.first[3], {-1, -1, -1, -1});
  check_dim_mapping(reverse_spmd_info.first[4], {-1, -1, -1, -1});
  check_dim_mapping(reverse_spmd_info.first[5], {0, -1});
  check_dim_mapping(reverse_spmd_info.second[0], {0, -1, -1, -1});
  check_dim_mapping(reverse_spmd_info.second[1], {0, -1, -1, -1});
  check_dim_mapping(reverse_spmd_info.second[2], {});
}

TEST(Reshape, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  auto build_input = [&](const std::vector<int64_t>& shape,
                         const std::vector<int64_t>& dim_mapping) {
    auto t_dist_attr = TensorDistAttr();
    t_dist_attr.set_process_mesh(process_mesh);
    t_dist_attr.set_dims_mapping(dim_mapping);
    t_dist_attr.set_dynamic_dims(std::vector<bool>(shape.size(), false));
    auto input =
        phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
    return input;
  };

  // b s h; dp , mp
  auto input = build_input({2, 1024, 1024}, {0, 1, -1});
  // [b, s, h] => [b, s, nh, h/nh]
  auto spmd = ReshapeInferSpmd(input, {2, 1024, 4, -1});
  EXPECT_EQ(spmd.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(spmd.second.size(), static_cast<size_t>(1));
  check_dim_mapping(spmd.first[0], {0, 1, -1});
  check_dim_mapping(spmd.second[0], {0, 1, -1, -1});

  auto out_grad = build_input({2, 1024, 4, 1024 / 4}, {-1, -1, -1, -1});
  auto x = build_input({2, 1024, 1024}, {0, 1, -1});
  auto spmd_grad = ReshapeGradInferSpmd(x, out_grad);
  EXPECT_EQ(spmd_grad.first.size(), static_cast<size_t>(2));
  EXPECT_EQ(spmd_grad.second.size(), static_cast<size_t>(1));
  check_dim_mapping(spmd_grad.first[0], {0, 1, -1});
  check_dim_mapping(spmd_grad.first[1], {0, 1, -1, -1});
  check_dim_mapping(spmd_grad.second[0], {0, 1, -1});
}

TEST(ElementwiseUnaryLike, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<std::string> dim_names = {"x", "y", "z"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  std::vector<int64_t> shape = {16, 16, 16};
  std::vector<std::vector<int64_t>> dims_mapping = {{0, 1}, {}, {2}};

  auto t_dist_attr = TensorDistAttr();
  t_dist_attr.set_process_mesh(process_mesh);
  t_dist_attr.set_dims_mapping(dims_mapping);
  t_dist_attr.set_dynamic_dims({false, false, false});

  auto check_element_unary_like = [&dims_mapping](auto& spmd_info) {
    EXPECT_EQ(spmd_info.first.size(), static_cast<size_t>(1));
    EXPECT_EQ(spmd_info.second.size(), static_cast<size_t>(1));
    check_multi_dims_mapping(spmd_info.first[0], dims_mapping);
    check_multi_dims_mapping(spmd_info.second[0], dims_mapping);
    check_partial_dims(spmd_info.second[0], {});
  };

  auto check_element_unary_like_backward = [&dims_mapping](auto& spmd_info) {
    EXPECT_GT(spmd_info.first.size(), static_cast<size_t>(1));
    EXPECT_EQ(spmd_info.second.size(), static_cast<size_t>(1));
    for (auto& dim_mapping : spmd_info.first) {
      check_multi_dims_mapping(dim_mapping, dims_mapping);
    }
    check_multi_dims_mapping(spmd_info.second[0], dims_mapping);
    check_partial_dims(spmd_info.second[0], {});
  };

  // cast
  auto input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  auto inferred_dist_attrs =
      phi::distributed::CastInferSpmd(input, phi::DataType::FLOAT32);

  check_element_unary_like(inferred_dist_attrs);
  // ElementwiseUnaryGradInferSpmd
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::ElementwiseUnaryGradInferSpmd(input);
  check_element_unary_like(inferred_dist_attrs);
  // full like
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs =
      phi::distributed::FullLikeInferSpmd(input, 1.0, phi::DataType::FLOAT32);
  check_element_unary_like(inferred_dist_attrs);
  // empty like
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs =
      phi::distributed::EmptyLikeInferSpmd(input, phi::DataType::FLOAT32);
  check_element_unary_like(inferred_dist_attrs);
  // pow
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::PowInferSpmd(input, 2);
  check_element_unary_like(inferred_dist_attrs);

  // pow backward
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::PowGradInferSpmd(input, input, 2);
  check_element_unary_like_backward(inferred_dist_attrs);

  // scale
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs =
      phi::distributed::ScaleInferSpmd(input, 1.0, 1.0, false);
  check_element_unary_like(inferred_dist_attrs);

  // round
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::RoundInfoSpmd(input, 0);
  check_element_unary_like(inferred_dist_attrs);

  // mish
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::MishInfoSpmd(input, 1.0);
  check_element_unary_like(inferred_dist_attrs);

  // mish backward
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::MishGradInfoSpmd(input, input, 1.0);
  check_element_unary_like_backward(inferred_dist_attrs);

  // elu
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::EluInfoSpmd(input, 1.0);
  check_element_unary_like(inferred_dist_attrs);

  // elu backward
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs =
      phi::distributed::EluGradInfoSpmd(input, input, input, 1.0);
  check_element_unary_like_backward(inferred_dist_attrs);

  // selu
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::SeluInfoSpmd(input, 1.0, 1.0);
  check_element_unary_like(inferred_dist_attrs);

  // selu backward
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs =
      phi::distributed::SeluGradInfoSpmd(input, input, 1.0, 1.0);
  check_element_unary_like_backward(inferred_dist_attrs);

  // celu
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::CeluInfoSpmd(input, 1.0);
  check_element_unary_like(inferred_dist_attrs);

  // celu backward
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::CeluGradInfoSpmd(input, input, 1.0);
  check_element_unary_like_backward(inferred_dist_attrs);

  // stanh
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::StanhInfoSpmd(input, 1.0, 1.0);
  check_element_unary_like(inferred_dist_attrs);

  // stanh backward
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs =
      phi::distributed::StanhGradInfoSpmd(input, input, 1.0, 1.0);
  check_element_unary_like_backward(inferred_dist_attrs);

  // softplus
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::SoftplusInfoSpmd(input, 1.0, 1.0);
  check_element_unary_like(inferred_dist_attrs);

  // softplus backward
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs =
      phi::distributed::SoftplusGradInfoSpmd(input, input, 1.0, 1.0);
  check_element_unary_like_backward(inferred_dist_attrs);

  // softshrink
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::SoftshrinkInfoSpmd(input, 1.0);
  check_element_unary_like(inferred_dist_attrs);

  // softshrink backward
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs =
      phi::distributed::SoftshrinkGradInfoSpmd(input, input, 1.0);
  check_element_unary_like_backward(inferred_dist_attrs);

  // thresholded_relu
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs =
      phi::distributed::ThresholdedReluInfoSpmd(input, 1.0, 1.0);
  check_element_unary_like(inferred_dist_attrs);

  // thresholded_relu backward
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs =
      phi::distributed::ThresholdedReluGradInfoSpmd(input, input, 1.0, 1.0);
  check_element_unary_like_backward(inferred_dist_attrs);

  // logit
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::LogitInfoSpmd(input, 1.0);
  check_element_unary_like(inferred_dist_attrs);

  // logit backward
  input =
      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
  inferred_dist_attrs = phi::distributed::LogitGradInfoSpmd(input, input, 1.0);
  check_element_unary_like_backward(inferred_dist_attrs);
}

TEST(EmbeddingGradInferSpmd, Ctor) {
  // build input data class
  std::vector<int64_t> x_shape = {4, 5};
  std::vector<int64_t> w_shape = {10, 3};
  std::vector<int64_t> out_grad_shape = {4, 5, 3};

  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  // indices is shard, embedding table is replicated,
  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({1, -1}));
  x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  TensorDistAttr w_dist_attr = TensorDistAttr();
  w_dist_attr.set_process_mesh(process_mesh);
  w_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1}));
  w_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  TensorDistAttr out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);
  out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, -1}));
  out_grad_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor w(phi::make_ddim(w_shape), w_dist_attr);
  phi::distributed::DistMetaTensor out_grad(phi::make_ddim(out_grad_shape),
                                            out_grad_dist_attr);

  auto spmdinfo = EmbeddingGradInferSpmd(x, w, out_grad, -1, false);

  EXPECT_EQ(spmdinfo.first.size(), 3UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]), std::vector<int64_t>({1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
            std::vector<int64_t>({-1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[2]),
            std::vector<int64_t>({1, -1, -1}));

  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({-1, -1}));
  EXPECT_EQ(
      PADDLE_GET_CONST(phi::distributed::TensorDistAttr, spmdinfo.second[0])
          .is_partial(),
      true);
  VLOG(4) << "Test EmbeddingGradInferSpmd with sharding indices and "
             "replicating weight"
          << std::endl
          << std::endl
          << std::endl;

  // Indices' rank is greater than 1, x and weight is replicated, out_grad is
  // sharded along axis 1
  x_dist_attr.set_dims_mapping({-1, -1});
  w_dist_attr.set_dims_mapping({-1, 1});
  out_grad_dist_attr.set_dims_mapping({-1, 1, -1});
  x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
  w = phi::distributed::DistMetaTensor(phi::make_ddim(w_shape), w_dist_attr);
  out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(out_grad_shape),
                                              out_grad_dist_attr);

  spmdinfo = EmbeddingGradInferSpmd(x, w, out_grad, -1, false);

  EXPECT_EQ(spmdinfo.first.size(), 3UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]), std::vector<int64_t>({-1, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]), std::vector<int64_t>({-1, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[2]),
            std::vector<int64_t>({-1, 1, -1}));

  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({-1, -1}));
  EXPECT_EQ(
      PADDLE_GET_CONST(phi::distributed::TensorDistAttr, spmdinfo.second[0])
          .is_partial(),
      true);
  VLOG(4) << "Test EmbeddingGradInferSpmd with replicating indices and "
             "sharding weight along col axis."
          << std::endl
          << std::endl
          << std::endl;

  // Indices' rank equals 1, indices and out_grad is sharded.
  x_shape = {5};
  w_shape = {10, 3};
  out_grad_shape = {5, 3};

  x_dist_attr.set_dims_mapping(std::vector<int64_t>({0}));
  w_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1}));
  out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 1}));

  x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
  w = phi::distributed::DistMetaTensor(phi::make_ddim(w_shape), w_dist_attr);
  out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(out_grad_shape),
                                              out_grad_dist_attr);

  spmdinfo = EmbeddingGradInferSpmd(x, w, out_grad, -1, false);

  EXPECT_EQ(spmdinfo.first.size(), 3UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]), std::vector<int64_t>({0}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
            std::vector<int64_t>({-1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[2]), std::vector<int64_t>({0, 1}));

  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({-1, 1}));
  EXPECT_EQ(
      PADDLE_GET_CONST(phi::distributed::TensorDistAttr, spmdinfo.second[0])
          .is_partial(),
      true);
  VLOG(4) << "Test EmbeddingGradInferSpmd with sharding weight and out_grad."
          << std::endl
          << std::endl
          << std::endl;

  x_shape = {12, 16};
  w_shape = {10, 4};
  out_grad_shape = {12, 16, 4};

  x_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1}));
  w_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -0}));
  out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, 0}));

  x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
  w = phi::distributed::DistMetaTensor(phi::make_ddim(w_shape), w_dist_attr);
  out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(out_grad_shape),
                                              out_grad_dist_attr);

  spmdinfo = EmbeddingGradInferSpmd(x, w, out_grad, -1, false);

  EXPECT_EQ(spmdinfo.first.size(), 3UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({-1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]), std::vector<int64_t>({-1, 0}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[2]),
            std::vector<int64_t>({-1, -1, 0}));

  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({-1, 0}));
  EXPECT_EQ(
      PADDLE_GET_CONST(phi::distributed::TensorDistAttr, spmdinfo.second[0])
          .is_partial(),
      false);
  VLOG(4) << "Test EmbeddingGradInferSpmd with sharding weight and out_grad."
          << std::endl
          << std::endl
          << std::endl;
}

TEST(SqueezeGradInferSpmd, Ctor) {
  std::vector<int64_t> x_shape = {1, 32, 1, 48};
  std::vector<int64_t> out_grad_shape = {32, 48};

  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 1, -1, -1}));
  x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false, false}));

  TensorDistAttr out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);
  out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 1}));
  out_grad_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor out_grad(phi::make_ddim(out_grad_shape),
                                            out_grad_dist_attr);

  auto spmdinfo = SqueezeGradInferSpmd(x, out_grad);

  EXPECT_EQ(spmdinfo.first.size(), 2UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({-1, 1, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]), std::vector<int64_t>({-1, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({-1, -1, -1, 1}));
  EXPECT_DOUBLE_EQ(
      PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);

  x_dist_attr.set_dims_mapping({-1, 0, -1, 1});
  out_grad_dist_attr.set_dims_mapping({0, 1});
  x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
  out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(out_grad_shape),
                                              out_grad_dist_attr);

  spmdinfo = SqueezeGradInferSpmd(x, out_grad);

  EXPECT_EQ(spmdinfo.first.size(), 2UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({-1, 0, -1, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]), std::vector<int64_t>({0, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({-1, 0, -1, 1}));
  EXPECT_DOUBLE_EQ(
      PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
}

TEST(UnsqueezeGradInferSpmd, Ctor) {
  std::vector<int64_t> x_shape = {32, 48};
  std::vector<int64_t> out_grad_shape = {1, 32, 1, 48};

  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({1, -1}));
  x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));

  TensorDistAttr out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);
  out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 1, -1, -1}));
  out_grad_dist_attr.set_dynamic_dims(
      std::vector<bool>({false, false, false, false}));

  phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor out_grad(phi::make_ddim(out_grad_shape),
                                            out_grad_dist_attr);

  auto spmdinfo = UnsqueezeGradInferSpmd(x, out_grad);

  EXPECT_EQ(spmdinfo.first.size(), 2UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]), std::vector<int64_t>({1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
            std::vector<int64_t>({-1, 1, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({1, -1}));
  EXPECT_DOUBLE_EQ(
      PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);

  x_dist_attr.set_dims_mapping({0, 1});
  out_grad_dist_attr.set_dims_mapping({-1, 0, -1, 1});
  x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
  out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(out_grad_shape),
                                              out_grad_dist_attr);

  spmdinfo = UnsqueezeGradInferSpmd(x, out_grad);

  EXPECT_EQ(spmdinfo.first.size(), 2UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]), std::vector<int64_t>({0, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
            std::vector<int64_t>({-1, 0, -1, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]), std::vector<int64_t>({0, 1}));
  EXPECT_DOUBLE_EQ(
      PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
}

TEST(ScatterGradInferSpmd, Ctor) {
  std::vector<int64_t> index_shape = {16};
  std::vector<int64_t> updates_shape = {32, 32, 48};
  std::vector<int64_t> out_grad_shape = {64, 32, 48};

  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr index_dist_attr = TensorDistAttr();
  index_dist_attr.set_process_mesh(process_mesh);
  TensorDistAttr updates_dist_attr = TensorDistAttr();
  updates_dist_attr.set_process_mesh(process_mesh);
  TensorDistAttr out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);

  // [0], [-1, -1, 1], [0, -1, 1] -->
  // inputs: [-1], [-1, -1, 1], [-1, -1, 1]
  // x_grad: [-1, -1, 1], updates_grad: [-1, -1, 1]
  index_dist_attr.set_dims_mapping(std::vector<int64_t>{0});
  updates_dist_attr.set_dims_mapping({-1, -1, 1});
  out_grad_dist_attr.set_dims_mapping({0, -1, 1});
  phi::distributed::DistMetaTensor index(phi::make_ddim(index_shape),
                                         index_dist_attr);
  phi::distributed::DistMetaTensor updates(phi::make_ddim(updates_shape),
                                           updates_dist_attr);
  phi::distributed::DistMetaTensor out_grad(phi::make_ddim(out_grad_shape),
                                            out_grad_dist_attr);
  auto spmdinfo = ScatterGradInferSpmd(index, updates, out_grad, false);
  EXPECT_EQ(spmdinfo.first.size(), 3UL);
  EXPECT_EQ(spmdinfo.second.size(), 2UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]), std::vector<int64_t>({-1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
            std::vector<int64_t>({-1, -1, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[2]),
            std::vector<int64_t>({-1, -1, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({-1, -1, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[1]),
            std::vector<int64_t>({-1, -1, 1}));

  // [0], [0, -1, 1], [-1, 0, 1] -->
  // inputs: [-1], [-1, -1, 1], [-1, 0, 1]
  // x_grad: [-1, 0, 1], updates_grad: [-1, 0, 1]
  index_dist_attr.set_dims_mapping(std::vector<int64_t>{0});
  updates_dist_attr.set_dims_mapping({0, -1, 1});
  out_grad_dist_attr.set_dims_mapping({-1, 0, 1});
  index = phi::distributed::DistMetaTensor(phi::make_ddim(index_shape),
                                           index_dist_attr);
  updates = phi::distributed::DistMetaTensor(phi::make_ddim(updates_shape),
                                             updates_dist_attr);
  out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(out_grad_shape),
                                              out_grad_dist_attr);
  spmdinfo = ScatterGradInferSpmd(index, updates, out_grad, false);
  EXPECT_EQ(spmdinfo.first.size(), 3UL);
  EXPECT_EQ(spmdinfo.second.size(), 2UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]), std::vector<int64_t>({-1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
            std::vector<int64_t>({-1, -1, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[2]),
            std::vector<int64_t>({-1, 0, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({-1, 0, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[1]),
            std::vector<int64_t>({-1, 0, 1}));
}

TEST(GatherGradInferSpmd, Ctor) {
  std::vector<int64_t> x_shape = {64, 32, 48};
  std::vector<int64_t> index_shape = {16};
  std::vector<int64_t> out_grad_shape = {16, 32, 48};
  phi::Scalar axis(0);

  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};

  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  TensorDistAttr index_dist_attr = TensorDistAttr();
  index_dist_attr.set_process_mesh(process_mesh);
  TensorDistAttr out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);

  // axis = 0
  // [0, -1, 1], [0], [0, -1, 1] -->
  // inputs: [0, -1, 1], [-1], [-1, -1, 1]
  // x_grad: [-1, -1, 1]
  axis = 0;
  x_dist_attr.set_dims_mapping({0, -1, 1});
  index_dist_attr.set_dims_mapping(std::vector<int64_t>{0});
  out_grad_dist_attr.set_dims_mapping({0, -1, 1});
  phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor index(phi::make_ddim(index_shape),
                                         index_dist_attr);
  phi::distributed::DistMetaTensor out_grad(phi::make_ddim(out_grad_shape),
                                            out_grad_dist_attr);
  auto spmdinfo = GatherGradInferSpmd(x, index, out_grad, axis);
  EXPECT_EQ(spmdinfo.first.size(), 3UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({0, -1, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]), std::vector<int64_t>({-1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[2]),
            std::vector<int64_t>({-1, -1, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({-1, -1, 1}));

  // 0-d tensor
  // axis = 1
  // [0, -1, 1], [-1], [0, 1] -->
  // inputs: [0, -1, 1], [-1], [0, 1]
  // x_grad: [0, -1, 1]
  axis = 1;
  index_shape = {};
  out_grad_shape = {64, 48};
  x_dist_attr.set_dims_mapping({0, -1, 1});
  index_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
  out_grad_dist_attr.set_dims_mapping({0, 1});
  x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
  index = phi::distributed::DistMetaTensor(phi::make_ddim(index_shape),
                                           index_dist_attr);
  out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(out_grad_shape),
                                              out_grad_dist_attr);
  spmdinfo = GatherGradInferSpmd(x, index, out_grad, axis);
  EXPECT_EQ(spmdinfo.first.size(), 3UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({0, -1, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]), std::vector<int64_t>({-1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[2]), std::vector<int64_t>({0, 1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({0, -1, 1}));
}

TEST(GatherNdGradInferSpmd, Ctor) {
  std::vector<int64_t> x_shape = {32};
  std::vector<int64_t> index_shape = {16};
  std::vector<int64_t> out_grad_shape = {16};

  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};

  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  TensorDistAttr index_dist_attr = TensorDistAttr();
  index_dist_attr.set_process_mesh(process_mesh);
  TensorDistAttr out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);

  // inputs: [-1], [0] --> [0]
  // x_grad: [-1]
  x_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
  index_dist_attr.set_dims_mapping(std::vector<int64_t>{0});
  out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>{0});
  phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor index(phi::make_ddim(index_shape),
                                         index_dist_attr);
  phi::distributed::DistMetaTensor out_grad(phi::make_ddim(out_grad_shape),
                                            out_grad_dist_attr);
  auto spmdinfo = GatherNdGradInferSpmd(x, index, out_grad);
  EXPECT_EQ(spmdinfo.first.size(), 3UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]), std::vector<int64_t>({-1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]), std::vector<int64_t>({0}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[2]), std::vector<int64_t>({0}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]), std::vector<int64_t>({-1}));

  // inputs: [-1, -1], [0, -1, -1] --> [0, -1, -1]
  // x_grad: [-1, -1]
  x_shape = {64, 32};
  index_shape = {16, 16, 1};
  out_grad_shape = {16, 16, 32};
  x_dist_attr.set_dims_mapping({-1, -1});
  index_dist_attr.set_dims_mapping({0, -1, -1});
  out_grad_dist_attr.set_dims_mapping({0, -1, -1});
  x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
  index = phi::distributed::DistMetaTensor(phi::make_ddim(index_shape),
                                           index_dist_attr);
  out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(out_grad_shape),
                                              out_grad_dist_attr);
  spmdinfo = GatherNdGradInferSpmd(x, index, out_grad);
  EXPECT_EQ(spmdinfo.first.size(), 3UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({-1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
            std::vector<int64_t>({0, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[2]),
            std::vector<int64_t>({0, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({-1, -1}));
}

TEST(CumSumGradInferSpmd, Ctor) {
  std::vector<int64_t> x_shape = {64, 32, 48};
  std::vector<int64_t> out_grad_shape = {64, 32, 48};

  std::vector<int64_t> mesh_shape = {2, 4};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};

  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  TensorDistAttr out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);

  // axis = 1
  // [0, 1, -1], [0, 1, -1] -->
  // inputs: [0, 1, -1], [0, -1, -1]
  // x_grad: [0, -1, -1]
  x_dist_attr.set_dims_mapping({0, 1, -1});
  out_grad_dist_attr.set_dims_mapping({0, 1, -1});
  phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor out_grad(phi::make_ddim(out_grad_shape),
                                            out_grad_dist_attr);
  auto spmdinfo = CumSumGradInferSpmd(x, out_grad, 1, false, false, false);
  EXPECT_EQ(spmdinfo.first.size(), 2UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({0, 1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
            std::vector<int64_t>({0, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({0, -1, -1}));

  // axis = -1
  // flatten = true
  // [0, 1, -1], [-1] -->
  // inputs: [0, 1, -1], [-1]
  // x_grad: [-1, -1, -1]
  x_dist_attr.set_dims_mapping({0, 1, -1});
  out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
  x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
  out_grad = phi::distributed::DistMetaTensor(phi::make_ddim({64 * 32 * 48}),
                                              out_grad_dist_attr);
  spmdinfo = CumSumGradInferSpmd(x, out_grad, -1, true, false, false);
  EXPECT_EQ(spmdinfo.first.size(), 2UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({0, 1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]), std::vector<int64_t>({-1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({-1, -1, -1}));
}

TEST(Flatten, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  auto build_input = [&](const std::vector<int64_t>& shape,
                         const std::vector<int64_t>& dim_mapping) {
    auto t_dist_attr = TensorDistAttr();
    t_dist_attr.set_process_mesh(process_mesh);
    t_dist_attr.set_dims_mapping(dim_mapping);
    t_dist_attr.set_dynamic_dims(std::vector<bool>(shape.size(), false));
    auto input =
        phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
    return input;
  };

  // [b, h/ph, w/pw, c, ph, pw]; dp
  auto input1 = build_input({4, 16, 16, 4, 2, 2}, {0, -1, -1, -1, -1, -1});
  // [b, h/ph, w/pw, c, ph, pw] => [b, h/ph, w/pw, hidden_size]
  auto spmd1 = FlattenInferSpmd(input1, -3, -1);
  EXPECT_EQ(spmd1.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(spmd1.second.size(), static_cast<size_t>(1));
  check_dim_mapping(spmd1.first[0], {0, -1, -1, -1, -1, -1});
  check_dim_mapping(spmd1.second[0], {0, -1, -1, -1});

  // [b, h/ph, w/pw, c, ph, pw]; dp, mp
  auto input2 = build_input({4, 16, 16, 4, 2, 2}, {-1, 0, -1, 1, -1, -1});
  auto spmd2 = FlattenInferSpmd(input2, 1, 4);
  EXPECT_EQ(spmd2.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(spmd2.second.size(), static_cast<size_t>(1));
  check_dim_mapping(spmd2.first[0], {-1, 0, -1, -1, -1, -1});
  check_dim_mapping(spmd2.second[0], {-1, 0, -1});

  // [b, s, nh, h/nh]; dp , mp
  auto input3 = build_input({2, 1024, 32, 32}, {0, -1, 1, -1});
  // [b, s, nh, h/nh] => [b, s, h]
  auto spmd3 = FlattenInferSpmd(input3, 2, 3);
  EXPECT_EQ(spmd3.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(spmd3.second.size(), static_cast<size_t>(1));
  check_dim_mapping(spmd3.first[0], {0, -1, 1, -1});
  check_dim_mapping(spmd3.second[0], {0, -1, 1});

  // [b, c, d, h, w]; dp, mp
  auto input4 = build_input({4, 16, 16, 4, 16}, {-1, -1, 0, 1, -1});
  auto spmd4 = FlattenInferSpmd(input4, 1, 4);
  EXPECT_EQ(spmd4.first.size(), static_cast<size_t>(1));
  EXPECT_EQ(spmd4.second.size(), static_cast<size_t>(1));
  check_dim_mapping(spmd4.first[0], {-1, -1, -1, -1, -1});
  check_dim_mapping(spmd4.second[0], {-1, -1});

  auto out_grad = build_input({2, 1024, 1024}, {0, -1, 1});
  auto x = build_input({2, 1024, 4, 1024 / 4}, {0, -1, 1, -1});
  auto spmd_grad = FlattenGradInferSpmd(x, out_grad);
  EXPECT_EQ(spmd_grad.first.size(), static_cast<size_t>(2));
  EXPECT_EQ(spmd_grad.second.size(), static_cast<size_t>(1));
  check_dim_mapping(spmd_grad.first[0], {0, -1, 1, -1});
  check_dim_mapping(spmd_grad.first[1], {0, -1, 1});
  check_dim_mapping(spmd_grad.second[0], {0, -1, 1, -1});
}

TEST(Conv2dSPMDRuleInferForward, Ctor) {
  // build input data class
  std::vector<int64_t> input_shape = {2, 4, 8, 8};
  std::vector<int64_t> filter_shape = {10, 4, 3, 3};

  std::vector<int64_t> mesh_shape = {2, 2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<std::string> dim_names = {"dp", "mp", "pp"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr input_dist_attr = TensorDistAttr();
  input_dist_attr.set_process_mesh(process_mesh);
  input_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, -1, -1}));
  input_dist_attr.set_dynamic_dims(
      std::vector<bool>({false, false, false, false}));

  TensorDistAttr filter_dist_attr = TensorDistAttr();
  filter_dist_attr.set_process_mesh(process_mesh);
  filter_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, -1, -1}));
  filter_dist_attr.set_dynamic_dims(
      std::vector<bool>({false, false, false, false}));

  phi::distributed::DistMetaTensor input(common::make_ddim(input_shape),
                                         input_dist_attr);
  phi::distributed::DistMetaTensor filter(common::make_ddim(filter_shape),
                                          filter_dist_attr);

  // test 1
  input_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1, -1}));
  filter_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, -1, -1}));

  input = phi::distributed::DistMetaTensor(common::make_ddim(input_shape),
                                           input_dist_attr);
  filter = phi::distributed::DistMetaTensor(common::make_ddim(filter_shape),
                                            filter_dist_attr);
  auto inferred_dist_attrs = phi::distributed::Conv2dInferSpmd(input, filter);

  check_dim_mapping(inferred_dist_attrs.first[0], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {0, -1, -1, -1});

  // test 2
  input_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, -1, -1}));
  filter_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1, -1}));

  input = phi::distributed::DistMetaTensor(common::make_ddim(input_shape),
                                           input_dist_attr);
  filter = phi::distributed::DistMetaTensor(common::make_ddim(filter_shape),
                                            filter_dist_attr);
  inferred_dist_attrs = phi::distributed::Conv2dInferSpmd(input, filter);

  check_dim_mapping(inferred_dist_attrs.first[0], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, 0, -1, -1});

  // test 3
  input_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1, -1}));
  filter_dist_attr.set_dims_mapping(std::vector<int64_t>({1, -1, -1, -1}));

  input = phi::distributed::DistMetaTensor(common::make_ddim(input_shape),
                                           input_dist_attr);
  filter = phi::distributed::DistMetaTensor(common::make_ddim(filter_shape),
                                            filter_dist_attr);
  inferred_dist_attrs = phi::distributed::Conv2dInferSpmd(input, filter);

  check_dim_mapping(inferred_dist_attrs.first[0], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {0, 1, -1, -1});

  // test 4
  input_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 0, -1, -1}));
  filter_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 0, -1, -1}));

  input = phi::distributed::DistMetaTensor(common::make_ddim(input_shape),
                                           input_dist_attr);
  filter = phi::distributed::DistMetaTensor(common::make_ddim(filter_shape),
                                            filter_dist_attr);

  inferred_dist_attrs = phi::distributed::Conv2dInferSpmd(input, filter);

  check_dim_mapping(inferred_dist_attrs.first[0], {-1, 0, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {-1, 0, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, -1, -1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), true);

  // test 5
  input_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 2, -1, -1}));
  filter_dist_attr.set_dims_mapping(std::vector<int64_t>({1, 2, -1, -1}));

  input = phi::distributed::DistMetaTensor(common::make_ddim(input_shape),
                                           input_dist_attr);
  filter = phi::distributed::DistMetaTensor(common::make_ddim(filter_shape),
                                            filter_dist_attr);

  inferred_dist_attrs = phi::distributed::Conv2dInferSpmd(input, filter);

  check_dim_mapping(inferred_dist_attrs.first[0], {0, 2, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {1, 2, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {0, 1, -1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), true);
}

TEST(Conv2dSPMDRuleInferBackward, Ctor) {
  // build input data class
  std::vector<int64_t> input_shape = {2, 4, 8, 8};
  std::vector<int64_t> filter_shape = {10, 4, 3, 3};
  std::vector<int64_t> output_shape = {2, 10, 6, 6};

  std::vector<int64_t> mesh_shape = {2, 2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<std::string> dim_names = {"dp", "mp", "pp"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr input_dist_attr = TensorDistAttr();
  input_dist_attr.set_process_mesh(process_mesh);
  input_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, -1, -1}));
  input_dist_attr.set_dynamic_dims(
      std::vector<bool>({false, false, false, false}));

  TensorDistAttr filter_dist_attr = TensorDistAttr();
  filter_dist_attr.set_process_mesh(process_mesh);
  filter_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, -1, -1}));
  filter_dist_attr.set_dynamic_dims(
      std::vector<bool>({false, false, false, false}));

  TensorDistAttr output_dist_attr = TensorDistAttr();
  output_dist_attr.set_process_mesh(process_mesh);
  output_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1, -1}));
  output_dist_attr.set_dynamic_dims(
      std::vector<bool>({false, false, false, false}));

  phi::distributed::DistMetaTensor input(common::make_ddim(input_shape),
                                         input_dist_attr);
  phi::distributed::DistMetaTensor filter(common::make_ddim(filter_shape),
                                          filter_dist_attr);
  phi::distributed::DistMetaTensor output(common::make_ddim(output_shape),
                                          output_dist_attr);

  auto inferred_dist_attrs =
      phi::distributed::Conv2dInferSpmdReverse(input, filter, output);

  check_dim_mapping(inferred_dist_attrs.first[0], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {0, 1, -1, -1});
}

TEST(Conv2dGradSPMDRule, Ctor) {
  // build input data class
  std::vector<int64_t> input_shape = {2, 4, 8, 8};
  std::vector<int64_t> filter_shape = {10, 4, 3, 3};
  std::vector<int64_t> output_grad_shape = {2, 10, 6, 6};

  std::vector<int64_t> mesh_shape = {2, 2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<std::string> dim_names = {"dp", "mp", "pp"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr input_dist_attr = TensorDistAttr();
  input_dist_attr.set_process_mesh(process_mesh);
  input_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, -1, -1}));
  input_dist_attr.set_dynamic_dims(
      std::vector<bool>({false, false, false, false}));

  TensorDistAttr filter_dist_attr = TensorDistAttr();
  filter_dist_attr.set_process_mesh(process_mesh);
  filter_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, -1, -1}));
  filter_dist_attr.set_dynamic_dims(
      std::vector<bool>({false, false, false, false}));

  TensorDistAttr output_grad_dist_attr = TensorDistAttr();
  output_grad_dist_attr.set_process_mesh(process_mesh);
  output_grad_dist_attr.set_dims_mapping(
      std::vector<int64_t>({-1, -1, -1, -1}));
  output_grad_dist_attr.set_dynamic_dims(
      std::vector<bool>({false, false, false, false}));

  phi::distributed::DistMetaTensor input(common::make_ddim(input_shape),
                                         input_dist_attr);
  phi::distributed::DistMetaTensor filter(common::make_ddim(filter_shape),
                                          filter_dist_attr);
  phi::distributed::DistMetaTensor output_grad(
      common::make_ddim(output_grad_shape), output_grad_dist_attr);

  // test 1
  auto inferred_dist_attrs =
      phi::distributed::Conv2dGradInferSpmd(input, filter, output_grad);
  EXPECT_EQ(inferred_dist_attrs.first.size(), (size_t)3);
  EXPECT_EQ(inferred_dist_attrs.second.size(), (size_t)2);
  check_dim_mapping(inferred_dist_attrs.first[0], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[2], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[1], {-1, -1, -1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.first[2]), false);
  VLOG(4) << "test 1 done";

  // test 2
  input_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1, -1}));
  filter_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, -1, -1}));
  output_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1, -1}));

  input = phi::distributed::DistMetaTensor(common::make_ddim(input_shape),
                                           input_dist_attr);
  filter = phi::distributed::DistMetaTensor(common::make_ddim(filter_shape),
                                            filter_dist_attr);
  output_grad = phi::distributed::DistMetaTensor(
      common::make_ddim(output_grad_shape), output_grad_dist_attr);
  inferred_dist_attrs =
      phi::distributed::Conv2dGradInferSpmd(input, filter, output_grad);

  check_dim_mapping(inferred_dist_attrs.first[0], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[2], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[1], {-1, -1, -1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.first[2]), false);
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[1]), true);

  // test 3
  input_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, -1, -1}));
  filter_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1, -1}));
  output_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 0, -1, -1}));

  input = phi::distributed::DistMetaTensor(common::make_ddim(input_shape),
                                           input_dist_attr);
  filter = phi::distributed::DistMetaTensor(common::make_ddim(filter_shape),
                                            filter_dist_attr);
  output_grad = phi::distributed::DistMetaTensor(
      common::make_ddim(output_grad_shape), output_grad_dist_attr);
  inferred_dist_attrs =
      phi::distributed::Conv2dGradInferSpmd(input, filter, output_grad);

  check_dim_mapping(inferred_dist_attrs.first[0], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[2], {-1, 0, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {-1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[1], {0, -1, -1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.first[2]), false);
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), true);

  // test 4
  input_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1, -1}));
  filter_dist_attr.set_dims_mapping(std::vector<int64_t>({1, -1, -1, -1}));
  output_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1, -1}));

  input = phi::distributed::DistMetaTensor(common::make_ddim(input_shape),
                                           input_dist_attr);
  filter = phi::distributed::DistMetaTensor(common::make_ddim(filter_shape),
                                            filter_dist_attr);
  output_grad = phi::distributed::DistMetaTensor(
      common::make_ddim(output_grad_shape), output_grad_dist_attr);
  inferred_dist_attrs =
      phi::distributed::Conv2dGradInferSpmd(input, filter, output_grad);

  check_dim_mapping(inferred_dist_attrs.first[0], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {1, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[2], {0, 1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {0, -1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[1], {1, -1, -1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.first[2]), false);
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), true);
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[1]), true);

  // test 5
  input_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 2, -1, -1}));
  filter_dist_attr.set_dims_mapping(std::vector<int64_t>({1, 2, -1, -1}));
  output_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1, -1}));
  output_grad_dist_attr.set_partial_status(std::vector<int64_t>({2}));

  input = phi::distributed::DistMetaTensor(common::make_ddim(input_shape),
                                           input_dist_attr);
  filter = phi::distributed::DistMetaTensor(common::make_ddim(filter_shape),
                                            filter_dist_attr);
  output_grad = phi::distributed::DistMetaTensor(
      common::make_ddim(output_grad_shape), output_grad_dist_attr);

  inferred_dist_attrs =
      phi::distributed::Conv2dGradInferSpmd(input, filter, output_grad);

  check_dim_mapping(inferred_dist_attrs.first[0], {0, 2, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[1], {1, 2, -1, -1});
  check_dim_mapping(inferred_dist_attrs.first[2], {0, 1, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[0], {0, 2, -1, -1});
  check_dim_mapping(inferred_dist_attrs.second[1], {1, 2, -1, -1});
  EXPECT_EQ(is_partial(inferred_dist_attrs.first[2]), true);
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[0]), true);
  EXPECT_EQ(is_partial(inferred_dist_attrs.second[1]), true);
}

TEST(Dropout, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"dp", "mp"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  // test forward
  auto t_dist_attr = TensorDistAttr();
  t_dist_attr.set_process_mesh(process_mesh);
  t_dist_attr.set_dims_mapping({0, -1, -1});
  t_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
      common::make_ddim({4, 6, 8}), t_dist_attr);
  phi::distributed::DistMetaTensor seed_tensor;
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::DropoutFwdInferSpmd(
          x, seed_tensor, 0.5, false, "", 0, true);
  check_dim_mapping(forward_info.first[0], {0, -1, -1});
  check_dim_mapping(forward_info.first[1], {});
  check_dim_mapping(forward_info.second[0], {0, -1, -1});
  check_dim_mapping(forward_info.second[1], {0, -1, -1});
  // test backward
  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
      common::make_ddim({4, 6, 8}), t_dist_attr);
  phi::distributed::DistMetaTensor mask = out_grad;
  phi::distributed::SpmdInfo backward_info =
      phi::distributed::DropoutBwdInferSpmd(mask, out_grad, 0.5, false, "");
  check_dim_mapping(backward_info.first[0], {0, -1, -1});
  check_dim_mapping(backward_info.first[1], {0, -1, -1});
  check_dim_mapping(backward_info.second[0], {0, -1, -1});
}

TEST(MeanAll, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  // test forward
  // [0, -1] --> [], partial_on_dim:[0]
  auto t_dist_attr = TensorDistAttr();
  t_dist_attr.set_process_mesh(process_mesh);
  t_dist_attr.set_dims_mapping({0, -1});
  t_dist_attr.set_dynamic_dims({false, false});
  phi::distributed::DistMetaTensor x =
      phi::distributed::DistMetaTensor(common::make_ddim({4, 8}), t_dist_attr);
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::MeanAllInferSpmd(x);

  EXPECT_EQ(forward_info.first.size(), 1UL);
  EXPECT_EQ(forward_info.second.size(), 1UL);

  check_dim_mapping(forward_info.first[0], {0, -1});
  check_dim_mapping(forward_info.second[0], {});
  check_partial_dims(forward_info.second[0], {0});

  // test backward
  // [] --> [-1, -1], []
  auto out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);
  out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>{});
  out_grad_dist_attr.set_dynamic_dims({});
  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
      common::make_ddim({}), out_grad_dist_attr);
  phi::distributed::SpmdInfo backward_info =
      phi::distributed::MeanAllGradInferSpmd(x, out_grad);

  EXPECT_EQ(backward_info.first.size(), 2UL);
  EXPECT_EQ(backward_info.second.size(), 1UL);

  check_dim_mapping(backward_info.first[0], {-1, -1});
  check_dim_mapping(backward_info.first[1], {});
  check_dim_mapping(backward_info.second[0], {-1, -1});
}
TEST(BatchNorm, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  // test forward
  // data_format = NCHW
  // [0, 1, -1, -1],[-1],[-1],[-1],[-1] ->[-1 , 1, -1, -1],[1],[1],[1],[1],[-1]
  auto x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping({0, 1, -1, -1});
  x_dist_attr.set_dynamic_dims({false, false, false, false});
  auto one_dim_dist_attr = TensorDistAttr();
  one_dim_dist_attr.set_process_mesh(process_mesh);
  one_dim_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
  one_dim_dist_attr.set_dynamic_dims({false});

  phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 16, 16, 16}), x_dist_attr);
  phi::distributed::DistMetaTensor mean = phi::distributed::DistMetaTensor(
      common::make_ddim({16}), one_dim_dist_attr);
  phi::distributed::DistMetaTensor variance = phi::distributed::DistMetaTensor(
      common::make_ddim({16}), one_dim_dist_attr);
  phi::distributed::DistMetaTensor scale = phi::distributed::DistMetaTensor(
      common::make_ddim({16}), one_dim_dist_attr);
  phi::distributed::DistMetaTensor bias = phi::distributed::DistMetaTensor(
      common::make_ddim({16}), one_dim_dist_attr);
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::BatchNormInferSpmdStatic(
          x, mean, variance, scale, bias);

  EXPECT_EQ(forward_info.first.size(), 5UL);
  EXPECT_EQ(forward_info.second.size(), 6UL);
  check_dim_mapping(forward_info.first[0], {-1, 1, -1, -1});
  check_dim_mapping(forward_info.first[1], {1});
  check_dim_mapping(forward_info.first[2], {1});
  check_dim_mapping(forward_info.first[3], {-1});
  check_dim_mapping(forward_info.first[4], {-1});
  check_dim_mapping(forward_info.second[0], {-1, 1, -1, -1});
  check_dim_mapping(forward_info.second[1], {1});
  check_dim_mapping(forward_info.second[2], {1});
  check_dim_mapping(forward_info.second[3], {1});
  check_dim_mapping(forward_info.second[4], {1});
  check_dim_mapping(forward_info.second[5], {-1});

  // test backward
  // data_format = NCHW
  // [0, 1, -1, -1],[-1],[-1],[-1],[-1],[-1],[-1],[-1],[0, 1, -1, -1]
  // ->[-1,1,-1,-1],[-1],[-1]
  // dst_input: [-1, 1, -1, -1],[-1],[-1],[1],[1],[1],[1],[-1],[-1, 1, -1, -1]

  x = phi::distributed::DistMetaTensor(common::make_ddim({16, 16, 16, 16}),
                                       x_dist_attr);
  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 16, 16, 16}), x_dist_attr);
  phi::distributed::DistMetaTensor mean_out = phi::distributed::DistMetaTensor(
      common::make_ddim({16}), one_dim_dist_attr);
  phi::distributed::DistMetaTensor variance_out =
      phi::distributed::DistMetaTensor(common::make_ddim({16}),
                                       one_dim_dist_attr);
  scale = phi::distributed::DistMetaTensor(common::make_ddim({16}),
                                           one_dim_dist_attr);
  bias = phi::distributed::DistMetaTensor(common::make_ddim({16}),
                                          one_dim_dist_attr);
  phi::distributed::DistMetaTensor saved_mean =
      phi::distributed::DistMetaTensor(common::make_ddim({16}),
                                       one_dim_dist_attr);
  phi::distributed::DistMetaTensor saved_variance =
      phi::distributed::DistMetaTensor(common::make_ddim({16}),
                                       one_dim_dist_attr);
  phi::distributed::DistMetaTensor reserve_space =
      phi::distributed::DistMetaTensor(common::make_ddim({16}),
                                       one_dim_dist_attr);
  phi::distributed::SpmdInfo backward_info =
      phi::distributed::BatchNormGradInferSpmd(x,
                                               scale,
                                               bias,
                                               mean_out,
                                               variance_out,
                                               saved_mean,
                                               saved_variance,
                                               reserve_space,
                                               out_grad,
                                               0.9,
                                               0.1,
                                               "NCHW",
                                               false,
                                               false,
                                               false);

  EXPECT_EQ(backward_info.first.size(), 9UL);
  EXPECT_EQ(backward_info.second.size(), 3UL);
  check_dim_mapping(backward_info.first[0], {-1, 1, -1, -1});
  check_dim_mapping(backward_info.first[1], {-1});
  check_dim_mapping(backward_info.first[2], {-1});
  check_dim_mapping(backward_info.first[3], {1});
  check_dim_mapping(backward_info.first[4], {1});
  check_dim_mapping(backward_info.first[5], {1});
  check_dim_mapping(backward_info.first[6], {1});
  check_dim_mapping(backward_info.first[7], {-1});
  check_dim_mapping(backward_info.first[8], {-1, 1, -1, -1});

  check_dim_mapping(backward_info.second[0], {-1, 1, -1, -1});
  check_dim_mapping(backward_info.second[1], {-1});
  check_dim_mapping(backward_info.second[2], {-1});
}
TEST(Topk, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
  phi::Scalar k(2);

  // test forward
  // axis = 1
  // [0, 1, -1] -> [0, -1, -1], [0, -1, -1]
  auto x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping({0, 1, -1});
  x_dist_attr.set_dynamic_dims({false, false, false});

  phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 16, 16}), x_dist_attr);
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::TopkInferSpmdDynamic(x, k, 1, true, true);

  EXPECT_EQ(forward_info.first.size(), 1UL);
  EXPECT_EQ(forward_info.second.size(), 2UL);
  check_dim_mapping(forward_info.first[0], {0, -1, -1});
  check_dim_mapping(forward_info.second[0], {0, -1, -1});
  check_dim_mapping(forward_info.second[1], {0, -1, -1});

  // test backward
  // axis = 1
  // [0, -1, 1], [0, -1, 1], [-1, 1, -1]-> [0, -1, 1], [0, -1, 1], [0, -1,
  // 1], [0, -1, 1]
  x_dist_attr.set_dims_mapping({0, -1, 1});
  x = phi::distributed::DistMetaTensor(common::make_ddim({16, 16, 16}),
                                       x_dist_attr);
  auto out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);
  out_grad_dist_attr.set_dims_mapping({-1, 1, -1});
  out_grad_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 2, 16}), out_grad_dist_attr);
  phi::distributed::SpmdInfo backward_info =
      phi::distributed::TopkGradInferSpmdDynamic(
          x, x, out_grad, k, 1, true, true);
  EXPECT_EQ(backward_info.first.size(), 3UL);
  EXPECT_EQ(backward_info.second.size(), 1UL);
  check_dim_mapping(backward_info.first[0], {0, -1, 1});
  check_dim_mapping(backward_info.first[1], {0, -1, 1});
  check_dim_mapping(backward_info.first[2], {0, -1, 1});
  check_dim_mapping(backward_info.second[0], {0, -1, 1});
}

TEST(Cummax, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
  int axis = 1;

  // test forward
  // axis = 1
  // [0, 1, -1] -> [0, -1, -1], [0, -1, -1]
  auto x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping({0, 1, -1});
  x_dist_attr.set_dynamic_dims({false, false, false});

  phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 16, 16}), x_dist_attr);
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::CummaxInferSpmd(x, axis, phi::DataType::INT64);

  EXPECT_EQ(forward_info.first.size(), 1UL);
  EXPECT_EQ(forward_info.second.size(), 2UL);
  check_dim_mapping(forward_info.first[0], {0, -1, -1});
  check_dim_mapping(forward_info.second[0], {0, -1, -1});
  check_dim_mapping(forward_info.second[1], {0, -1, -1});

  // test backward
  // axis = 1
  // [0, -1, 1], [0, -1, 1], [-1, 1, -1]-> [0, -1, 1], [0, -1, 1], [0, -1,
  // 1], [0, -1, 1]
  x_dist_attr.set_dims_mapping({0, -1, 1});
  x = phi::distributed::DistMetaTensor(common::make_ddim({16, 16, 16}),
                                       x_dist_attr);
  auto out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);
  out_grad_dist_attr.set_dims_mapping({-1, 1, -1});
  out_grad_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 16, 16}), out_grad_dist_attr);
  phi::distributed::SpmdInfo backward_info =
      phi::distributed::CummaxGradInferSpmd(
          x, x, out_grad, axis, phi::DataType::INT64);
  EXPECT_EQ(backward_info.first.size(), 3UL);
  EXPECT_EQ(backward_info.second.size(), 1UL);
  check_dim_mapping(backward_info.first[0], {0, -1, 1});
  check_dim_mapping(backward_info.first[1], {0, -1, 1});
  check_dim_mapping(backward_info.first[2], {0, -1, 1});
  check_dim_mapping(backward_info.second[0], {0, -1, 1});
}

TEST(Cummin, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
  int axis = 1;

  // test forward
  // axis = 1
  // [0, 1, -1] -> [0, -1, -1], [0, -1, -1]
  auto x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping({0, 1, -1});
  x_dist_attr.set_dynamic_dims({false, false, false});

  phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 16, 16}), x_dist_attr);
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::CumminInferSpmd(x, axis, phi::DataType::INT64);

  EXPECT_EQ(forward_info.first.size(), 1UL);
  EXPECT_EQ(forward_info.second.size(), 2UL);
  check_dim_mapping(forward_info.first[0], {0, -1, -1});
  check_dim_mapping(forward_info.second[0], {0, -1, -1});
  check_dim_mapping(forward_info.second[1], {0, -1, -1});

  // test backward
  // axis = 1
  // [0, -1, 1], [0, -1, 1], [-1, 1, -1]-> [0, -1, 1], [0, -1, 1], [0, -1,
  // 1], [0, -1, 1]
  x_dist_attr.set_dims_mapping({0, -1, 1});
  x = phi::distributed::DistMetaTensor(common::make_ddim({16, 16, 16}),
                                       x_dist_attr);
  auto out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);
  out_grad_dist_attr.set_dims_mapping({-1, 1, -1});
  out_grad_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 16, 16}), out_grad_dist_attr);
  phi::distributed::SpmdInfo backward_info =
      phi::distributed::CumminGradInferSpmd(
          x, x, out_grad, axis, phi::DataType::INT64);
  EXPECT_EQ(backward_info.first.size(), 3UL);
  EXPECT_EQ(backward_info.second.size(), 1UL);
  check_dim_mapping(backward_info.first[0], {0, -1, 1});
  check_dim_mapping(backward_info.first[1], {0, -1, 1});
  check_dim_mapping(backward_info.first[2], {0, -1, 1});
  check_dim_mapping(backward_info.second[0], {0, -1, 1});
}
TEST(ArgSortGradInferSpmd, Ctor) {
  // Sharding along axes besides argsort axis.
  std::vector<int64_t> x_shape = {16, 32, 48};
  std::vector<int64_t> indices_shape = {16, 32, 48};
  std::vector<int64_t> out_grad_shape = {16, 32, 48};

  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1}));
  x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));

  TensorDistAttr indices_dist_attr = TensorDistAttr();
  indices_dist_attr.set_process_mesh(process_mesh);
  indices_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1}));
  indices_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));

  TensorDistAttr out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);
  out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1}));
  out_grad_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));

  phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
  phi::distributed::DistMetaTensor indices(phi::make_ddim(x_shape),
                                           indices_dist_attr);
  phi::distributed::DistMetaTensor out_grad(phi::make_ddim(x_shape),
                                            out_grad_dist_attr);
  int axis = -1;

  auto spmdinfo =
      ArgSortGradInferSpmd(indices, x, out_grad, axis, false, false);

  EXPECT_EQ(spmdinfo.first.size(), 3UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({0, 1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
            std::vector<int64_t>({0, 1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[2]),
            std::vector<int64_t>({0, 1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({0, 1, -1}));
  VLOG(4) << "Test ArgSortGradInferSpmd sharding on other axes." << std::endl
          << std::endl
          << std::endl;

  // Sharding along argsort axis.
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, 1}));
  indices_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, 1}));
  out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, 1}));
  x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
  indices = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape),
                                             indices_dist_attr);
  out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape),
                                              out_grad_dist_attr);
  axis = 2;

  spmdinfo = ArgSortGradInferSpmd(indices, x, out_grad, axis, false, false);

  EXPECT_EQ(spmdinfo.first.size(), 3UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({0, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
            std::vector<int64_t>({0, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[2]),
            std::vector<int64_t>({0, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({0, -1, -1}));
  VLOG(4) << "Test ArgSortGradInferSpmd sharding on softmax axis." << std::endl
          << std::endl
          << std::endl;

  // Sharding on multi axes.
  x_shape = {10, 32, 48, 24};
  indices_shape = {10, 32, 48, 24};
  out_grad_shape = {10, 32, 48, 24};
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1, -1}));
  indices_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1, -1}));
  out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1, -1}));
  x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
  indices = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape),
                                             indices_dist_attr);
  out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape),
                                              out_grad_dist_attr);
  axis = 1;

  spmdinfo = ArgSortGradInferSpmd(indices, x, out_grad, axis, false, false);

  EXPECT_EQ(spmdinfo.first.size(), 3UL);
  EXPECT_EQ(spmdinfo.second.size(), 1UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({0, -1, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
            std::vector<int64_t>({0, -1, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.first[2]),
            std::vector<int64_t>({0, -1, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({0, -1, -1, -1}));
  VLOG(4) << "Test ArgSortGradInferSpmd sharding on multi axes." << std::endl
          << std::endl
          << std::endl;
}

TEST(ArgSortInferSpmd, Ctor) {
  // Sharding along axes besides argsort axis.
  std::vector<int64_t> x_shape = {16, 32, 48};

  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);

  TensorDistAttr x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1}));
  x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));

  phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
  int axis = -1;

  auto spmdinfo = ArgSortInferSpmd(x, axis, false, false);

  EXPECT_EQ(spmdinfo.first.size(), 1UL);
  EXPECT_EQ(spmdinfo.second.size(), 2UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({0, 1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({0, 1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[1]),
            std::vector<int64_t>({0, 1, -1}));
  VLOG(4) << "Test ArgSortGradInferSpmd sharding on other axes." << std::endl
          << std::endl
          << std::endl;

  // Sharding along argsort axis.
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, 1}));
  x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
  axis = 2;

  spmdinfo = ArgSortInferSpmd(x, axis, false, false);

  EXPECT_EQ(spmdinfo.first.size(), 1UL);
  EXPECT_EQ(spmdinfo.second.size(), 2UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({0, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({0, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[1]),
            std::vector<int64_t>({0, -1, -1}));
  VLOG(4) << "Test ArgSortGradInferSpmd sharding on softmax axis." << std::endl
          << std::endl
          << std::endl;

  // Sharding on multi axes.
  x_shape = {10, 32, 48, 24};
  x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1, -1}));
  x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);

  axis = 1;

  spmdinfo = ArgSortInferSpmd(x, axis, false, false);

  EXPECT_EQ(spmdinfo.first.size(), 1UL);
  EXPECT_EQ(spmdinfo.second.size(), 2UL);

  EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
            std::vector<int64_t>({0, -1, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
            std::vector<int64_t>({0, -1, -1, -1}));
  EXPECT_EQ(get_dims_mapping(spmdinfo.second[1]),
            std::vector<int64_t>({0, -1, -1, -1}));
  VLOG(4) << "Test ArgSortGradInferSpmd sharding on multi axes." << std::endl
          << std::endl
          << std::endl;
}

TEST(Roll, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
  phi::IntArray shifts = {1};
  std::vector<int64_t> axis = {};

  // test forward
  // axis = [], shifts = [1]
  // [0, 1, -1] --> [-1, -1, -1], [-1, -1, -1]
  auto x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping({0, 1, -1});
  x_dist_attr.set_dynamic_dims({false, false, false});

  phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 16, 16}), x_dist_attr);
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::RollInferSpmdDynamic(x, shifts, axis);

  EXPECT_EQ(forward_info.first.size(), 1UL);
  EXPECT_EQ(forward_info.second.size(), 1UL);
  check_dim_mapping(forward_info.first[0], {-1, -1, -1});
  check_dim_mapping(forward_info.second[0], {-1, -1, -1});
  // axis = [0, 2], shifts = [1, 2]
  // [0, 1, -1] --> [-1, 1, -1], [-1, 1, -1]
  shifts = {1, 2};
  axis = {0, 2};
  forward_info = phi::distributed::RollInferSpmdDynamic(x, shifts, axis);

  EXPECT_EQ(forward_info.first.size(), 1UL);
  EXPECT_EQ(forward_info.second.size(), 1UL);
  check_dim_mapping(forward_info.first[0], {-1, 1, -1});
  check_dim_mapping(forward_info.second[0], {-1, 1, -1});

  // test backward
  // axis = [], shifts = [1]
  // [0, 1, -1], [-1, -1, -1] --> [-1, -1, -1], [-1, -1, -1], [-1, -1, -1]
  shifts = {1};
  axis = {};
  x_dist_attr.set_dims_mapping({-1, -1, -1});
  x = phi::distributed::DistMetaTensor(common::make_ddim({16, 16, 16}),
                                       x_dist_attr);
  auto out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);
  out_grad_dist_attr.set_dims_mapping({0, 1, -1});
  out_grad_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 16, 16}), out_grad_dist_attr);
  phi::distributed::SpmdInfo backward_info =
      phi::distributed::RollGradInferSpmdDynamic(x, out_grad, shifts, axis);
  EXPECT_EQ(backward_info.first.size(), 2UL);
  EXPECT_EQ(backward_info.second.size(), 1UL);
  check_dim_mapping(backward_info.first[0], {-1, -1, -1});
  check_dim_mapping(backward_info.first[1], {-1, -1, -1});
  check_dim_mapping(backward_info.second[0], {-1, -1, -1});

  // axis = [0, 2], shifts = [1, 2]
  // [-1, -1, -1], [0, 1, -1] --> [-1, 1, -1], [-1, 1, -1], [-1, 1, -1]
  shifts = {1, 2};
  axis = {0, 2};
  backward_info =
      phi::distributed::RollGradInferSpmdDynamic(x, out_grad, shifts, axis);
  EXPECT_EQ(backward_info.first.size(), 2UL);
  EXPECT_EQ(backward_info.second.size(), 1UL);
  check_dim_mapping(backward_info.first[0], {-1, 1, -1});
  check_dim_mapping(backward_info.first[1], {-1, 1, -1});
  check_dim_mapping(backward_info.second[0], {-1, 1, -1});
}

TEST(IndexSelect, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 3};
  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
  int axis = 1;

  // test forward
  // axis = 1
  // [-1, -1, -1], [0] --> [-1, -1, -1], [0], [-1, 0, -1]
  auto x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping({-1, -1, -1});
  x_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 16, 16}), x_dist_attr);
  auto index_dist_attr = TensorDistAttr();
  index_dist_attr.set_process_mesh(process_mesh);
  index_dist_attr.set_dims_mapping(std::vector<int64_t>{0});
  index_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor index =
      phi::distributed::DistMetaTensor(common::make_ddim({3}), index_dist_attr);

  phi::distributed::SpmdInfo forward_info =
      phi::distributed::IndexSelectInferSpmd(x, index, axis);

  EXPECT_EQ(forward_info.first.size(), 2UL);
  EXPECT_EQ(forward_info.second.size(), 1UL);
  check_dim_mapping(forward_info.first[0], {-1, -1, -1});
  check_dim_mapping(forward_info.first[1], {0});
  check_dim_mapping(forward_info.second[0], {-1, 0, -1});

  // test backward
  // axis = 1
  // [-1, -1, -1], [0], [-1, 0, -1], axis=1 --> [-1, -1, -1], [0], [-1, 0, -1],
  // [-1, -1, -1](partial on axis=1 with 0)
  auto out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);
  out_grad_dist_attr.set_dims_mapping({-1, 0, -1});
  out_grad_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 3, 16}), out_grad_dist_attr);
  phi::distributed::SpmdInfo backward_info =
      phi::distributed::IndexSelectGradInferSpmd(x, index, out_grad, axis);
  EXPECT_EQ(backward_info.first.size(), 3UL);
  EXPECT_EQ(backward_info.second.size(), 1UL);
  check_dim_mapping(backward_info.first[0], {-1, -1, -1});
  check_dim_mapping(backward_info.first[1], {0});
  check_dim_mapping(backward_info.first[2], {-1, 0, -1});
  check_dim_mapping(backward_info.second[0], {-1, -1, -1});
  EXPECT_EQ(is_partial(backward_info.second[0]), true);
  check_partial_dims(backward_info.second[0], {0});
}

TEST(Unique, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
  bool return_index = true;
  bool return_inverse = true;
  bool return_counts = true;
  std::vector<int> axis = {};

  // test forward
  // return_index=True, return_inverse=True, return_counts=True, axis={}
  // [0, -1] --> [-1,-1], [-1], [-1], [-1], [-1]
  auto x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping({0, -1});
  x_dist_attr.set_dynamic_dims({false, false});
  phi::distributed::DistMetaTensor x =
      phi::distributed::DistMetaTensor(common::make_ddim({4, 8}), x_dist_attr);

  phi::distributed::SpmdInfo forward_info =
      phi::distributed::UniqueInferSpmd(x,
                                        return_index,
                                        return_inverse,
                                        return_counts,
                                        axis,
                                        phi::DataType::FLOAT32);

  EXPECT_EQ(forward_info.first.size(), 1UL);
  EXPECT_EQ(forward_info.second.size(), 4UL);

  check_dim_mapping(forward_info.first[0], {-1, -1});
  check_dim_mapping(forward_info.second[0], {-1});
  check_dim_mapping(forward_info.second[1], {-1});
  check_dim_mapping(forward_info.second[2], {-1});
  check_dim_mapping(forward_info.second[3], {-1});

  // return_index=True, return_inverse=True, return_counts=True, axis={0}
  // [0, -1] --> [-1, -1], [-1, -1], [-1], [-1], [-1]
  axis = {0};
  forward_info = phi::distributed::UniqueInferSpmd(x,
                                                   return_index,
                                                   return_inverse,
                                                   return_counts,
                                                   axis,
                                                   phi::DataType::FLOAT32);

  EXPECT_EQ(forward_info.first.size(), 1UL);
  EXPECT_EQ(forward_info.second.size(), 4UL);

  check_dim_mapping(forward_info.first[0], {-1, -1});
  check_dim_mapping(forward_info.second[0], {-1, -1});
  check_dim_mapping(forward_info.second[1], {-1});
  check_dim_mapping(forward_info.second[2], {-1});
  check_dim_mapping(forward_info.second[3], {-1});
}
TEST(LabelSmooth, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
  float epsilon = 0.1;

  // test forward
  // [0, 1, -1], [Fake] --> [0, 1, -1], [Fake], [0, 1, -1]
  auto label_dist_attr = TensorDistAttr();
  label_dist_attr.set_process_mesh(process_mesh);
  label_dist_attr.set_dims_mapping({0, 1, -1});
  label_dist_attr.set_dynamic_dims({false, false, false});

  auto prior_dist_dist_attr = TensorDistAttr();
  prior_dist_dist_attr.set_process_mesh(process_mesh);
  prior_dist_dist_attr.set_dims_mapping({-1, -1});
  prior_dist_dist_attr.set_dynamic_dims({false, false});

  phi::distributed::DistMetaTensor label = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 16, 16}), label_dist_attr);
  phi::distributed::DistMetaTensor prior_dist =
      phi::distributed::DistMetaTensor(common::make_ddim({1, 16}),
                                       prior_dist_dist_attr);
  phi::distributed::SpmdInfo forward_info =
      phi::distributed::LabelSmoothInferSpmd(
          label, phi::distributed::DistMetaTensor(), epsilon);

  EXPECT_EQ(forward_info.first.size(), 2UL);
  EXPECT_EQ(forward_info.second.size(), 1UL);
  check_dim_mapping(forward_info.first[0], {0, 1, -1});
  check_empty_dist_attr(forward_info.first[1]);
  check_dim_mapping(forward_info.second[0], {0, 1, -1});

  // shape: [16, 16, 16], [1, 16]. [0, 1, -1], [-1, -1] --> [0, 1, -1], [-1,
  // -1], [0, 1, -1]
  forward_info =
      phi::distributed::LabelSmoothInferSpmd(label, prior_dist, epsilon);
  EXPECT_EQ(forward_info.first.size(), 2UL);
  EXPECT_EQ(forward_info.second.size(), 1UL);
  check_dim_mapping(forward_info.first[0], {0, 1, -1});
  check_dim_mapping(forward_info.first[1], {-1, -1});
  check_dim_mapping(forward_info.second[0], {0, 1, -1});

  // test backward
  // [0, -1, 1] --> [0, -1, 1], [0, -1, 1]
  auto out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);
  out_grad_dist_attr.set_dims_mapping({0, -1, 1});
  out_grad_dist_attr.set_dynamic_dims({false, false, false});
  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
      common::make_ddim({16, 16, 16}), out_grad_dist_attr);
  phi::distributed::SpmdInfo backward_info =
      phi::distributed::LabelSmoothGradInferSpmd(out_grad, epsilon);
  EXPECT_EQ(backward_info.first.size(), 1UL);
  EXPECT_EQ(backward_info.second.size(), 1UL);
  check_dim_mapping(backward_info.first[0], {0, -1, 1});
  check_dim_mapping(backward_info.second[0], {0, -1, 1});
}
TEST(RolAlign, Ctor) {
  std::vector<int64_t> mesh_shape = {2, 2};
  std::vector<int64_t> process_ids = {0, 1, 2, 3};
  std::vector<std::string> dim_names = {"x", "y"};
  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
  int pooled_height = 3;
  int pooled_width = 3;
  float spatial_scale = 0.5;
  int sampling_ratio = -1;
  bool aligned = true;

  // test forward
  // [0, 1, -1, -1], [-1, 1], [0] --> [-1, 1, -1, -1],[-1, -1],[-1],[-1,1,-1,-1]
  auto x_dist_attr = TensorDistAttr();
  x_dist_attr.set_process_mesh(process_mesh);
  x_dist_attr.set_dims_mapping({0, 1, -1, -1});
  x_dist_attr.set_dynamic_dims({false, false, false, false});
  auto boxes_dist_attr = TensorDistAttr();
  boxes_dist_attr.set_process_mesh(process_mesh);
  boxes_dist_attr.set_dims_mapping({-1, 1});
  boxes_dist_attr.set_dynamic_dims({false, false});
  auto boxes_num_dist_attr = TensorDistAttr();
  boxes_num_dist_attr.set_process_mesh(process_mesh);
  boxes_num_dist_attr.set_dims_mapping(std::vector<int64_t>{0});
  boxes_num_dist_attr.set_dynamic_dims({false});

  phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
      common::make_ddim({2, 4, 16, 16}), x_dist_attr);
  phi::distributed::DistMetaTensor boxes = phi::distributed::DistMetaTensor(
      common::make_ddim({6, 6}), boxes_dist_attr);
  phi::distributed::DistMetaTensor boxes_num = phi::distributed::DistMetaTensor(
      common::make_ddim({2}), boxes_num_dist_attr);

  phi::distributed::SpmdInfo forward_info =
      phi::distributed::RoiAlignInferSpmd(x,
                                          boxes,
                                          boxes_num,
                                          pooled_height,
                                          pooled_width,
                                          spatial_scale,
                                          sampling_ratio,
                                          aligned);

  EXPECT_EQ(forward_info.first.size(), 3UL);
  EXPECT_EQ(forward_info.second.size(), 1UL);
  check_dim_mapping(forward_info.first[0], {-1, 1, -1, -1});
  check_dim_mapping(forward_info.first[1], {-1, -1});
  check_dim_mapping(forward_info.first[2], {-1});
  check_dim_mapping(forward_info.second[0], {-1, 1, -1, -1});
  forward_info =
      phi::distributed::RoiAlignInferSpmd(x,
                                          boxes,
                                          phi::distributed::DistMetaTensor(),
                                          pooled_height,
                                          pooled_width,
                                          spatial_scale,
                                          sampling_ratio,
                                          aligned);
  EXPECT_EQ(forward_info.first.size(), 3UL);
  EXPECT_EQ(forward_info.second.size(), 1UL);
  check_dim_mapping(forward_info.first[0], {-1, 1, -1, -1});
  check_dim_mapping(forward_info.first[1], {-1, -1});
  check_empty_dist_attr(forward_info.first[2]);
  check_dim_mapping(forward_info.second[0], {-1, 1, -1, -1});

  // test backward
  // [0, 1, -1, -1], [-1, 1], [0], [0, -1, -1, -1] --> [-1, 1, -1, -1],[-1,
  // -1],[-1],[-1,1,-1,-1],[-1, 1, -1, -1]
  auto out_grad_dist_attr = TensorDistAttr();
  out_grad_dist_attr.set_process_mesh(process_mesh);
  out_grad_dist_attr.set_dims_mapping({0, -1, -1, -1});
  out_grad_dist_attr.set_dynamic_dims({false, false, false, false});
  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
      common::make_ddim({6, 2, 3, 3}), out_grad_dist_attr);
  phi::distributed::SpmdInfo backward_info =
      phi::distributed::RoiAlignGradInferSpmd(x,
                                              boxes,
                                              boxes_num,
                                              out_grad,
                                              pooled_height,
                                              pooled_width,
                                              spatial_scale,
                                              sampling_ratio,
                                              aligned);
  EXPECT_EQ(backward_info.first.size(), 4UL);
  EXPECT_EQ(backward_info.second.size(), 1UL);
  check_dim_mapping(backward_info.first[0], {-1, 1, -1, -1});
  check_dim_mapping(backward_info.first[1], {-1, -1});
  check_dim_mapping(backward_info.first[2], {-1});
  check_dim_mapping(backward_info.first[3], {-1, 1, -1, -1});
  check_dim_mapping(backward_info.second[0], {-1, 1, -1, -1});

  backward_info = phi::distributed::RoiAlignGradInferSpmd(
      x,
      boxes,
      phi::distributed::DistMetaTensor(),
      out_grad,
      pooled_height,
      pooled_width,
      spatial_scale,
      sampling_ratio,
      aligned);
  EXPECT_EQ(backward_info.first.size(), 4UL);
  EXPECT_EQ(backward_info.second.size(), 1UL);
  check_dim_mapping(backward_info.first[0], {-1, 1, -1, -1});
  check_dim_mapping(backward_info.first[1], {-1, -1});
  check_empty_dist_attr(backward_info.first[2]);
  check_dim_mapping(backward_info.first[3], {-1, 1, -1, -1});
  check_dim_mapping(backward_info.first[0], {-1, 1, -1, -1});
}
}  // namespace auto_parallel
}  // namespace distributed
}  // namespace paddle