diff --git a/paddle/phi/infermeta/spmd_rules/label_smooth.cc b/paddle/phi/infermeta/spmd_rules/label_smooth.cc new file mode 100644 index 00000000000000..3146e27b98012b --- /dev/null +++ b/paddle/phi/infermeta/spmd_rules/label_smooth.cc @@ -0,0 +1,85 @@ +/* Copyright (c) 2025 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/infermeta/spmd_rules/label_smooth.h" +#include "glog/logging.h" +#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h" +#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h" +#include "paddle/phi/core/distributed/type_defs.h" +#include "paddle/phi/infermeta/spmd_rules/elementwise.h" +#include "paddle/phi/infermeta/spmd_rules/spmd_rule_macro_define.h" +#include "paddle/phi/infermeta/spmd_rules/utils.h" + +namespace phi { +namespace distributed { + +std::vector ReplicateWithExcludeAxes( + const std::vector& x, const std::vector& exclude_axes) { + std::vector res(x); + for (size_t i = 0; i < x.size(); ++i) { + if (res[i] == -1) { + continue; + } + + if (std::find(exclude_axes.begin(), exclude_axes.end(), res[i]) != + exclude_axes.end()) { + res[i] = -1; + } + } + return res; +} + +SpmdInfo LabelSmoothInferSpmd(const DistMetaTensor& label, + const DistMetaTensor& prior_dist, + float epsilon) { + if (prior_dist.initialized()) { + const auto& label_dims_attr = label.dist_attr(); + const auto& prior_dims_attr = prior_dist.dist_attr(); + const auto& label_dims_mapping = label_dims_attr.dims_mapping(); + const auto& prior_dims_mapping = prior_dims_attr.dims_mapping(); + size_t label_rank = label.dims().size(); + size_t prior_rank = prior_dist.dims().size(); + + if (label_rank > prior_rank) { + size_t extra_dims_count = label_rank - prior_rank; + std::vector extra_dims_mapping_values; + extra_dims_mapping_values.reserve(extra_dims_count); + for (size_t i = 0; i < extra_dims_count; ++i) { + extra_dims_mapping_values.push_back(label_dims_mapping[i]); + } + const auto& new_prior_dims_mapping = ReplicateWithExcludeAxes( + prior_dims_mapping, extra_dims_mapping_values); + TensorDistAttr new_prior_dist_attr = prior_dims_attr; + new_prior_dist_attr.set_dims_mapping(new_prior_dims_mapping); + DistMetaTensor modified_prior_dist(prior_dist.dims(), + new_prior_dist_attr); + VLOG(4) << "LabelSmoothInferSpmd call ElementwiseBinaryInferSpmd:"; + return ElementwiseBinaryInferSpmd(label, modified_prior_dist); + } + VLOG(4) << "LabelSmoothInferSpmd call ElementwiseBinaryInferSpmd:"; + return ElementwiseBinaryInferSpmd(label, prior_dist); + } + VLOG(4) << "LabelSmoothInferSpmd call ElementwiseUnaryInferSpmd:"; + SpmdInfo unary_spmd_info = ElementwiseUnaryInferSpmd(label); + unary_spmd_info.first.push_back(TensorDistAttr()); + return unary_spmd_info; +} + +SpmdInfo LabelSmoothGradInferSpmd(const DistMetaTensor& out_grad, + float epsilon) { + VLOG(4) << "LabelSmoothGradInferSpmd call ElementwiseUnaryGradInferSpmd:"; + return ElementwiseUnaryGradInferSpmd(out_grad); +} +} // namespace distributed +} // namespace phi diff --git a/paddle/phi/infermeta/spmd_rules/label_smooth.h b/paddle/phi/infermeta/spmd_rules/label_smooth.h new file mode 100644 index 00000000000000..ccb7195c125616 --- /dev/null +++ b/paddle/phi/infermeta/spmd_rules/label_smooth.h @@ -0,0 +1,31 @@ +/* Copyright (c) 2025 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. */ + +#pragma once + +#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h" +#include "paddle/phi/core/distributed/type_defs.h" + +namespace phi { +namespace distributed { + +SpmdInfo LabelSmoothInferSpmd(const DistMetaTensor& label, + const DistMetaTensor& prior_dist, + float epsilon); + +SpmdInfo LabelSmoothGradInferSpmd(const DistMetaTensor& out_grad, + float epsilon); + +} // namespace distributed +} // namespace phi diff --git a/paddle/phi/infermeta/spmd_rules/rules.cc b/paddle/phi/infermeta/spmd_rules/rules.cc index 2af63d6dbaa446..2f50c4b7cd5c9b 100644 --- a/paddle/phi/infermeta/spmd_rules/rules.cc +++ b/paddle/phi/infermeta/spmd_rules/rules.cc @@ -218,8 +218,8 @@ PD_REGISTER_SPMD_RULE( PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmdReverse)); PD_REGISTER_SPMD_RULE( label_smooth, - PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmd), - PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmdReverse)); + PD_INFER_SPMD(phi::distributed::LabelSmoothInferSpmd), + PD_INFER_SPMD(phi::distributed::LabelSmoothGradInferSpmd)); PD_REGISTER_SPMD_RULE( leaky_relu, PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmd), diff --git a/paddle/phi/infermeta/spmd_rules/rules.h b/paddle/phi/infermeta/spmd_rules/rules.h index a637dbc7ca6b42..88a3708841cdbc 100644 --- a/paddle/phi/infermeta/spmd_rules/rules.h +++ b/paddle/phi/infermeta/spmd_rules/rules.h @@ -46,6 +46,7 @@ limitations under the License. */ #include "paddle/phi/infermeta/spmd_rules/gather.h" #include "paddle/phi/infermeta/spmd_rules/gather_nd.h" #include "paddle/phi/infermeta/spmd_rules/index_select.h" +#include "paddle/phi/infermeta/spmd_rules/label_smooth.h" #include "paddle/phi/infermeta/spmd_rules/layer_norm.h" #include "paddle/phi/infermeta/spmd_rules/logsumexp.h" #include "paddle/phi/infermeta/spmd_rules/matmul.h" diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml index 4d402b97c2661d..af381b43ad694c 100644 --- a/paddle/phi/ops/yaml/backward.yaml +++ b/paddle/phi/ops/yaml/backward.yaml @@ -1859,6 +1859,7 @@ infer_meta : func : UnchangedInferMeta param : [out_grad] + spmd_rule : LabelSmoothGradInferSpmd kernel : func : label_smooth_grad diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml index e0ab4dd0632644..2779f2cf463e5f 100644 --- a/paddle/phi/ops/yaml/ops.yaml +++ b/paddle/phi/ops/yaml/ops.yaml @@ -2995,6 +2995,7 @@ infer_meta : func : UnchangedInferMeta param : [label] + spmd_rule : LabelSmoothInferSpmd kernel : func : label_smooth data_type : label diff --git a/test/auto_parallel/spmd_rules/CMakeLists.txt b/test/auto_parallel/spmd_rules/CMakeLists.txt index cd14c163b16415..154043a9cae45c 100644 --- a/test/auto_parallel/spmd_rules/CMakeLists.txt +++ b/test/auto_parallel/spmd_rules/CMakeLists.txt @@ -54,6 +54,7 @@ if(WITH_DISTRIBUTE) py_test_modules(test_cummax_rule MODULES test_cummax_rule) py_test_modules(test_cummin_rule MODULES test_cummin_rule) py_test_modules(test_argsort_rule MODULES test_argsort_rule) + py_test_modules(test_label_smooth_rule MODULES test_label_smooth_rule) py_test_modules(test_index_select_rule MODULES test_index_select_rule) py_test_modules(test_put_along_axis_rule MODULES test_put_along_axis_rule) endif() diff --git a/test/auto_parallel/spmd_rules/test_label_smooth_rule.py b/test/auto_parallel/spmd_rules/test_label_smooth_rule.py new file mode 100644 index 00000000000000..ce995e0b1a6781 --- /dev/null +++ b/test/auto_parallel/spmd_rules/test_label_smooth_rule.py @@ -0,0 +1,216 @@ +# Copyright (c) 2025 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. + +import unittest +from collections import OrderedDict + +from paddle.distributed.auto_parallel.static.dist_attribute import ( + DistTensorSpec, + TensorDistAttr, +) +from paddle.distributed.fleet import auto +from paddle.framework import core + + +class TestLabelSmoothSPMDRule(unittest.TestCase): + def setUp(self): + label_shape = [16, 16, 16] + out_shape = [16, 16, 16] + prior_dist_shape = [1, 16] + process_mesh = auto.ProcessMesh(mesh=[[0, 1], [2, 3]]) + + label_tensor_dist_attr = TensorDistAttr() + label_tensor_dist_attr.dims_mapping = [-1, -1, -1] + label_tensor_dist_attr.process_mesh = process_mesh + + self.label_dist_tensor_spec = DistTensorSpec( + label_shape, label_tensor_dist_attr + ) + prior_dist_tensor_dist_attr = TensorDistAttr() + prior_dist_tensor_dist_attr.dims_mapping = [-1, -1] + prior_dist_tensor_dist_attr.process_mesh = process_mesh + + self.prior_dist_dist_tensor_spec = DistTensorSpec( + prior_dist_shape, prior_dist_tensor_dist_attr + ) + + out_tensor_dist_attr = TensorDistAttr() + out_tensor_dist_attr.dims_mapping = [-1, -1, -1] + out_tensor_dist_attr.process_mesh = process_mesh + self.out_dist_tensor_spec = DistTensorSpec( + out_shape, out_tensor_dist_attr + ) + + self.rule = core.get_phi_spmd_rule("label_smooth") + self.attrs = OrderedDict() + self.attrs['epsilon'] = 0.1 + + def test_label_smooth_forward(self): + # [0, 1, -1], [Fake] --> [0, 1, -1], [Fake], [0, 1, -1] + self.label_dist_tensor_spec.set_dims_mapping([0, 1, -1]) + + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + DistTensorSpec(), + self.attrs['epsilon'], + ) + + self.assertEqual(len(result_dist_attrs), 2) + inferred_input_dist_attrs = result_dist_attrs[0] + inferred_output_dist_attrs = result_dist_attrs[1] + + self.assertEqual(len(inferred_input_dist_attrs), 2) + self.assertEqual(len(inferred_output_dist_attrs), 1) + + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, 1, -1]) + self.assertEqual(inferred_input_dist_attrs[1], TensorDistAttr()) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1]) + + # shape: [16, 16, 16], [1, 16]. [0, 1, -1], [-1, -1] --> [0, 1, -1], [-1, -1], [0, 1, -1] + self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, -1]) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + + self.assertEqual(len(result_dist_attrs), 2) + inferred_input_dist_attrs = result_dist_attrs[0] + inferred_output_dist_attrs = result_dist_attrs[1] + + self.assertEqual(len(inferred_input_dist_attrs), 2) + self.assertEqual(len(inferred_output_dist_attrs), 1) + + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, 1, -1]) + self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1]) + + # shape: [16, 16, 16], [1, 16]. [0, 1, -1], [-1, 1] --> [0, 1, -1], [-1, -1], [0, 1, -1] + self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, 1]) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + self.assertEqual(len(result_dist_attrs), 2) + inferred_input_dist_attrs = result_dist_attrs[0] + inferred_output_dist_attrs = result_dist_attrs[1] + self.assertEqual(len(inferred_input_dist_attrs), 2) + self.assertEqual(len(inferred_output_dist_attrs), 1) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, 1, -1]) + self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1]) + + # shape: [16, 16, 16], [1, 16]. [0, 1, -1], [-1, 0] --> [0, 1, -1], [-1, -1], [0, 1, -1] + self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, 0]) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + self.assertEqual(len(result_dist_attrs), 2) + inferred_input_dist_attrs = result_dist_attrs[0] + inferred_output_dist_attrs = result_dist_attrs[1] + self.assertEqual(len(inferred_input_dist_attrs), 2) + self.assertEqual(len(inferred_output_dist_attrs), 1) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, 1, -1]) + self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1]) + + # shape: [16, 16, 16], [1, 16]. [0, -1, 1], [-1, 1] --> [0, -1, 1], [-1, 1], [0, -1, 1] + self.label_dist_tensor_spec.set_dims_mapping([0, -1, 1]) + self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, 1]) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + self.assertEqual(len(result_dist_attrs), 2) + inferred_input_dist_attrs = result_dist_attrs[0] + inferred_output_dist_attrs = result_dist_attrs[1] + self.assertEqual(len(inferred_input_dist_attrs), 2) + self.assertEqual(len(inferred_output_dist_attrs), 1) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, 1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1]) + + # shape: [16, 16, 16], [1, 16]. [0, -1, 1], [-1, 0] --> [0, -1, 1], [-1, 1], [0, -1, 1] + self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, 0]) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + self.assertEqual(len(result_dist_attrs), 2) + inferred_input_dist_attrs = result_dist_attrs[0] + inferred_output_dist_attrs = result_dist_attrs[1] + self.assertEqual(len(inferred_input_dist_attrs), 2) + self.assertEqual(len(inferred_output_dist_attrs), 1) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, 1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1]) + + # shape: [16, 16, 16], [1, 16]. [0, -1, 1], [-1, -1] --> [0, -1, 1], [-1, 1], [0, -1, 1] + self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, -1]) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + self.assertEqual(len(result_dist_attrs), 2) + inferred_input_dist_attrs = result_dist_attrs[0] + inferred_output_dist_attrs = result_dist_attrs[1] + self.assertEqual(len(inferred_input_dist_attrs), 2) + self.assertEqual(len(inferred_output_dist_attrs), 1) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, 1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1]) + + # shape: [16, 16, 16], [16, 16, 16]. [0, 1, -1], [0, -1, 1] --> [0, 1, -1], [0, 1, -1], [0, 1, -1] + self.label_dist_tensor_spec.set_dims_mapping([0, 1, -1]) + self.prior_dist_dist_tensor_spec.shape = [16, 16, 16] + self.prior_dist_dist_tensor_spec.set_dims_mapping([0, -1, 1]) + result_dist_attrs = self.rule.infer_forward( + self.label_dist_tensor_spec, + self.prior_dist_dist_tensor_spec, + self.attrs['epsilon'], + ) + self.assertEqual(len(result_dist_attrs), 2) + inferred_input_dist_attrs = result_dist_attrs[0] + inferred_output_dist_attrs = result_dist_attrs[1] + self.assertEqual(len(inferred_input_dist_attrs), 2) + self.assertEqual(len(inferred_output_dist_attrs), 1) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, 1, -1]) + self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [0, 1, -1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1]) + + def test_label_smooth_backward(self): + # shape: [16, 16 ,16]. [0, -1, 1] --> [0, -1, 1], [0, -1, 1] + self.out_dist_tensor_spec.set_dims_mapping([0, -1, 1]) + result_dist_attrs = self.rule.infer_backward( + self.out_dist_tensor_spec, + self.attrs['epsilon'], + ) + + self.assertEqual(len(result_dist_attrs), 2) + inferred_input_dist_attrs = result_dist_attrs[0] + inferred_output_dist_attrs = result_dist_attrs[1] + self.assertEqual(len(inferred_input_dist_attrs), 1) + self.assertEqual(len(inferred_output_dist_attrs), 1) + self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1]) + self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc index bb7113641036f8..17c5eec4833be3 100644 --- a/test/cpp/auto_parallel/spmd_rule_test.cc +++ b/test/cpp/auto_parallel/spmd_rule_test.cc @@ -3147,7 +3147,71 @@ TEST(Unique, Ctor) { check_dim_mapping(forward_info.second[2], {-1}); check_dim_mapping(forward_info.second[3], {-1}); } +TEST(LabelSmooth, Ctor) { + std::vector mesh_shape = {2, 2}; + std::vector process_ids = {0, 1, 2, 3}; + std::vector 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}); + const phi::distributed::ArgDistAttr& attr = forward_info.first[1]; + if (paddle::holds_alternative(attr)) { + EXPECT_EQ(paddle::get(attr), + phi::distributed::TensorDistAttr()); + } else { + FAIL() << "forward_info.first[1] is not TensorDistAttr"; + } + 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 mesh_shape = {2, 2}; std::vector process_ids = {0, 1, 2, 3};