/* Copyright 2020 The TensorFlow 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 "tensorflow/compiler/xla/service/sharding_propagation.h"

#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "tensorflow/compiler/xla/protobuf_util.h"
#include "tensorflow/compiler/xla/service/hlo_matchers.h"
#include "tensorflow/compiler/xla/service/hlo_op_metadata.h"
#include "tensorflow/compiler/xla/service/hlo_parser.h"
#include "tensorflow/compiler/xla/status_macros.h"
#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"

namespace op = xla::testing::opcode_matchers;

namespace xla {
namespace {

using ShardingPropagationTest = HloTestBase;

void ClearMetadata(HloModule* module) {
  for (HloComputation* computation : module->computations()) {
    for (HloInstruction* instruction : computation->instructions()) {
      if (instruction->metadata().ByteSizeLong() != 0) {
        instruction->set_metadata(OpMetadata());
      }
      if (!instruction->has_sharding()) {
        continue;
      }
      instruction->set_sharding(instruction->sharding().WithoutMetadata());
    }
  }
}

struct MetadataTestParameter {
  explicit MetadataTestParameter(bool propagate_metadata, bool clear_metadata)
      : propagate_metadata(propagate_metadata),
        clear_metadata(clear_metadata) {}

  bool propagate_metadata = false;
  bool clear_metadata = false;
};

class ParameterizedMetadataTest
    : public HloTestBase,
      public ::testing::WithParamInterface<MetadataTestParameter> {};

std::string OpMetadataListToString(absl::Span<const OpMetadata> metadata) {
  std::vector<std::string> metadata_strings;
  metadata_strings.reserve(metadata.size());
  for (const OpMetadata& element : metadata) {
    metadata_strings.push_back(
        absl::StrCat("{", OpMetadataToString(element), "}"));
  }
  return absl::StrCat("{", absl::StrJoin(metadata_strings, ", "), "}");
}

class HloShardingMetadataMatcher
    : public ::testing::MatcherInterface<const HloSharding&> {
 public:
  explicit HloShardingMetadataMatcher(absl::Span<const OpMetadata> metadata)
      : metadata_(metadata.begin(), metadata.end()) {}

  bool MatchAndExplain(
      const HloSharding& sharding,
      ::testing::MatchResultListener* listener) const override {
    if (sharding.metadata().size() != metadata_.size()) {
      *listener << sharding.ToString(/*include_metadata=*/true)
                << " has incorrect sharding metadata (expected: "
                << OpMetadataListToString(metadata_) << ")";
      return false;
    }

    for (int i = 0, e = metadata_.size(); i < e; ++i) {
      if (!protobuf_util::ProtobufEquals(sharding.metadata()[i],
                                         metadata_[i])) {
        *listener << sharding.ToString(/*include_metadata=*/true)
                  << " has incorrect sharding metadata (expected: "
                  << OpMetadataListToString(metadata_) << ")";
        return false;
      }
    }

    return true;
  }

  void DescribeTo(std::ostream* os) const override {
    *os << OpMetadataListToString(metadata_);
  }

 private:
  std::vector<OpMetadata> metadata_;
};

::testing::Matcher<const HloSharding&> ShardingMetadata(
    absl::Span<const OpMetadata> metadata) {
  return ::testing::MakeMatcher(new HloShardingMetadataMatcher(metadata));
}

OpMetadata CreateMetadata(const std::string& op_name) {
  OpMetadata metadata;
  metadata.set_op_name(op_name);
  return metadata;
}

INSTANTIATE_TEST_SUITE_P(
    ShardingPropagation, ParameterizedMetadataTest,
    ::testing::Values(MetadataTestParameter(/*propagate_metadata=*/false,
                                            /*clear_metadata=*/false),
                      MetadataTestParameter(/*propagate_metadata=*/false,
                                            /*clear_metadata=*/true),
                      MetadataTestParameter(/*propagate_metadata=*/true,
                                            /*clear_metadata=*/false),
                      MetadataTestParameter(/*propagate_metadata=*/true,
                                            /*clear_metadata=*/true)),
    [](const ::testing::TestParamInfo<MetadataTestParameter>& info) {
      return absl::StrCat(info.param.propagate_metadata
                              ? "MetadataPropagation"
                              : "NoMetadataPropagation",
                          "_",
                          info.param.clear_metadata ? "NoMetadataInModule"
                                                    : "MetadataInModule");
    });

TEST_P(ParameterizedMetadataTest, ShardingMetadataFromInstruction) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %elementwise {
  %param0 = f32[5,7,11,13]{3,2,1,0} parameter(0),
    sharding={devices=[1,2,2,1]0,1,2,3},
    metadata={op_name="test"}
  ROOT %copy = f32[5,7,11,13]{3,2,1,0} copy(%param0)
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_EQ(changed,
            GetParam().propagate_metadata && !GetParam().clear_metadata);
  auto* instruction = FindInstruction(module.get(), "param0");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,2,1]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("test")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_F(ShardingPropagationTest, ShardingMetadataFromInstructionNoOverwrite) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %elementwise {
  %param0 = f32[5,7,11,13]{3,2,1,0} parameter(0),
    sharding={devices=[1,2,2,1]0,1,2,3 metadata={op_name="name"}},
    metadata={op_name="test"}
  ROOT %copy = f32[5,7,11,13]{3,2,1,0} copy(%param0)
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  TF_ASSERT_OK_AND_ASSIGN(bool changed,
                          ShardingPropagation(/*is_spmd=*/false,
                                              /*propagate_metadata=*/true)
                              .Run(module.get()));
  EXPECT_FALSE(changed);
  auto* instruction = FindInstruction(module.get(), "param0");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,2,1]0,1,2,3}"));
  EXPECT_THAT(instruction->sharding(),
              ShardingMetadata({CreateMetadata("name")}));
}

TEST_F(ShardingPropagationTest, ShardingMetadataFromInstructionNoMetadata) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %elementwise {
  %param0 = f32[5,7,11,13]{3,2,1,0} parameter(0),
    sharding={devices=[1,2,2,1]0,1,2,3 metadata={op_name="name"}}
  ROOT %copy = f32[5,7,11,13]{3,2,1,0} copy(%param0)
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  TF_ASSERT_OK_AND_ASSIGN(bool changed,
                          ShardingPropagation(/*is_spmd=*/false,
                                              /*propagate_metadata=*/true)
                              .Run(module.get()));
  EXPECT_FALSE(changed);
  auto* instruction = FindInstruction(module.get(), "param0");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,2,1]0,1,2,3}"));
  EXPECT_THAT(instruction->sharding(),
              ShardingMetadata({CreateMetadata("name")}));
}

TEST_F(ShardingPropagationTest, ShardingNoMetadataAndInstructionNoMetadata) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %elementwise {
  %param0 = f32[5,7,11,13]{3,2,1,0} parameter(0),
    sharding={devices=[1,2,2,1]0,1,2,3}
  ROOT %copy = f32[5,7,11,13]{3,2,1,0} copy(%param0)
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  TF_ASSERT_OK_AND_ASSIGN(bool changed,
                          ShardingPropagation(/*is_spmd=*/false,
                                              /*propagate_metadata=*/true)
                              .Run(module.get()));
  EXPECT_FALSE(changed);
  auto* instruction = FindInstruction(module.get(), "param0");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,2,1]0,1,2,3}"));
  EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
}

TEST_P(ParameterizedMetadataTest, ElementwiseOperationForwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %elementwise {
  %param0 = f32[5,7,11,13]{3,2,1,0} parameter(0),
    sharding={devices=[1,2,2,1]0,1,2,3 metadata={op_name="a"}}
  %param1 = f32[5,7,11,13]{3,2,1,0} parameter(1)
  %add = f32[5,7,11,13]{3,2,1,0} add(%param0, %param1)
  ROOT %copy = f32[5,7,11,13]{3,2,1,0} copy(%add)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "add");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,2,1]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ElementwiseOperationBackwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %elementwise {
  %param0 = f32[5,7,11,13]{3,2,1,0} parameter(0)
  %param1 = f32[5,7,11,13]{3,2,1,0} parameter(1)
  %add = f32[5,7,11,13]{3,2,1,0} add(%param0, %param1)
  ROOT %copy = f32[5,7,11,13]{3,2,1,0} copy(%add),
    sharding={devices=[1,2,2,1]0,1,2,3 metadata={op_name="a"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "add");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,2,1]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

// Regression Test for b/129569657.
TEST_P(ParameterizedMetadataTest, BroadcastForwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %broadcast {
  %param0 = f32[3,2048,2048]{2,1,0} parameter(0),
    sharding={devices=[1,2,2]0,1,2,3 metadata={op_name="a"}}
  %broadcast = f32[3,2048,2048,3]{3,2,1,0} broadcast(%param0), dimensions={0,1,2}
  ROOT %copy = f32[3,2048,2048,3]{3,2,1,0} copy(%broadcast)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "broadcast");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,2,1]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, BroadcastBackwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %broadcast {
  %param0 = f32[13]{0} parameter(0)
  %broadcast = f32[5,7,11,13]{3,2,1,0} broadcast(%param0), dimensions={3}
  ROOT %copy = f32[5,7,11,13]{3,2,1,0} copy(%broadcast),
    sharding={devices=[1,2,2,1]0,1,2,3 metadata={op_name="a"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "broadcast");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,2,1]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, BroadcastForwardPartial) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %broadcast {
  %param0 = f32[3,2048]parameter(0),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="a"}}
  %broadcast = f32[3,2048,3] broadcast(%param0), dimensions={0,1}
  ROOT %copy = f32[3,2048,3] copy(%broadcast)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "broadcast");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(
      instruction,
      op::Sharding("{devices=[1,2,1,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, BroadcastMerge) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %broadcast {
  %param0 = f32[3,2048]parameter(0),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="a"}}
  %broadcast = f32[3,2048,3] broadcast(%param0), dimensions={0,1}
  ROOT %copy = f32[3,2048,3] copy(%broadcast),
    sharding={devices=[1,1,2,2]0,2,1,3 last_tile_dim_replicate metadata={op_name="b"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "broadcast");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,2]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a"), CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, BroadcastUser) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %broadcast {
  %param0 = f32[24,8]{0,1} parameter(0)
  %copy = f32[24,8]{0,1} copy(%param0)
  ROOT %broadcast = f32[4,24,6,8]{3,2,1,0} broadcast(%copy), dimensions={1,3},
    sharding={devices=[1,2,1,4]0,1,2,3,4,5,6,7 metadata={op_name="a"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "copy");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,4]0,1,2,3,4,5,6,7}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, BroadcastUserPartial) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %broadcast {
  %param0 = f32[24,8]{0,1} parameter(0)
  %copy = f32[24,8]{0,1} copy(%param0)
  ROOT %broadcast = f32[4,24,6,8] broadcast(%copy), dimensions={1,3},
    sharding={devices=[4,2,1,1]0,1,2,3,4,5,6,7 metadata={op_name="a"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "copy");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(
      instruction,
      op::Sharding("{devices=[2,1,4]0,2,4,6,1,3,5,7 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, MaximalReduceForwardPass) {
  const char* const hlo_string = R"(
HloModule module
%add {
  %lhs = f32[] parameter(0)
  %rhs = f32[] parameter(1)
  ROOT %add = f32[] add(%lhs, %rhs)
}
ENTRY %reduce {
  %param0 = f32[5,7,11,13]{3,2,1,0} parameter(0),
    sharding={devices=[1,2,2,1]0,1,2,3 metadata={op_name="a"}}
  %init = f32[] parameter(1)
  %reduce = f32[5,7]{1,0} reduce(%param0, %init), dimensions={2,3}, to_apply=%add
  ROOT %copy = f32[5,7]{0,1} copy(%reduce)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "reduce");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{replicated}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ShardedReduceForwardPass) {
  const char* const hlo_string = R"(
HloModule module
%add {
  %lhs = f32[] parameter(0)
  %rhs = f32[] parameter(1)
  ROOT %add = f32[] add(%lhs, %rhs)
}
ENTRY %reduce {
  %param0 = f32[5,7,11,13]{3,2,1,0} parameter(0),
    sharding={devices=[1,2,2,1]0,1,2,3 metadata={op_name="a"}}
  %init = f32[] parameter(1)
  %reduce = f32[7,11]{1,0} reduce(%param0, %init), dimensions={0,3}, to_apply=%add
  ROOT %copy = f32[7,11]{0,1} copy(f32[7,11]{1,0} %reduce)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "reduce");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,2]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ReducePartiallyOnTiledDims) {
  const char* const hlo_string = R"(
HloModule module
%add {
  %lhs = f32[] parameter(0)
  %rhs = f32[] parameter(1)
  ROOT %add = f32[] add(%lhs, %rhs)
}
ENTRY %reduce {
  %param0 = f32[8,8] parameter(0),
    sharding={devices=[2,2]0,1,2,3 metadata={op_name="a"}}
  %init = f32[] parameter(1)
  %reduce = f32[8] reduce(%param0, %init), dimensions={0}, to_apply=%add
  ROOT %copy = f32[8] copy(%reduce)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "reduce");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[2,2]0,2,1,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ReducePartiallyOnTiledDims2) {
  const char* const hlo_string = R"(
HloModule module
%add {
  %lhs = f32[] parameter(0)
  %rhs = f32[] parameter(1)
  ROOT %add = f32[] add(%lhs, %rhs)
}
ENTRY %reduce {
  %param0 = f32[8,8] parameter(0),
    sharding={devices=[2,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate metadata={op_name="a"}}
  %init = f32[] parameter(1)
  %reduce = f32[8] reduce(%param0, %init), dimensions={0}, to_apply=%add
  ROOT %copy = f32[8] copy(%reduce)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "reduce");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(
      instruction,
      op::Sharding("{devices=[2,4]0,1,4,5,2,3,6,7 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ReducePartiallyBackward) {
  const char* const hlo_string = R"(
HloModule module
%add {
  %lhs = f32[] parameter(0)
  %rhs = f32[] parameter(1)
  ROOT %add = f32[] add(%lhs, %rhs)
}
ENTRY %reduce {
  %param0 = f32[8,8] parameter(0)
  %input = f32[8,8] copy(%param0)
  %init = f32[] parameter(1)
  %reduce = f32[8] reduce(%input, %init), dimensions={0}, to_apply=%add,
    sharding={devices=[2,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="a"}}
  ROOT %copy = f32[8] copy(%reduce)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "input");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ShardedTupleReduceForwardAndBackwardPass) {
  const char* const hlo_string = R"(
HloModule module

%minmax_func {
  %lhs_value = f32[] parameter(0)
  %rhs_value = f32[] parameter(2)
  %compare.2 = pred[] compare(%lhs_value, %rhs_value), direction=GT
  %select.4 = f32[] select(%compare.2, %lhs_value, %rhs_value)
  %lhs_index = s32[] parameter(1)
  %rhs_index = s32[] parameter(3)
  %select.5 = s32[] select(%compare.2, %lhs_index, %rhs_index)
  ROOT %tuple.2 = (f32[], s32[]) tuple(%select.4, %select.5)
}

ENTRY %main {
  %param0 = f32[28,10] parameter(0)
  %param1 = s32[28,10] parameter(1), sharding={devices=[2,1]0,1 metadata={op_name="a"}}
  %copy_param0 = f32[28,10] copy(%param0)
  %init0 = f32[] parameter(2)
  %init1 = s32[] parameter(3)
  %reduce = (f32[28], s32[28]) reduce(%copy_param0, %param1, %init0, %init1),
    dimensions={1}, to_apply=%minmax_func
  %gte0 = f32[28] get-tuple-element(%reduce), index=0
  %gte1 = s32[28] get-tuple-element(%reduce), index=1
  %copy0 = f32[28] copy(%gte0)
  %copy1 = s32[28] copy(%gte1)
  ROOT %tuple = (f32[28], s32[28]) tuple(%copy0, %copy1)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* reduce = FindInstruction(module.get(), "reduce");
  ASSERT_NE(reduce, nullptr);
  EXPECT_THAT(reduce, op::Sharding("{{devices=[2]0,1},{devices=[2]0,1}}"));
  auto* copy_param0 = FindInstruction(module.get(), "copy_param0");
  ASSERT_NE(copy_param0, nullptr);
  EXPECT_THAT(copy_param0, op::Sharding("{devices=[2,1]0,1}"));
  for (const HloSharding& sharding :
       {copy_param0->sharding(), reduce->sharding().tuple_elements()[0],
        reduce->sharding().tuple_elements()[1]}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(sharding, ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(sharding, ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, GetTupleElementForwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %gte {
  %param0 = f32[5,7,11,13]{3,2,1,0} parameter(0)
  %tuple = (f32[5,7,11,13]{3,2,1,0}, f32[5,7,11,13]{3,2,1,0}) tuple(
    %param0, %param0)
  %tuple.1 = (f32[5,7,11,13]{3,2,1,0},
              (f32[5,7,11,13]{3,2,1,0}, f32[5,7,11,13]{3,2,1,0})) tuple(
    %param0, %tuple),
    sharding={{devices=[1,2,2,1]0,1,2,3 metadata={op_name="a"}},
              {replicated metadata={op_name="b"}},
              {devices=[1,2,2,1]0,1,2,3 metadata={op_name="c"}}}
  %gte = f32[5,7,11,13]{3,2,1,0} get-tuple-element(%tuple.1), index=0
  %gte.1 = (f32[5,7,11,13]{3,2,1,0}, f32[5,7,11,13]{3,2,1,0}) get-tuple-element(
    %tuple.1), index=1
  %gte.2 = f32[5,7,11,13]{3,2,1,0} get-tuple-element(%gte.1), index=0
  ROOT %copy = f32[5,7,11,13]{3,2,1,0} copy(%gte.2)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* gte = FindInstruction(module.get(), "gte");
  ASSERT_NE(gte, nullptr);
  EXPECT_THAT(gte, op::Sharding("{devices=[1,2,2,1]0,1,2,3}"));
  auto* gte1 = FindInstruction(module.get(), "gte.1");
  ASSERT_NE(gte1, nullptr);
  EXPECT_THAT(gte1, op::Sharding("{{replicated}, {devices=[1,2,2,1]0,1,2,3}}"));
  auto* gte2 = FindInstruction(module.get(), "gte.2");
  ASSERT_NE(gte2, nullptr);
  EXPECT_THAT(gte2, op::Sharding("{replicated}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(gte->sharding(), ShardingMetadata({CreateMetadata("a")}));
    EXPECT_THAT(gte1->sharding().tuple_elements()[0],
                ShardingMetadata({CreateMetadata("b")}));
    EXPECT_THAT(gte1->sharding().tuple_elements()[1],
                ShardingMetadata({CreateMetadata("c")}));
    EXPECT_THAT(gte2->sharding(), ShardingMetadata({CreateMetadata("b")}));
  } else {
    for (const HloSharding& sharding :
         {gte->sharding(), gte1->sharding().tuple_elements()[0],
          gte1->sharding().tuple_elements()[1], gte2->sharding()}) {
      EXPECT_THAT(sharding, ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, TupleForwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %tuple {
  %param0 = f32[5,7,11,13]{3,2,1,0} parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %param1 = f32[5,7,11,13]{3,2,1,0} parameter(1),
    sharding={devices=[1,2,2,1]0,1,2,3 metadata={op_name="b"}}
  %param2 = f32[5,7,11,13]{3,2,1,0} parameter(2)
  %tuple = (f32[5,7,11,13]{3,2,1,0}, f32[5,7,11,13]{3,2,1,0}) tuple(
    %param1, %param2)
  %tuple.1 = (f32[5,7,11,13]{3,2,1,0},
              (f32[5,7,11,13]{3,2,1,0}, f32[5,7,11,13]{3,2,1,0})) tuple(
    %param0, %tuple)
  ROOT %copy = (f32[5,7,11,13]{3,2,1,0},
                (f32[5,7,11,13]{3,2,1,0}, f32[5,7,11,13]{3,2,1,0})) copy(
    %tuple.1)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* tuple = FindInstruction(module.get(), "tuple");
  ASSERT_NE(tuple, nullptr);
  EXPECT_THAT(tuple, op::Sharding("{{devices=[1,2,2,1]0,1,2,3},"
                                  " {replicated}}"));
  auto* tuple1 = FindInstruction(module.get(), "tuple.1");
  ASSERT_NE(tuple1, nullptr);
  EXPECT_THAT(tuple1, op::Sharding("{{replicated},"
                                   " {devices=[1,2,2,1]0,1,2,3},"
                                   " {replicated}}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(tuple->sharding().tuple_elements()[0],
                ShardingMetadata({CreateMetadata("b")}));
    EXPECT_THAT(tuple->sharding().tuple_elements()[1], ShardingMetadata({}));
    EXPECT_THAT(tuple1->sharding().tuple_elements()[0], ShardingMetadata({}));
    EXPECT_THAT(tuple1->sharding().tuple_elements()[1],
                ShardingMetadata({CreateMetadata("b")}));
    EXPECT_THAT(tuple1->sharding().tuple_elements()[2], ShardingMetadata({}));
  } else {
    for (const HloSharding& tuple_sharding :
         {tuple->sharding(), tuple1->sharding()}) {
      for (const HloSharding& sub_sharding : tuple_sharding.tuple_elements()) {
        EXPECT_THAT(sub_sharding, ShardingMetadata({}));
      }
    }
  }
}

TEST_P(ParameterizedMetadataTest, ForwardConvolutionForwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %lhs = f32[5,7,11,13]{3,2,1,0} parameter(0),
    sharding={devices=[2,2,2,1]0,1,2,3,4,5,6,7 metadata={op_name="a"}}
  %rhs = f32[3,3,13,17]{3,2,1,0} parameter(1)
  %convolution = f32[5,7,11,17]{3,2,1,0} convolution(%lhs, %rhs),
    window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f
  ROOT %copy = f32[5,7,11,17]{3,2,1,0} copy(%convolution)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "convolution");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,2,2,1]0,1,2,3,4,5,6,7}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ForwardConvolutionLargeDilationForwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %lhs = f32[8,64,2]{2,1,0} parameter(0),
    sharding={devices=[1,4,1]0,1,2,3 metadata={op_name="a"}}
  %rhs = f32[3,2,2]{2,1,0} parameter(1)
  %convolution = f32[8,32,2]{2,1,0} convolution(%lhs, %rhs),
    window={size=3 rhs_dilate=16}, dim_labels=b0f_0io->b0f
  ROOT %copy = f32[8,32,2]{2,1,0} copy(%convolution)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "convolution");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,4,1]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ForwardConvolution3DSmallKernel) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %lhs = bf16[32,32,8,7,128]{4,3,2,1,0} parameter(0),
    sharding={devices=[1,4,1,1,1]0,1,2,3 metadata={op_name="a"}}
  %rhs = bf16[3,3,3,128,256]{4,3,2,1,0} parameter(1)
  %convolution = bf16[16,16,8,3,256]{4,3,2,1,0}
    convolution(bf16[32,32,8,7,128]{4,3,2,1,0} %lhs,
    bf16[3,3,3,128,256]{4,3,2,1,0} %rhs),
    window={size=3x3x3 stride=2x2x2 pad=1_1x1_1x0_0},
    dim_labels=01b2f_012io->01b2f
  ROOT %copy = bf16[16,16,8,3,256]{4,3,2,1,0} copy(%convolution)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "convolution");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,4,1,1,1]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, TransposeForwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %transpose {
  %param = f32[7,11,13]{2,1,0} parameter(0),
    sharding={devices=[2,1,2]0,1,2,3 metadata={op_name="a"}}
  %transpose = f32[11,13,7]{2,1,0} transpose(%param), dimensions={1,2,0}
  ROOT %copy = f32[11,13,7]{2,1,0} copy(%transpose)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "transpose");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,2]0,2,1,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, TransposeBackwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %transpose {
  %param = f32[7,11,13]{2,1,0} parameter(0)
  %copy = f32[7,11,13]{2,1,0} copy(%param)
  ROOT %transpose = f32[11,13,7]{2,1,0} transpose(%copy), dimensions={1,2,0},
    sharding={devices=[1,2,2]0,1,2,3 metadata={op_name="a"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "copy");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,1,2]0,2,1,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ReshapeForwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %reshape {
  %param0 = f32[1430,1]{1,0} parameter(0),
    sharding={devices=[2,1]0,1 metadata={op_name="a"}}
  %reshape = f32[10,11,13]{2,1,0} reshape(%param0)
  ROOT %copy = f32[10,11,13]{2,1,0} copy(%reshape)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "reshape");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,1,1]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ReshapeBackwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %reshape {
  %param0 = f32[2002,1]{1,0} parameter(0)
  %copy = f32[2002,1]{1,0} copy(f32[2002,1]{1,0} %param0)
  ROOT %reshape = f32[14,11,13]{2,1,0} reshape(%copy),
    sharding={devices=[2,1,1]0,1 metadata={op_name="a"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "copy");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,1]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, PadForwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %pad {
  %input = f32[11,17]{1,0} parameter(0),
    sharding={devices=[2,2]0,1,2,3 metadata={op_name="a"}}
  %pad_value = f32[] parameter(1)
  %pad = f32[27,51]{1,0} pad(%input, %pad_value), padding=2_4_1x1_1_2
  ROOT %copy = f32[27,51]{1,0} copy(%pad)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "pad");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,2]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, PadBackwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %pad {
  %input = f32[11,17]{1,0} parameter(0)
  %copy = f32[11,17]{1,0} copy(%input)
  %pad_value = f32[] parameter(1)
  %pad = f32[27,51]{1,0} pad(%copy, %pad_value), padding=2_4_1x1_1_2,
    sharding={devices=[2,2]0,1,2,3 metadata={op_name="a"}}
  ROOT %result = f32[27,51]{1,0} copy(%pad)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "copy");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,2]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, PartialReplicatedPadForwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %pad {
  %input = f32[11,17]{1,0} parameter(0),
    sharding={devices=[2,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate metadata={op_name="a"}}
  %pad_value = f32[] parameter(1)
  %pad = f32[27,51]{1,0} pad(%input, %pad_value), padding=2_4_1x1_1_2
  ROOT %copy = f32[27,51]{1,0} copy(%pad)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "pad");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(
      instruction,
      op::Sharding("{devices=[2,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ShardedPreferredOverReplicated) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %replicated {
  %param0 = f32[5,7,11,13]{3,2,1,0} parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %copy = f32[5,7,11,13]{3,2,1,0} copy(%param0)
  %param1 = f32[5,7,11,13]{3,2,1,0} parameter(1),
    sharding={devices=[1,2,2,1]0,1,2,3 metadata={op_name="b"}}
  %copy.1 = f32[5,7,11,13]{3,2,1,0} copy(%param1)
  %add = f32[5,7,11,13]{3,2,1,0} add(%copy, %copy.1)
  ROOT %copy.2 = f32[5,7,11,13]{3,2,1,0} copy(%add)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* copy = FindInstruction(module.get(), "copy");
  ASSERT_NE(copy, nullptr);
  EXPECT_THAT(copy, op::Sharding("{devices=[1,2,2,1]0,1,2,3}"));
  auto* copy1 = FindInstruction(module.get(), "copy.1");
  ASSERT_NE(copy1, nullptr);
  EXPECT_THAT(copy1, op::Sharding("{devices=[1,2,2,1]0,1,2,3}"));
  auto* add = FindInstruction(module.get(), "add");
  ASSERT_NE(add, nullptr);
  EXPECT_THAT(add, op::Sharding("{devices=[1,2,2,1]0,1,2,3}"));
  for (const HloSharding& sharding :
       {copy->sharding(), copy1->sharding(), add->sharding()}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(sharding, ShardingMetadata({CreateMetadata("b")}));
    } else {
      EXPECT_THAT(sharding, ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, PartialReplicateReshapeForwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %reshape {
  %param0 = f32[1430,1]{1,0} parameter(0),
    sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="a"}}
  %reshape = f32[10,11,13]{2,1,0} reshape(%param0)
  ROOT %copy = f32[10,11,13]{2,1,0} copy(%reshape)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "reshape");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(
      instruction,
      op::Sharding("{devices=[2,1,1,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, PartialReplicateReshapeBackwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %reshape {
  %param0 = f32[2002,1]{1,0} parameter(0)
  %copy = f32[2002,1]{1,0} copy(f32[2002,1]{1,0} %param0)
  ROOT %reshape = f32[14,11,13]{2,1,0} reshape(%copy),
    sharding={devices=[2,1,1,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="a"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "copy");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, DontShardTuplesIfAllInputIsMaximal) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %tuple {
  %param0 = f32[5,7,11,13]{3,2,1,0} parameter(0),
    sharding={maximal device=0 metadata={op_name="a"}}
  %param1 = f32[5,7,11,13]{3,2,1,0} parameter(1),
    sharding={maximal device=1 metadata={op_name="b"}}
  %tuple = (f32[5,7,11,13]{3,2,1,0}, f32[5,7,11,13]{3,2,1,0}) tuple(
    %param0, %param1)
  ROOT %copy = (f32[5,7,11,13]{3,2,1,0}, f32[5,7,11,13]{3,2,1,0}) copy(%tuple)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_EQ(changed,
            !GetParam().propagate_metadata && !GetParam().clear_metadata);
  auto* instruction = FindInstruction(module.get(), "tuple");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::NoSharding());
}

TEST_P(ParameterizedMetadataTest, ValidConvolution) {
  const char* const hlo_string = R"(
HloModule module

ENTRY conv {
  %lhs = f32[13,17,19]{2,1,0} parameter(0),
    sharding={devices=[1,2,1]0,1 metadata={op_name="a"}}
  %rhs = f32[19,5,19]{2,1,0} parameter(1)
  %conv = f32[13,13,19]{2,1,0} convolution(%lhs, %rhs),
    window={size=5}, dim_labels=b0f_i0o->b0f
  ROOT %tuple = (f32[13,13,19]{2,1,0}) tuple(%conv)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "conv");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,1]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, StridedSlice) {
  const char* const hlo_string = R"(
HloModule module

ENTRY %slice {
  %param = f32[17,13]{1,0} parameter(0),
    sharding={devices=[2,1]0,1 metadata={op_name="a"}}
  %slice = f32[7,5]{1,0} slice(%param), slice={[1:15:2], [5:10:1]}
  ROOT %tuple = (f32[7,5]{1,0}) tuple(%slice)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "slice");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,1]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, PartialReplicatedStridedSlice) {
  const char* const hlo_string = R"(
HloModule module

ENTRY %slice {
  %param = f32[17,13]{1,0} parameter(0),
    sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="a"}}
  %slice = f32[7,5]{1,0} slice(%param), slice={[1:15:2], [5:10:1]}
  ROOT %tuple = (f32[7,5]{1,0}) tuple(%slice)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "slice");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ReduceWindowBackwardPass) {
  const char* const hlo_string = R"(
HloModule module
%add (lhs: f32[], rhs: f32[]) -> f32[] {
  %lhs = f32[] parameter(0)
  %rhs = f32[] parameter(1)
  ROOT %add = f32[] add(%lhs, %rhs)
}
ENTRY %reduce_window {
  %param = f32[13,17]{1,0} parameter(0)
  %param.copy = f32[13,17]{1,0} copy(%param)
  %init = f32[] parameter(1)
  ROOT %reduce-window = f32[7,17]{1,0} reduce-window(%param.copy, %init),
    window={size=3x2 stride=2x1 pad=1_1x0_1}, to_apply=%add,
    sharding={devices=[2,1]0,1 metadata={op_name="a"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* param_copy = FindInstruction(module.get(), "param.copy");
  ASSERT_NE(param_copy, nullptr);
  EXPECT_THAT(param_copy, op::Sharding("{devices=[2,1]0,1}"));
  auto* reduce_window = FindInstruction(module.get(), "reduce-window");
  ASSERT_NE(reduce_window, nullptr);
  EXPECT_THAT(reduce_window, op::Sharding("{devices=[2,1]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(param_copy->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
    EXPECT_THAT(reduce_window->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(param_copy->sharding(), ShardingMetadata({}));
    EXPECT_THAT(reduce_window->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, VariadicReduceWindowBackwardPass) {
  const char* const hlo_string = R"(
HloModule module
%add (a: f32[], b: s32[], c: f32[], d: s32[]) -> (f32[], s32[]) {
  %a = f32[] parameter(0)
  %b = s32[] parameter(1)
  %c = f32[] parameter(2)
  %d = s32[] parameter(3)
  %add.0 = f32[] add(%a, %c)
  %add.1 = s32[] add(%b, %d)
  ROOT %t = tuple(%add.0, %add.1)
}
ENTRY %reduce_window {
  %param.0 = f32[13,17]{1,0} parameter(0)
  %param.0.copy = f32[13,17]{1,0} copy(%param.0)
  %param.1 = s32[13,17]{1,0} parameter(1)
  %param.1.copy = s32[13,17]{1,0} copy(%param.1)
  %init.0 = f32[] parameter(2)
  %init.1 = s32[] parameter(3)
  ROOT %reduce-window = (f32[7,17]{1,0}, s32[7,17]{1,0}) reduce-window(%param.0.copy, %param.1.copy, %init.0, %init.1),
    window={size=3x2 stride=2x1 pad=1_1x0_1}, to_apply=%add,
    sharding={{devices=[2,1]0,1 metadata={op_name="a"}}, {devices=[2,1]0,1 metadata={op_name="b"}}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* param_0_copy = FindInstruction(module.get(), "param.0.copy");
  ASSERT_NE(param_0_copy, nullptr);
  EXPECT_THAT(param_0_copy, op::Sharding("{devices=[2,1]0,1}"));
  auto* param_1_copy = FindInstruction(module.get(), "param.1.copy");
  ASSERT_NE(param_1_copy, nullptr);
  EXPECT_THAT(param_1_copy, op::Sharding("{devices=[2,1]0,1}"));
  auto* reduce_window = FindInstruction(module.get(), "reduce-window");
  ASSERT_NE(reduce_window, nullptr);
  EXPECT_THAT(reduce_window,
              op::Sharding("{{devices=[2,1]0,1}, {devices=[2,1]0,1}}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(param_0_copy->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
    EXPECT_THAT(param_1_copy->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
    EXPECT_THAT(reduce_window->sharding().tuple_elements()[0],
                ShardingMetadata({CreateMetadata("a")}));
    EXPECT_THAT(reduce_window->sharding().tuple_elements()[1],
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(param_0_copy->sharding(), ShardingMetadata({}));
    EXPECT_THAT(param_1_copy->sharding(), ShardingMetadata({}));
    EXPECT_THAT(reduce_window->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ReplicatedConvolutionLhs) {
  const char* const hlo_string = R"(
HloModule module

ENTRY conv {
  %lhs = f32[3,2,3]{2,1,0} parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %rhs = f32[2,2,1]{2,1,0} parameter(1)
  %conv = f32[3,2,3]{2,1,0} convolution(%lhs, %rhs),
    window={size=1}, dim_labels=bf0_oi0->bf0
  ROOT %tuple = f32[3,2,3]{2,1,0} tuple(%conv)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* lhs = FindInstruction(module.get(), "lhs");
  ASSERT_NE(lhs, nullptr);
  EXPECT_THAT(lhs, op::Sharding("{replicated}"));
  auto* conv = FindInstruction(module.get(), "conv");
  ASSERT_NE(conv, nullptr);
  EXPECT_THAT(conv, op::Sharding("{replicated}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(lhs->sharding(), ShardingMetadata({CreateMetadata("a")}));
    EXPECT_THAT(conv->sharding(), ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(lhs->sharding(), ShardingMetadata({}));
    EXPECT_THAT(conv->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ConvolutionShardedFeature) {
  const char* const hlo_string = R"(
HloModule module

ENTRY conv {
  %lhs = f32[3,2,3]{2,1,0} parameter(0),
    sharding={devices=[1,2,1]0,1 metadata={op_name="a"}}
  %rhs = f32[2,2,1]{2,1,0} parameter(1)
  %conv = f32[3,2,3]{2,1,0} convolution(%lhs, %rhs),
    window={size=1}, dim_labels=bf0_oi0->bf0
  ROOT %tuple = (f32[3,2,3]{2,1,0}) tuple(%conv)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "conv");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{replicated}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ConvolutionDifferentDimensionNumbers) {
  const char* const hlo_string = R"(
HloModule module

ENTRY conv {
  %lhs = f32[8,16,512] parameter(0),
    sharding={devices=[1,2,1]0,1 metadata={op_name="a"}}
  %rhs = f32[8,2,512] parameter(1)
  %conv = f32[3,512,512] convolution(%lhs, %rhs),
    window={size=2 stride=5},
    dim_labels=f0b_i0o->0bf
  ROOT %tuple = f32[3,512,512] tuple(%conv)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "conv");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,1,1]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, Concatenate) {
  const char* const hlo_string = R"(
HloModule module

ENTRY %concat {
  %param.0 = f32[5,7] parameter(0),
    sharding={devices=[2,1]0,1 metadata={op_name="a"}}
  %param.1 = f32[5,9] parameter(1),
    sharding={devices=[2,1]0,1 metadata={op_name="b"}}
  %concat = f32[5,16] concatenate(%param.0, %param.1),
    dimensions={1}
  ROOT %tuple = (f32[5,16]) tuple(%concat)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "concat");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,1]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, TupleBackwardPass) {
  const char* const hlo_string = R"(
HloModule module

ENTRY %tuple {
  %param.0 = f32[1] parameter(0)
  %param.1 = f32[3] parameter(1)
  %copy.0 = f32[1] copy(%param.0)
  %copy.1 = f32[3] copy(%param.1)
  ROOT %tuple = (f32[1], f32[3]) tuple(%copy.0, %copy.1),
    sharding={{replicated metadata={op_name="a"}},
              {devices=[2]0,1 metadata={op_name="b"}}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* copy0 = FindInstruction(module.get(), "copy.0");
  ASSERT_NE(copy0, nullptr);
  EXPECT_THAT(copy0, op::Sharding("{replicated}"));
  auto* copy1 = FindInstruction(module.get(), "copy.1");
  ASSERT_NE(copy1, nullptr);
  EXPECT_THAT(copy1, op::Sharding("{devices=[2]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(copy0->sharding(), ShardingMetadata({CreateMetadata("a")}));
    EXPECT_THAT(copy1->sharding(), ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(copy0->sharding(), ShardingMetadata({}));
    EXPECT_THAT(copy1->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, AllReduce) {
  const char* const hlo_string = R"(
HloModule module

%add (lhs: f32[], rhs: f32[]) -> f32[] {
  %add_lhs = f32[] parameter(0)
  %add_rhs = f32[] parameter(1)
  ROOT %add = f32[] add(f32[] %add_lhs, f32[] %add_rhs)
}

ENTRY %entry {
  %param.0 = f32[3] parameter(0)
  %param.1 = f32[3] parameter(1)

  %copy_f_t = f32[3] copy(%param.1),
    sharding={devices=[2]0,1 metadata={op_name="a"}}
  %crs_f.tiled = f32[3] all-reduce(%copy_f_t), to_apply=%add
  %crs_f.none = f32[3] all-reduce(%copy_f_t), to_apply=%add,
    channel_id=1

  %crs_b.replicated = f32[3] all-reduce(%param.0), to_apply=%add
  %copy_b_r = f32[3] copy(%crs_b.replicated),
    sharding={replicated metadata={op_name="b"}}

  ROOT %tuple = (f32[3], f32[3], f32[3], f32[3]) tuple(
    %crs_f.tiled, crs_f.none, %copy_b_r)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* crs_f_tiled = FindInstruction(module.get(), "crs_f.tiled");
  ASSERT_NE(crs_f_tiled, nullptr);
  EXPECT_THAT(crs_f_tiled, op::Sharding("{devices=[2]0,1}"));
  auto* crs_f_none = FindInstruction(module.get(), "crs_f.none");
  ASSERT_NE(crs_f_none, nullptr);
  EXPECT_THAT(crs_f_none, op::NoSharding());
  auto* crs_b_replicated = FindInstruction(module.get(), "crs_b.replicated");
  ASSERT_NE(crs_b_replicated, nullptr);
  EXPECT_THAT(crs_b_replicated, op::Sharding("{replicated}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(crs_f_tiled->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
    EXPECT_THAT(crs_b_replicated->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(crs_f_tiled->sharding(), ShardingMetadata({}));
    EXPECT_THAT(crs_b_replicated->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, While) {
  const char* const hlo_string = R"(
HloModule module

%cond {
  %vars.cond = (u32[], f32[10,10]) parameter(0)
  %count.cond = u32[] get-tuple-element((u32[], f32[10,10]) %vars.cond), index=0
  %limit = u32[] constant(10)
  ROOT %lt = pred[] compare(u32[] %count.cond, u32[] %limit), direction=LT
}

%body {
  %vars = (u32[], f32[10,10]) parameter(0)
  %count = u32[] get-tuple-element(%vars), index=0
  %acc = f32[10,10] get-tuple-element((u32[], f32[10,10]) %vars), index=1

  %one = u32[] constant(1)
  %count.1 = u32[] add(u32[] %count, u32[] %one), sharding={replicated}
  %acc.1 = f32[10,10] add(f32[10,10] %acc, f32[10,10] %acc)
  ROOT %tuple = (u32[], f32[10,10]) tuple(u32[] %count.1, f32[10,10] %acc.1)
}

ENTRY %entry {
  %p0 = f32[10,10] parameter(0)
  %p0.copy = f32[10,10] copy(f32[10,10] %p0)
  %p1 = f32[10,10] parameter(1)
  %zero = u32[] constant(0)
  %init = (u32[], f32[10,10]) tuple(u32[] %zero, f32[10,10] %p0.copy)
  %while = (u32[], f32[10,10]) while((u32[], f32[10,10]) %init),
    body=%body, condition=%cond
  %res = f32[10,10] get-tuple-element((u32[], f32[10,10]) %while), index=1
  %prev = f32[10,10] get-tuple-element((u32[], f32[10,10]) %init), index=1
  %res.1 = f32[10,10] multiply(f32[10,10] %res, %prev)
  ROOT %res_tuple = (f32[10,10]) tuple(f32[10,10] %res.1)
})";

  auto while_is_sharded = [this](HloModule* module, const HloSharding& sharding,
                                 absl::Span<const absl::Span<const OpMetadata>>
                                     sharding_metadata) {
    if (GetParam().clear_metadata) {
      ClearMetadata(module);
    }
    TF_ASSERT_OK_AND_ASSIGN(
        bool changed,
        ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
            .Run(module));
    EXPECT_TRUE(changed);
    auto while_instr = FindInstruction(module, "while");
    EXPECT_NE(nullptr, while_instr);
    std::vector<const HloInstruction*> instructions{
        while_instr, while_instr->while_body()->root_instruction(),
        while_instr->while_body()->parameter_instruction(0),
        while_instr->while_condition()->parameter_instruction(0)};

    for (auto instr : instructions) {
      ASSERT_TRUE(instr->has_sharding());
      EXPECT_EQ(sharding, instr->sharding());
      ASSERT_EQ(instr->sharding().tuple_elements().size(),
                sharding_metadata.size());
      for (int i = 0, e = sharding_metadata.size(); i < e; ++i) {
        if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
          EXPECT_THAT(instr->sharding().tuple_elements()[i],
                      ShardingMetadata(sharding_metadata[i]));
        } else {
          EXPECT_THAT(instr->sharding().tuple_elements()[i],
                      ShardingMetadata({}));
        }
      }
    }
  };
  {
    // Propagation of user-defined partial sharding of while-related instruction
    // (body root in this test).
    TF_ASSERT_OK_AND_ASSIGN(auto module,
                            ParseAndReturnVerifiedModule(hlo_string));
    auto body_root = FindInstruction(module.get(), "tuple");
    EXPECT_NE(nullptr, body_root);
    auto sharding = ParseSharding(
                        "{{replicated metadata={op_name=\"b\"}}, "
                        "{devices=[2,1]0,1 metadata={op_name=\"c\"}}}")
                        .ConsumeValueOrDie();
    body_root->set_sharding(sharding);
    while_is_sharded(module.get(), sharding.WithoutMetadata(),
                     {{CreateMetadata("b")}, {CreateMetadata("c")}});
  }
  {
    // Propagation from acc.1 to the rest of the loop.
    TF_ASSERT_OK_AND_ASSIGN(auto module,
                            ParseAndReturnVerifiedModule(hlo_string));
    auto acc_1 = FindInstruction(module.get(), "acc.1");
    EXPECT_NE(nullptr, acc_1);
    acc_1->set_sharding(
        ParseSharding("{devices=[2,1]0,1 metadata={op_name=\"b\"}}")
            .ConsumeValueOrDie());

    while_is_sharded(
        module.get(),
        ParseSharding("{{replicated}, {devices=[2,1]0,1}}").ConsumeValueOrDie(),
        {{}, {CreateMetadata("b")}});
  }
  {
    // Merge partial sharding from operand and body.
    TF_ASSERT_OK_AND_ASSIGN(auto module,
                            ParseAndReturnVerifiedModule(hlo_string));
    auto acc_1 = FindInstruction(module.get(), "acc.1");
    EXPECT_NE(nullptr, acc_1);
    acc_1->set_sharding(
        ParseSharding("{devices=[2,1,2]0,1,2,3 last_tile_dim_replicate "
                      "metadata={op_name=\"b\"}}")
            .ConsumeValueOrDie());
    auto p0 = FindInstruction(module.get(), "p0");
    p0->set_sharding(
        ParseSharding("{devices=[1,2,2]0,2,1,3 last_tile_dim_replicate "
                      "metadata={op_name=\"c\"}}")
            .ConsumeValueOrDie());

    while_is_sharded(module.get(),
                     ParseSharding("{{replicated}, "
                                   "{devices=[2,2]0,1,2,3}}")
                         .ConsumeValueOrDie(),
                     {{}, {CreateMetadata("c"), CreateMetadata("b")}});
  }
}

TEST_P(ParameterizedMetadataTest, WhileGetShardingFromRecvInBody) {
  const char* const hlo_string = R"(
HloModule module

%cond {
  %vars.cond = (u32[], f32[]) parameter(0)
  %count.cond = u32[] get-tuple-element(%vars.cond), index=0
  %limit = u32[] constant(10)
  ROOT %lt = pred[] compare(%count.cond, %limit), direction=LT
}

%body {
  %param = (u32[], f32[]) parameter(0)
  %count = u32[] get-tuple-element(%param), index=0
  %after-all = token[] after-all()
  %recv = (f32[], u32[], token[]) recv(%after-all), channel_id=1,
    sharding={maximal device=1 metadata={op_name="a"}}
  %recv-done = (f32[], token[]) recv-done(%recv), channel_id=1
  %data = f32[] get-tuple-element(%recv-done), index=0
  ROOT %tuple = (u32[], f32[]) tuple(%count, %data)
}

ENTRY %entry {
  %p0 = f32[] parameter(0)
  %zero = u32[] constant(0)
  %init = (u32[], f32[]) tuple(%zero, %p0)
  %while = (u32[], f32[]) while(%init), body=%body, condition=%cond
  ROOT %result = f32[] get-tuple-element(%while), index=1
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  // The change happens before the fixpt loop
  EXPECT_EQ(changed,
            !GetParam().propagate_metadata && !GetParam().clear_metadata);
  auto sharding = ParseSharding("{{maximal device=1}, {maximal device=1}}")
                      .ConsumeValueOrDie();
  auto while_instr = FindInstruction(module.get(), "while");
  ASSERT_NE(nullptr, while_instr);
  std::vector<const HloInstruction*> instructions{
      while_instr, while_instr->while_body()->root_instruction(),
      while_instr->while_body()->parameter_instruction(0),
      while_instr->while_condition()->parameter_instruction(0)};
  for (auto instr : instructions) {
    ASSERT_TRUE(instr->has_sharding());
    EXPECT_EQ(sharding, instr->sharding());
    for (const HloSharding& sub_sharding : instr->sharding().tuple_elements()) {
      EXPECT_THAT(sub_sharding, ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, WhileConflictingShardingInBodyBeforeRecv) {
  const char* const hlo_string = R"(
HloModule module

%cond {
  %vars.cond = (u32[], f32[]) parameter(0)
  %count.cond = u32[] get-tuple-element(%vars.cond), index=0
  %limit = u32[] constant(10)
  ROOT %lt = pred[] compare(%count.cond, %limit), direction=LT
}

%body {
  %param = (u32[], f32[]) parameter(0)
  %count = u32[] get-tuple-element(%param), index=0,
    sharding={maximal device=0 metadata={op_name="a"}}
  %after-all = token[] after-all()
  %recv = (f32[], u32[], token[]) recv(%after-all), channel_id=1,
    sharding={maximal device=1 metadata={op_name="b"}}
  %recv-done = (f32[], token[]) recv-done(%recv), channel_id=1
  %data = f32[] get-tuple-element(%recv-done), index=0
  ROOT %tuple = (u32[], f32[]) tuple(%count, %data)
}

ENTRY %entry {
  %p0 = f32[] parameter(0)
  %zero = u32[] constant(0)
  %init = (u32[], f32[]) tuple(%zero, %p0)
  %while = (u32[], f32[]) while(%init), body=%body, condition=%cond
  ROOT %result = f32[] get-tuple-element(%while), index=1
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  auto result =
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get());
  EXPECT_THAT(result.status().error_message(),
              ::testing::HasSubstr(
                  "Instruction: count is on device: 0, which conflicts with "
                  "device: 1 of channel instruction: recv"));
}

TEST_P(ParameterizedMetadataTest, WhileConflictingShardingInBodyAfterRecv) {
  const char* const hlo_string = R"(
HloModule module

%cond {
  %vars.cond = (u32[], f32[]) parameter(0)
  %count.cond = u32[] get-tuple-element(%vars.cond), index=0
  %limit = u32[] constant(10)
  ROOT %lt = pred[] compare(%count.cond, %limit), direction=LT
}

%body {
  %param = (u32[], f32[]) parameter(0)
  %count = u32[] get-tuple-element(%param), index=0
  %after-all = token[] after-all()
  %recv = (f32[], u32[], token[]) recv(%after-all), channel_id=1,
    sharding={maximal device=1 metadata={op_name="a"}}
  %recv-done = (f32[], token[]) recv-done(%recv), channel_id=1
  %data = f32[] get-tuple-element(%recv-done), index=0,
    sharding={maximal device=0 metadata={op_name="b"}}
  ROOT %tuple = (u32[], f32[]) tuple(%count, %data)
}

ENTRY %entry {
  %p0 = f32[] parameter(0)
  %zero = u32[] constant(0)
  %init = (u32[], f32[]) tuple(%zero, %p0)
  %while = (u32[], f32[]) while(%init), body=%body, condition=%cond
  ROOT %result = f32[] get-tuple-element(%while), index=1
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  auto result =
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get());
  EXPECT_THAT(result.status().error_message(),
              ::testing::HasSubstr(
                  "Instruction: data is on device: 0, which conflicts with "
                  "device: 1 of channel instruction: recv"));
}

TEST_P(ParameterizedMetadataTest, WhileConflictingShardingOnWhileInstruction) {
  const char* const hlo_string = R"(
HloModule module

%cond {
  %vars.cond = (u32[], f32[]) parameter(0)
  %count.cond = u32[] get-tuple-element(%vars.cond), index=0
  %limit = u32[] constant(10)
  ROOT %lt = pred[] compare(%count.cond, %limit), direction=LT
}

%body {
  %param = (u32[], f32[]) parameter(0)
  %count = u32[] get-tuple-element(%param), index=0
  %after-all = token[] after-all()
  %recv = (f32[], u32[], token[]) recv(%after-all), channel_id=1,
    sharding={maximal device=1 metadata={op_name="a"}}
  %recv-done = (f32[], token[]) recv-done(%recv), channel_id=1
  %data = f32[] get-tuple-element(%recv-done), index=0
  ROOT %tuple = (u32[], f32[]) tuple(%count, %data)
}

ENTRY %entry {
  %p0 = f32[] parameter(0)
  %zero = u32[] constant(0)
  %init = (u32[], f32[]) tuple(%zero, %p0)
  %while = (u32[], f32[]) while(%init), body=%body, condition=%cond,
    sharding={maximal device=0 metadata={op_name="b"}}
  ROOT %result = f32[] get-tuple-element(%while), index=1
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  auto result =
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get());
  EXPECT_THAT(result.status().error_message(),
              ::testing::HasSubstr(
                  "Instruction: while is on device: 0, which conflicts with "
                  "device: 1 of channel instruction: recv"));
}

TEST_P(ParameterizedMetadataTest, WhileConv) {
  const char* const hlo_string = R"(
HloModule module

%cond {
  %vars.cond = (u32[], bf16[2048, 768], bf16[128,512,2048], bf16[128,512,768]) parameter(0)
  %count.cond = u32[] get-tuple-element(%vars.cond), index=0
  %limit = u32[] constant(2)
  ROOT %lt = pred[] compare(%count.cond, %limit), direction=LT
}

%body {
  %param = (u32[], bf16[2048, 768], bf16[128,512,2048], bf16[128,512,768]) parameter(0)
  %count = u32[] get-tuple-element(%param), index=0
  %kernel = bf16[2048, 768]{1,0}
   get-tuple-element(%param), index=1
  %lhs = bf16[128,512,2048]{2,1,0}
   get-tuple-element(%param), index=2,
   sharding={devices=[8,1,2]0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}
  %reshape = bf16[2048,768,1]{2,1,0} reshape(bf16[2048,768]{1,0} %kernel)
  %convolution = bf16[128,512,768]{2,1,0}
    convolution(bf16[128,512,2048]{2,1,0} %lhs,
    bf16[2048,768,1]{2,1,0} %reshape), window={size=1},
    dim_labels=0bf_io0->0bf, sharding={devices=[8,1,2]0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}
  ROOT %tuple = (u32[], bf16[2048, 768], bf16[128,512,2048], bf16[128,512,768]) tuple(%count, %kernel, %lhs, %convolution)
}

ENTRY %entry {
  %p0 = bf16[2048,768] parameter(0),
    sharding={devices=[2,1,8]0,2,4,6,8,10,12,14,1,3,5,7,9,11,13,15 last_tile_dim_replicate}
  %p1 = bf16[128,512,2048] parameter(1),
   sharding={devices=[8,1,2]0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}
  %p2 = bf16[128,512,768] parameter(2)
  %zero = u32[] constant(0)
  %init = (u32[], bf16[2048, 768], bf16[128,512,2048], bf16[128,512,768]) tuple(%zero, %p0, %p1, %p2)
  %while = (u32[], bf16[2048, 768], bf16[128,512,2048], bf16[128,512,768]) while(%init), body=%body, condition=%cond
  ROOT %result = bf16[128,512,768] get-tuple-element(%while), index=3, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* kernel = FindInstruction(module.get(), "kernel");
  ASSERT_NE(kernel, nullptr);
  EXPECT_THAT(kernel, op::Sharding("{devices=[2,1,8]0,2,4,6,8,10,12,14,1,3,5,"
                                   "7,9,11,13,15 last_tile_dim_replicate}"));
}

TEST_P(ParameterizedMetadataTest, Dot) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %param.0 = f32[8,256,128] parameter(0)
  %param.1 = f32[8,128,512] parameter(1)
  %param.2 = f32[8,128] parameter(2)

  %p0_copy_0 = f32[8,256,128] copy(%param.0),
    sharding={devices=[1,4,1]0,1,2,3 metadata={op_name="a"}}
  %p1_copy_0 = f32[8,128,512] copy(%param.1),
    sharding={devices=[1,1,4]0,1,2,3 metadata={op_name="b"}}
  %p2_copy = f32[8,128] copy(%param.2)
  %dot_prop_rhs = f32[8,256,512] dot(%p0_copy_0, %p1_copy_0),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={1}
  %dot_prop_lhs = f32[8,512,256] dot(%p1_copy_0, %p0_copy_0),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={1}, rhs_contracting_dims={2}
  %dot_mat_vec = f32[8,256] dot(%p0_copy_0, %p2_copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={1}

  %p0_copy_1 = f32[8,256,128] copy(%param.0)
  %p1_copy_1 = f32[8,128,512] copy(%param.1)
  %dot_back_prop_rhs = f32[8,256,512] dot(%p0_copy_1, %p1_copy_1),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={1}
  %copy_back_prop_rhs = f32[8,256,512] copy(%dot_back_prop_rhs),
    sharding={devices=[1,2,2]0,1,2,3 metadata={op_name="c"}}

  ROOT %tuple = (f32[8,256,256], f32[8,256,256], f32[8,256])
    tuple(%dot_prop_lhs, %dot_prop_rhs, %dot_mat_vec, %copy_back_prop_rhs)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* dot_prop_rhs = FindInstruction(module.get(), "dot_prop_rhs");
  ASSERT_NE(dot_prop_rhs, nullptr);
  EXPECT_THAT(dot_prop_rhs, op::Sharding("{devices=[1,1,4]0,1,2,3}"));
  auto* dot_prop_lhs = FindInstruction(module.get(), "dot_prop_lhs");
  ASSERT_NE(dot_prop_lhs, nullptr);
  EXPECT_THAT(dot_prop_lhs, op::Sharding("{devices=[1,4,1]0,1,2,3}"));
  auto* dot_mat_vec = FindInstruction(module.get(), "dot_mat_vec");
  ASSERT_NE(dot_mat_vec, nullptr);
  EXPECT_THAT(dot_mat_vec, op::Sharding("{devices=[1,4]0,1,2,3}"));

  auto* p0_copy_1 = FindInstruction(module.get(), "p0_copy_1");
  ASSERT_NE(p0_copy_1, nullptr);
  EXPECT_THAT(
      p0_copy_1,
      op::Sharding("{devices=[1,2,1,2]0,1,2,3 last_tile_dim_replicate}"));
  auto* p1_copy_1 = FindInstruction(module.get(), "p1_copy_1");
  ASSERT_NE(p1_copy_1, nullptr);
  EXPECT_THAT(
      p1_copy_1,
      op::Sharding("{devices=[1,1,2,2]0,2,1,3  last_tile_dim_replicate}"));
  auto* dot_back_prop_rhs = FindInstruction(module.get(), "dot_back_prop_rhs");
  ASSERT_NE(dot_back_prop_rhs, nullptr);
  EXPECT_THAT(dot_back_prop_rhs, op::Sharding("{devices=[1,2,2]0,1,2,3}"));

  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(dot_prop_rhs->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
    EXPECT_THAT(dot_prop_lhs->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
    EXPECT_THAT(dot_mat_vec->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
    EXPECT_THAT(p0_copy_1->sharding(), ShardingMetadata({CreateMetadata("c")}));
    EXPECT_THAT(p1_copy_1->sharding(), ShardingMetadata({CreateMetadata("c")}));
    EXPECT_THAT(dot_back_prop_rhs->sharding(),
                ShardingMetadata({CreateMetadata("c")}));
  } else {
    for (HloInstruction* instruction :
         {dot_prop_rhs, dot_prop_lhs, dot_mat_vec, p0_copy_1, p1_copy_1,
          dot_back_prop_rhs}) {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, DotTiledBatchDim) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %p0 = f32[8,256,512] parameter(0)
  %p1 = f32[8,512,128] parameter(1)

  %add = f32[8,256,512] add(%p0, %p0)
  %dot = f32[8,256,128] dot(%add, %p1),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={1}
  %res = f32[8,32768] reshape(%dot),
    sharding={devices=[2,2]0,1,2,3 metadata={op_name="a"}}

  ROOT %tuple = (f32[8,32768]) tuple(%res)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "add");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,2,1]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, DotMergeOperands) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %p0 = f32[8,256,512] parameter(0),
    sharding={devices=[2,2,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate metadata={op_name="a"}}
  %p1 = f32[8,128,512] parameter(1),
    sharding={devices=[2,2,1,2]0,2,1,3,4,6,5,7 last_tile_dim_replicate metadata={op_name="b"}}
  %dot = f32[8,256,128] dot(%p0, %p1),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2}
  ROOT %copy = f32[8,256,128] copy(%dot)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "dot");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,2,2]0,1,2,3,4,5,6,7}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b"), CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, DotMergeOperands2) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %p0 = f32[8,256,512] parameter(0),
    sharding={devices=[2,2,2]0,1,2,3,4,5,6,7 metadata={op_name="a"}}
  %p1 = f32[8,128,512] parameter(1),
    sharding={devices=[2,2,2]0,1,2,3,4,5,6,7 metadata={op_name="b"}}
  %dot = f32[8,256,128] dot(%p0, %p1),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2}
  ROOT %copy = f32[8,256,128] copy(%dot)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "dot");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(
      instruction,
      op::Sharding(
          "{devices=[2,2,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, DotMergeOperands3) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %p0 = f32[256,512] parameter(0),
    sharding={devices=[2,4]0,1,2,3,4,5,6,7 metadata={op_name="a"}}
  %p1 = f32[128,512] parameter(1),
    sharding={devices=[4,2]0,4,2,6,3,7,1,5 metadata={op_name="b"}}
  %dot = f32[256,128] dot(%p0, %p1),
    lhs_contracting_dims={1}, rhs_contracting_dims={1}
  ROOT %copy = f32[256,128] copy(%dot)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "dot");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,4]0,2,3,1,4,6,7,5}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b"), CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, BackwardDotFromContracting) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %p0 = f32[8,256,512] parameter(0),
    sharding={devices=[2,2,2]0,1,2,3,4,5,6,7 metadata={op_name="a"}}
  %p1 = f32[8,128,512] parameter(1)
  %copy1 = f32[8,128,512] copy(%p1)
  %dot = f32[8,256,128] dot(%p0, %copy1),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2},
    sharding={devices=[2,1,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate metadata={op_name="b"}}
  ROOT %copy = f32[8,256,128] copy(%dot)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "copy1");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,2,2]0,1,2,3,4,5,6,7}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a"), CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ConvAsDotOnTrivialDims) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %lhs = f32[128,1,1,1001] parameter(0),
    sharding={devices=[1,2,1,1]0,1 metadata={op_name="a"}}
  %rhs = f32[1,1,1024,1001] parameter(1),
    sharding={devices=[1,2,1,1]0,1 metadata={op_name="b"}}
  %convolution = f32[128,1,1,1024] convolution(%lhs, %rhs),
    window={size=1x1 rhs_reversal=1x1}, dim_labels=b01f_01oi->b01f
  ROOT %copy = f32[128,1,1,1024] copy(%convolution)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "convolution");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,1,2,1]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ConvAsDotOnTrivialDimsBackward) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %p0 = f32[128,5,5,128] parameter(0)
  %lhs = f32[128,5,5,128] copy(%p0)
  %p1 = f32[5,5,128,768] parameter(1)
  %rhs = f32[5,5,128,768] copy(%p1)
  %convolution = f32[128,1,1,768] convolution(%lhs, %rhs), window={size=5x5},
    dim_labels=b01f_01io->b01f,
    sharding={devices=[1,2,1,1]0,1 metadata={op_name="a"}}
  ROOT %copy = f32[128,1,1,768] copy(%convolution)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* lhs = FindInstruction(module.get(), "lhs");
  ASSERT_NE(lhs, nullptr);
  auto* rhs = FindInstruction(module.get(), "rhs");
  ASSERT_NE(rhs, nullptr);
  for (HloInstruction* instruction : {lhs, rhs}) {
    EXPECT_THAT(instruction, op::Sharding("{replicated}"));
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest,
       ConvolutionFilterIFOFPartitionedInputPartialReplicate) {
  const char* const hlo_string = R"(
  HloModule module

ENTRY entry {
  %lhs = f32[128,112,112,12] parameter(0)
  %lhs.copy = f32[128,112,112,12] copy(f32[128,112,112,12] %lhs),
    sharding={devices=[1,1,1,2,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="a"}}
  %rhs = f32[7,7,12,64] parameter(1)
  %rhs.copy = f32[7,7,12,64] copy(f32[7,7,12,64] %rhs),
    sharding={devices=[1,1,2,2]0,1,2,3 metadata={op_name="b"}}
  %conv = f32[128,56,56,64] convolution(
    f32[128,112,112,12] %lhs.copy,
    f32[7,7,12,64] %rhs.copy),
    window={size=7x7 stride=2x2 pad=3_3x3_3},
    dim_labels=b01f_01io->b01f
  ROOT %copy = f32[128,56,56,64] copy(conv)
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "conv");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(
      instruction,
      op::Sharding("{devices=[1,1,1,2,2]0,2,1,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ConcatFromUserUnshardedDim) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %p0 = f32[8,128] parameter(0)
  %p1 = f32[8,128] parameter(1)
  %c0 = f32[8,128] copy(%p0)
  %c1 = f32[8,128] copy(%p1)

  %concat = f32[16,128] concatenate(%c0, %c1),
    dimensions={0},
    sharding={devices=[1,2]0,1 metadata={op_name="a"}}
  ROOT %tuple = (f32[16,128]) tuple(%concat)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* c0 = FindInstruction(module.get(), "c0");
  ASSERT_NE(c0, nullptr);
  auto* c1 = FindInstruction(module.get(), "c1");
  ASSERT_NE(c1, nullptr);
  for (HloInstruction* instruction : {c0, c1}) {
    EXPECT_THAT(instruction, op::Sharding("{devices=[1,2]0,1}"));
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, ConcatFromUserShardedDim) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %p0 = f32[8,128] parameter(0)
  %p1 = f32[8,128] parameter(1)
  %c0 = f32[8,128] copy(%p0)
  %c1 = f32[8,128] copy(%p1)

  %concat = f32[16,128] concatenate(%c0, %c1),
    dimensions={0},
    sharding={devices=[3,1]0,1,2 metadata={op_name="a"}}
  ROOT %tuple = (f32[16,128]) tuple(%concat)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* c0 = FindInstruction(module.get(), "c0");
  EXPECT_THAT(c0, op::Sharding("{devices=[2,1]0,1}"));
  ASSERT_NE(c0, nullptr);
  auto* c1 = FindInstruction(module.get(), "c1");
  ASSERT_NE(c1, nullptr);
  EXPECT_THAT(c1, op::Sharding("{devices=[2,1]1,2}"));
  for (HloInstruction* instruction : {c0, c1}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, ConcatFromUserShardedDimMaximalOperand) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %conv {
  %p0 = f32[8,128] parameter(0)
  %p1 = f32[24,128] parameter(1)
  %c0 = f32[8,128] copy(%p0)
  %c1 = f32[24,128] copy(%p1)

  %concat = f32[32,128] concatenate(%c0, %c1),
    dimensions={0},
    sharding={devices=[4,1]0,1,2,3 metadata={op_name="a"}}
  ROOT %tuple = (f32[32,128]) tuple(%concat)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* c0 = FindInstruction(module.get(), "c0");
  ASSERT_NE(c0, nullptr);
  EXPECT_THAT(c0, op::NoSharding());
  auto* c1 = FindInstruction(module.get(), "c1");
  ASSERT_NE(c1, nullptr);
  EXPECT_THAT(c1, op::Sharding("{devices=[3,1]1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(c1->sharding(), ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(c1->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ReplicatedToSideEffecting) {
  const char* const hlo_string = R"(
HloModule module
ENTRY entry_computation {
  %const.0 = s32[] constant(0),
    sharding={replicated metadata={op_name="a"}}
  %const.1 = s32[] constant(2147483647),
    sharding={replicated metadata={op_name="b"}}
  %rng = s32[4]{0} rng(%const.0, %const.1),
    distribution=rng_uniform
  ROOT %root = (s32[4]{0}) tuple(%rng)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_EQ(changed,
            !GetParam().propagate_metadata && !GetParam().clear_metadata);
  auto* instruction = FindInstruction(module.get(), "rng");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::NoSharding());
}

TEST_P(ParameterizedMetadataTest, PartReplicatedTupleUser) {
  const char* const hlo_string = R"(
HloModule module
ENTRY entry_computation {
  %param.0 = f32[5] parameter(0)
  %param.1 = f32[7] parameter(1)
  %param.2 = f32[9] parameter(2)
  %tuple.0 = (f32[5], f32[7]) tuple(%param.0, %param.1)
  ROOT %tuple.1 = ((f32[5], f32[7]), f32[9]) tuple(%tuple.0, %param.2),
    sharding={{maximal device=0 metadata={op_name="a"}},
              {replicated metadata={op_name="b"}},
              {maximal device=1 metadata={op_name="c"}}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "tuple.0");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{{maximal device=0}, {replicated}}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding().tuple_elements()[0],
                ShardingMetadata({CreateMetadata("a")}));
    EXPECT_THAT(instruction->sharding().tuple_elements()[1],
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    for (const HloSharding& sub_sharding :
         instruction->sharding().tuple_elements()) {
      EXPECT_THAT(sub_sharding, ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, Conditional) {
  const char* const hlo_string = R"(
HloModule module

%true_comp {
  %tp = (f32[3,5]) parameter(0)
  %tgte = f32[3,5] get-tuple-element(%tp), index=0
  %ttr = f32[5,3] transpose(%tgte), dimensions={1,0}
  ROOT %tr = (f32[5,3]) tuple(%ttr)
}

%false_comp {
  %fp = (f32[5,3]) parameter(0)
  %fgte = f32[5,3] get-tuple-element(%fp), index=0
  ROOT %fr = (f32[5,3]) tuple(%fgte)
}

ENTRY entry {
  %cond = pred[] parameter(0)
  %true_param = (f32[3,5]) parameter(1),
    sharding={{devices=[1,2]0,1 metadata={op_name="a"}}}
  %false_param = (f32[5,3]) parameter(2),
    sharding={{devices=[1,3]0,1,2 metadata={op_name="b"}}}
  %conditional = (f32[5,3]) conditional(
      %cond, %true_param, %false_param),
    true_computation=%true_comp,
    false_computation=%false_comp
  ROOT %root = f32[5,3] get-tuple-element(%conditional), index=0
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* tp = FindInstruction(module.get(), "tp");
  ASSERT_NE(tp, nullptr);
  EXPECT_THAT(tp, op::Sharding("{{devices=[1,2]0,1}}"));
  auto* tgte = FindInstruction(module.get(), "tgte");
  ASSERT_NE(tgte, nullptr);
  EXPECT_THAT(tgte, op::Sharding("{devices=[1,2]0,1}"));
  auto* ttr = FindInstruction(module.get(), "ttr");
  ASSERT_NE(ttr, nullptr);
  EXPECT_THAT(ttr, op::Sharding("{devices=[2,1]0,1}"));
  auto* tr = FindInstruction(module.get(), "tr");
  ASSERT_NE(tr, nullptr);
  EXPECT_THAT(tr, op::Sharding("{{devices=[1,3]0,1,2}}"));
  auto* fp = FindInstruction(module.get(), "fp");
  ASSERT_NE(fp, nullptr);
  EXPECT_THAT(fp, op::Sharding("{{devices=[1,3]0,1,2}}"));
  auto* fgte = FindInstruction(module.get(), "fgte");
  ASSERT_NE(fgte, nullptr);
  EXPECT_THAT(fgte, op::Sharding("{devices=[1,3]0,1,2}"));
  auto* fr = FindInstruction(module.get(), "fr");
  ASSERT_NE(fr, nullptr);
  EXPECT_THAT(fr, op::Sharding("{{devices=[1,3]0,1,2}}"));
  auto* conditional = FindInstruction(module.get(), "conditional");
  ASSERT_NE(conditional, nullptr);
  EXPECT_THAT(conditional, op::Sharding("{{devices=[1,3]0,1,2}}"));

  auto check_metadata = [](const HloSharding& sharding,
                           const OpMetadata& metadata) {
    if (sharding.IsTuple()) {
      EXPECT_THAT(sharding.tuple_elements()[0], ShardingMetadata({metadata}));
    } else {
      EXPECT_THAT(sharding, ShardingMetadata({metadata}));
    }
  };

  auto check_empty_metadata = [](const HloSharding& sharding) {
    if (sharding.IsTuple()) {
      EXPECT_THAT(sharding.tuple_elements()[0], ShardingMetadata({}));
    } else {
      EXPECT_THAT(sharding, ShardingMetadata({}));
    }
  };

  for (HloInstruction* instruction : {tp, tgte, ttr}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      check_metadata(instruction->sharding(), CreateMetadata("a"));
    } else {
      check_empty_metadata(instruction->sharding());
    }
  }
  for (HloInstruction* instruction : {tr, fp, fgte, fr, conditional}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      check_metadata(instruction->sharding(), CreateMetadata("b"));
    } else {
      check_empty_metadata(instruction->sharding());
    }
  }
}

TEST_P(ParameterizedMetadataTest, TupleFromUser) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %entry {
  %p0 = f32[13] parameter(0)
  %p1 = f32[15] parameter(1)
  %p2 = f32[17] parameter(2)
  %t0 = (f32[13], f32[15]) tuple(%p0, %p1)
  %t1 = ((f32[13], f32[15]), f32[17]) tuple(%t0, %p2)
  %gte.0 = (f32[13], f32[15]) get-tuple-element(%t1), index=0
  %gte.1 = f32[13] get-tuple-element(%gte.0), index=0
  %gte.2 = f32[15] get-tuple-element(%gte.0), index=1
  %gte.3 = f32[17] get-tuple-element(%t1), index=1
  ROOT %t2 = (f32[13], f32[15], f32[17]) tuple(%gte.1, %gte.2, %gte.3),
    sharding={{replicated metadata={op_name="a"}},
              {devices=[2]0,1 metadata={op_name="b"}},
              {devices=[3]1,2,3 metadata={op_name="c"}}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* t0 = FindInstruction(module.get(), "t0");
  ASSERT_NE(t0, nullptr);
  EXPECT_THAT(t0, op::Sharding("{{replicated}, {devices=[2]0,1}}"));
  auto* t1 = FindInstruction(module.get(), "t1");
  ASSERT_NE(t1, nullptr);
  EXPECT_THAT(
      t1, op::Sharding("{{replicated}, {devices=[2]0,1}, {devices=[3]1,2,3}}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(t0->sharding().tuple_elements()[0],
                ShardingMetadata({CreateMetadata("a")}));
    EXPECT_THAT(t0->sharding().tuple_elements()[1],
                ShardingMetadata({CreateMetadata("b")}));
    EXPECT_THAT(t1->sharding().tuple_elements()[0],
                ShardingMetadata({CreateMetadata("a")}));
    EXPECT_THAT(t1->sharding().tuple_elements()[1],
                ShardingMetadata({CreateMetadata("b")}));
    EXPECT_THAT(t1->sharding().tuple_elements()[2],
                ShardingMetadata({CreateMetadata("c")}));
  } else {
    for (HloInstruction* instruction : {t0, t1}) {
      for (const HloSharding& sub_sharding :
           instruction->sharding().tuple_elements()) {
        EXPECT_THAT(sub_sharding, ShardingMetadata({}));
      }
    }
  }
}

TEST_P(ParameterizedMetadataTest, DynamicSliceForwardPass) {
  const char* hlo_string = R"(
HloModule module
ENTRY %entry {
  %p0 = f32[11,13,15] parameter(0)
  %c0 = f32[11,13,15] copy(%p0),
    sharding={devices=[1,1,2]0,1 metadata={op_name="a"}}
  %p1 = s32[] parameter(1)
  %i0 = s32[] constant(0)
  %ds = f32[11,1,15] dynamic-slice(%c0, %i0, %p1, %i0),
    dynamic_slice_sizes={11,1,15}
  ROOT %root = (f32[11,1,15]) tuple(%ds)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "ds");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,1,2]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, DynamicSliceBackwardPass) {
  const char* hlo_string = R"(
HloModule module
ENTRY %entry {
  %p0 = f32[11,13,15] parameter(0)
  %c0 = f32[11,13,15] copy(%p0)
  %p1 = s32[] parameter(1)
  %i0 = s32[] constant(0)
  %ds = f32[11,1,15] dynamic-slice(%c0, %i0, %p1, %i0),
    dynamic_slice_sizes={11,1,15},
    sharding={devices=[1,1,2]0,1 metadata={op_name="a"}}
  ROOT %root = (f32[11,1,15]) tuple(%ds)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "c0");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,1,2]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, DynamicUpdateSliceForwardPassBase) {
  const char* hlo_string = R"(
HloModule module
ENTRY %entry {
  %p0 = f32[11,13,15] parameter(0)
  %c0 = f32[11,13,15] copy(%p0),
    sharding={devices=[1,1,2]0,1 metadata={op_name="a"}}
  %p1 = f32[11,1,15] parameter(1)
  %c1 = f32[11,1,15] copy(%p1)
  %p2 = s32[] parameter(2)
  %i0 = s32[] constant(0)
  %dus = f32[11,13,15] dynamic-update-slice(%c0, %c1, %i0, %p2, %i0)
  ROOT %root = (f32[11,13,15]) tuple(%dus)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* dus = FindInstruction(module.get(), "dus");
  ASSERT_NE(dus, nullptr);
  EXPECT_THAT(dus, op::Sharding("{devices=[1,1,2]0,1}"));
  auto* c1 = FindInstruction(module.get(), "c1");
  ASSERT_NE(c1, nullptr);
  EXPECT_THAT(c1, op::Sharding("{devices=[1,1,2]0,1}"));
  for (HloInstruction* instruction : {dus, c1}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, DynamicUpdateSliceForwardPassUpdate) {
  const char* hlo_string = R"(
HloModule module
ENTRY %entry {
  %p0 = f32[11,13,15] parameter(0)
  %c0 = f32[11,13,15] copy(%p0)
  %p1 = f32[11,1,15] parameter(1)
  %c1 = f32[11,1,15] copy(%p1),
    sharding={devices=[1,1,2]0,1 metadata={op_name="a"}}
  %p2 = s32[] parameter(2)
  %i0 = s32[] constant(0)
  %dus = f32[11,13,15] dynamic-update-slice(%c0, %c1, %i0, %p2, %i0)
  ROOT %root = (f32[11,13,15]) tuple(%dus)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* dus = FindInstruction(module.get(), "dus");
  ASSERT_NE(dus, nullptr);
  EXPECT_THAT(dus, op::Sharding("{devices=[1,1,2]0,1}"));
  auto* c0 = FindInstruction(module.get(), "c0");
  ASSERT_NE(c0, nullptr);
  EXPECT_THAT(c0, op::Sharding("{devices=[1,1,2]0,1}"));
  for (HloInstruction* instruction : {dus, c0}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, DynamicUpdateSliceBackwardPass) {
  const char* hlo_string = R"(
HloModule module
ENTRY %entry {
  %p0 = f32[11,13,15] parameter(0)
  %c0 = f32[11,13,15] copy(%p0)
  %p1 = f32[11,1,15] parameter(1)
  %c1 = f32[11,1,15] copy(%p1)
  %p2 = s32[] parameter(2)
  %i0 = s32[] constant(0)
  %dus = f32[11,13,15] dynamic-update-slice(%c0, %c1, %i0, %p2, %i0),
    sharding={devices=[1,1,2]0,1 metadata={op_name="a"}}
  ROOT %root = (f32[11,13,15]) tuple(%dus)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* c0 = FindInstruction(module.get(), "c0");
  ASSERT_NE(c0, nullptr);
  EXPECT_THAT(c0, op::Sharding("{devices=[1,1,2]0,1}"));
  auto* c1 = FindInstruction(module.get(), "c1");
  ASSERT_NE(c1, nullptr);
  EXPECT_THAT(c1, op::Sharding("{devices=[1,1,2]0,1}"));
  for (HloInstruction* instruction : {c0, c1}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, EinsumLHSBatchPartitioned) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64] parameter(0)
  %lhs.copy = f32[32,24,64] copy(%lhs),
    sharding={devices=[2,1,1]0,1 metadata={op_name="a"}}
  %rhs = f32[32,39296,64] parameter(1)
  %rhs.copy = f32[32,39296,64] copy(%rhs)
  %conv = f32[32,24,39296] convolution(%lhs.copy, %rhs.copy),
    dim_labels=0bf_0oi->0bf, window={size=32 stride=31 lhs_dilate=32}
  ROOT %copy = f32[32,24,39296] copy(%conv)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* rhs_copy = FindInstruction(module.get(), "rhs.copy");
  ASSERT_NE(rhs_copy, nullptr);
  EXPECT_THAT(rhs_copy, op::Sharding("{devices=[2,1,1]0,1}"));
  auto* conv = FindInstruction(module.get(), "conv");
  ASSERT_NE(conv, nullptr);
  EXPECT_THAT(conv, op::Sharding("{devices=[2,1,1]0,1}"));
  for (HloInstruction* instruction : {rhs_copy, conv}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, EinsumOutputBatchPartitioned) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64] parameter(0)
  %lhs.copy = f32[32,24,64] copy(%lhs)
  %rhs = f32[32,39296,64] parameter(1)
  %rhs.copy = f32[32,39296,64] copy(%rhs)
  %conv = f32[32,24,39296] convolution(%lhs.copy, %rhs.copy),
    dim_labels=0bf_0oi->0bf, window={size=32 stride=31 lhs_dilate=32},
    sharding={devices=[2,1,1]0,1 metadata={op_name="a"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* lhs_copy = FindInstruction(module.get(), "lhs.copy");
  ASSERT_NE(lhs_copy, nullptr);
  EXPECT_THAT(lhs_copy, op::Sharding("{devices=[2,1,1]0,1}"));
  auto* rhs_copy = FindInstruction(module.get(), "rhs.copy");
  ASSERT_NE(rhs_copy, nullptr);
  EXPECT_THAT(rhs_copy, op::Sharding("{devices=[2,1,1]0,1}"));
  for (HloInstruction* instruction : {lhs_copy, rhs_copy}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, EinsumLHSNonContractingPartitioned) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64,128] parameter(0)
  %lhs.copy = f32[32,24,64,128] copy(%lhs),
    sharding={devices=[1,2,1,2]0,1,2,3 metadata={op_name="a"}}
  %rhs = f32[32,39296,64,1] parameter(1)
  %rhs.copy = f32[32,39296,64,1] copy(%rhs)
  %conv = f32[32,24,39296,128] convolution(%lhs.copy, %rhs.copy),
    dim_labels=0bf1_0oi1->0bf1, window={size=32x1 stride=31x1 lhs_dilate=32x1}
  ROOT %copy = f32[32,24,39296,128] copy(%conv)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "conv");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,1,2]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, EinsumOutputLHSNonContractingPartitioned) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64,128] parameter(0)
  %lhs.copy = f32[32,24,64,128] copy(%lhs)
  %rhs = f32[32,39296,64,1] parameter(1)
  %rhs.copy = f32[32,39296,64,1] copy(%rhs)
  ROOT %conv = f32[32,24,39296,128] convolution(%lhs.copy, %rhs.copy),
    dim_labels=0bf1_0oi1->0bf1, window={size=32x1 stride=31x1 lhs_dilate=32x1},
    sharding={devices=[1,2,1,2]0,1,2,3 metadata={op_name="a"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "lhs.copy");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,1,2]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, EinsumRHSNonContractingPartitioned) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64,1] parameter(0)
  %lhs.copy = f32[32,24,64,1] copy(%lhs)
  %rhs = f32[32,39296,64,128] parameter(1)
  %rhs.copy = f32[32,39296,64,128] copy(%rhs),
    sharding={devices=[1,2,1,2]0,1,2,3 metadata={op_name="a"}}
  %conv = f32[32,24,39296,128] convolution(%lhs.copy, %rhs.copy),
    dim_labels=0bf1_0oi1->0bf1,
    window={size=32x128 stride=31x1 pad=0_0x127_127 lhs_dilate=32x1 rhs_reversal=0x1}
  ROOT %copy = f32[32,24,39296,128] copy(%conv)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "conv");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,1,2,2]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, EinsumOutputRHSNonContractingPartitioned) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64,1] parameter(0)
  %lhs.copy = f32[32,24,64,1] copy(%lhs)
  %rhs = f32[32,39296,64,128] parameter(1)
  %rhs.copy = f32[32,39296,64,128] copy(%rhs)
  ROOT %conv = f32[32,24,39296,128] convolution(%lhs.copy, %rhs.copy),
    dim_labels=0bf1_0oi1->0bf1,
    window={size=32x128 stride=31x1 pad=0_0x127_127 lhs_dilate=32x1 rhs_reversal=0x1},
    sharding={devices=[1,1,2,2]0,1,2,3 metadata={op_name="a"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "rhs.copy");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,1,2]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, EinsumChooseLargerOperand) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64,1] parameter(0)
  %lhs.copy = f32[32,24,64,1] copy(%lhs),
    sharding={devices=[1,4,1,1]0,1,2,3 metadata={op_name="a"}}
  %rhs = f32[32,39296,64,128] parameter(1)
  %rhs.copy = f32[32,39296,64,128] copy(%rhs),
    sharding={devices=[1,2,1,2]0,1,2,3 metadata={op_name="b"}}
  %conv = f32[32,24,39296,128] convolution(%lhs.copy, %rhs.copy),
    dim_labels=0bf1_0oi1->0bf1,
    window={size=32x128 stride=31x1 pad=0_0x127_127 lhs_dilate=32x1 rhs_reversal=0x1}
  ROOT %copy = f32[32,24,39296,128] copy(%conv)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "conv");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,1,2,2]0,1,2,3}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, EinsumChooseBatchFirst) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64,1] parameter(0)
  %lhs.copy = f32[32,24,64,1] copy(%lhs),
    sharding={devices=[1,2,1,1]0,1 metadata={op_name="a"}}
  %rhs = f32[32,39296,64,128] parameter(1)
  %rhs.copy = f32[32,39296,64,128] copy(%rhs),
    sharding={devices=[2,1,1,1]0,1 metadata={op_name="b"}}
  %conv = f32[32,24,39296,128] convolution(%lhs.copy, %rhs.copy),
    dim_labels=0bf1_0oi1->0bf1,
    window={size=32x128 stride=31x1 pad=0_0x127_127 lhs_dilate=32x1 rhs_reversal=0x1}
  ROOT %copy = f32[32,24,39296,128] copy(%conv)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "conv");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,1,1,1]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, GatherFromIndex) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %input = f32[2,2,9] parameter(0),
    sharding={replicated  metadata={op_name="a"}}
  %indices = s32[2,3,4] parameter(1),
    sharding={devices=[1,2,1]0,1 metadata={op_name="b"}}
  %gather = f32[3,4,9] gather(%input, %indices), offset_dims={2},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,9}
  ROOT %copy = f32[3,4,9] copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "gather");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,1,1]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, GatherFromIndex_PartialReplicate) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %indices = s32[3] parameter(1),
   sharding={devices=[2,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="b"}}
  %gather = f32[3,9] gather(%input, %indices), offset_dims={1},
    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1,
    slice_sizes={1,9}
  ROOT %copy = f32[3,9] copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "gather");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, GatherFromDataOperand) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={devices=[1,2]0,1 metadata={op_name="a"}}
  %indices = s32[3] parameter(1),
    sharding={replicated metadata={op_name="b"}}
  %gather = f32[3,9] gather(%input, %indices), offset_dims={1},
    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1,
    slice_sizes={1,9}
  ROOT %copy = f32[3,9] copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "gather");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, GatherFromDataOperand_PartialReplicate) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="a"}}
  %indices = s32[3] parameter(1),
    sharding={replicated metadata={op_name="b"}}
  %gather = f32[3,9] gather(%input, %indices), offset_dims={1},
    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1,
    slice_sizes={1,9}
  ROOT %copy = f32[3,9] copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "gather");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, GatherToIndex) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %p1 = s32[3] parameter(1)
  %indices = s32[3] copy(%p1)
  ROOT %gather = f32[3,9] gather(%input, %indices), offset_dims={1},
    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1,
    slice_sizes={1,9},
    sharding={devices=[2,1]0,1 metadata={op_name="b"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "indices");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, GatherToIndex_PartialReplicate) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %p1 = s32[3] parameter(1)
  %indices = s32[3] copy(%p1)
  ROOT %gather = f32[3,9] gather(%input, %indices), offset_dims={1},
    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1,
    slice_sizes={1,9},
    sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="b"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "indices");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[2,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, GatherToIndex2) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %input = bf16[2,4819,4] parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %p1 = s32[2,1000,2] parameter(1)
  %indices = s32[2,1000,2] copy(%p1)
  ROOT %gather = bf16[2,1000,4]
    gather(bf16[2,4819,4] %input, s32[2,1000,2] %indices),
    offset_dims={2}, collapsed_slice_dims={0,1},
    start_index_map={0,1}, index_vector_dim=2, slice_sizes={1,1,4},
    sharding={devices=[1,2,1]0,1 metadata={op_name="b"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "indices");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2,1]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, GatherToIndex2_PartialReplicate) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %input = bf16[2,4819,4] parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %p1 = s32[2,1000,2] parameter(1)
  %indices = s32[2,1000,2] copy(%p1)
  ROOT %gather = bf16[2,1000,4]
    gather(bf16[2,4819,4] %input, s32[2,1000,2] %indices),
    offset_dims={2}, collapsed_slice_dims={0,1},
    start_index_map={0,1}, index_vector_dim=2, slice_sizes={1,1,4},
    sharding={devices=[1,2,1,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="b"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "indices");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(
      instruction,
      op::Sharding("{devices=[1,2,1,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, GatherToIndex3) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %input = bf16[2,4819,4] parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %p1 = s32[2,2,1000] parameter(1)
  %indices = s32[2,2,1000] copy(%p1)
  ROOT %gather = bf16[2,1000,4]
    gather(bf16[2,4819,4] %input, s32[2,2,1000] %indices),
    offset_dims={2}, collapsed_slice_dims={0,1},
    start_index_map={0,1}, index_vector_dim=1, slice_sizes={1,1,4},
    sharding={devices=[1,2,1]0,1 metadata={op_name="b"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "indices");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,1,2]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, GatherToDataOperand) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %p0 = f32[2,9] parameter(0)
  %input = f32[2,9] copy(%p0)
  %indices = s32[3] parameter(1),
    sharding={replicated metadata={op_name="a"}}
  ROOT %gather = f32[3,9] gather(%input, %indices), offset_dims={1},
    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1,
    slice_sizes={1,9},
    sharding={devices=[1,2]0,1 metadata={op_name="b"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "input");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, GatherToDataOperand_PartialReplicate) {
  const char* hlo_string = R"(
HloModule module

ENTRY entry {
  %p0 = f32[2,9] parameter(0)
  %input = f32[2,9] copy(%p0)
  %indices = s32[3] parameter(1),
    sharding={replicated metadata={op_name="a"}}
  ROOT %gather = f32[3,9] gather(%input, %indices), offset_dims={1},
    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1,
    slice_sizes={1,9},
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="b"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "input");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, DataOperandToScatter) {
  const char* const hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={devices=[1,2]0,1 metadata={op_name="a"}}
  %indices = s32[3] parameter(1),
    sharding={replicated metadata={op_name="b"}}
  %updates = f32[3,9] parameter(2),
    sharding={replicated metadata={op_name="c"}}
  %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1
  ROOT %copy = f32[2,9] copy(%scatter)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "scatter");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, DataOperandToScatter_PartialReplicate) {
  const char* const hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="a"}}
  %indices = s32[3] parameter(1),
    sharding={replicated metadata={op_name="b"}}
  %updates = f32[3,9] parameter(2),
    sharding={replicated metadata={op_name="c"}}
  %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1
  ROOT %copy = f32[2,9] copy(%scatter)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "scatter");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, UpdateOperandToScatter) {
  const char* const hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %indices = s32[3] parameter(1),
    sharding={replicated metadata={op_name="b"}}
  %updates = f32[3,9] parameter(2),
    sharding={devices=[1,2]0,1 metadata={op_name="c"}}
  %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1
  ROOT %copy = f32[2,9] copy(%scatter)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "scatter");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("c")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, UpdateOperandToScatter_PartialReplicate) {
  const char* const hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %indices = s32[3] parameter(1),
    sharding={replicated metadata={op_name="b"}}
  %updates = f32[3,9] parameter(2),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="c"}}
  %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1
  ROOT %copy = f32[2,9] copy(%scatter)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "scatter");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("c")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ScatterToDataOperand_PartialReplicate) {
  const char* const hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %p0 = f32[2,9] parameter(0)
  %input = f32[2,9] copy(%p0)
  %indices = s32[3] parameter(1),
    sharding={replicated metadata={op_name="a"}}
  %updates = f32[3,9] parameter(2),
    sharding={replicated metadata={op_name="b"}}
  ROOT %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1,
      sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="c"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "input");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("c")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ScatterToDataOperand) {
  const char* const hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %p0 = f32[2,9] parameter(0)
  %input = f32[2,9] copy(%p0)
  %indices = s32[3] parameter(1),
    sharding={replicated metadata={op_name="a"}}
  %updates = f32[3,9] parameter(2),
    sharding={replicated metadata={op_name="b"}}
  ROOT %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1,
      sharding={devices=[1,2]0,1 metadata={op_name="c"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "input");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("c")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ScatterToUpdateOperand_PartialReplicate) {
  const char* const hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9] parameter(0)
  %indices = s32[3] parameter(1),
    sharding={replicated metadata={op_name="a"}}
  %p2 = f32[3,9] parameter(2)
  %updates = f32[3,9] copy(%p2)
  ROOT %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1,
      sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="b"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "updates");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ScatterToUpdateOperand) {
  const char* const hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9] parameter(0)
  %indices = s32[3] parameter(1),
    sharding={replicated metadata={op_name="a"}}
  %p2 = f32[3,9] parameter(2)
  %updates = f32[3,9] copy(%p2)
  ROOT %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1,
      sharding={devices=[1,2]0,1 metadata={op_name="b"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "updates");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[1,2]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ScatterUpdateToIndex) {
  const char* const hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %p1 = s32[3] parameter(1),
    sharding={replicated metadata={op_name="b"}}
  %indices = s32[3] copy(%p1)
  %updates = f32[3,9] parameter(2),
    sharding={devices=[2,1]0,1 metadata={op_name="c"}}
  ROOT %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1,
      sharding={replicated metadata={op_name="d"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "indices");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("c")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ScatterUpdateToIndex_PartialReplicate) {
  const char* const hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %p1 = s32[3] parameter(1),
    sharding={replicated metadata={op_name="b"}}
  %indices = s32[3] copy(%p1)
  %updates = f32[3,9] parameter(2),
    sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="c"}}
  ROOT %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1,
      sharding={replicated metadata={op_name="d"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "indices");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[2,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("c")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ScatterIndexToUpdate) {
  const char* const hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %indices = s32[3] parameter(1),
    sharding={devices=[2]0,1 metadata={op_name="b"}}
  %p2 = f32[3,9] parameter(2),
    sharding={replicated metadata={op_name="c"}}
  %updates = f32[3,9] copy(%p2)
  ROOT %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1,
      sharding={replicated metadata={op_name="d"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "updates");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[2,1]0,1}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ScatterIndexToUpdate_PartialReplicate) {
  const char* const hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={replicated metadata={op_name="a"}}
  %indices = s32[3] parameter(1),
    sharding={devices=[2,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="b"}}
  %p2 = f32[3,9] parameter(2),
    sharding={replicated metadata={op_name="c"}}
  %updates = f32[3,9] copy(%p2)
  ROOT %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1,
      sharding={replicated metadata={op_name="d"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/false, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "updates");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction,
              op::Sharding("{devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("b")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, PartialShardingOnElementwise) {
  const char* const hlo_string = R"(
HloModule module

ENTRY entry {
  %p0 = f32[2,9] parameter(0),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate metadata={op_name="a"}}
  %p1 = f32[2,9] parameter(1),
    sharding={devices=[2,1,2]0,2,1,3 last_tile_dim_replicate metadata={op_name="b"}}
  %lhs = f32[2,9] copy(%p0)
  %rhs = f32[2,9] copy(%p1)
  %add = f32[2,9] add(%lhs, %rhs)
  ROOT %copy = f32[2,9] copy(%add)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* lhs = FindInstruction(module.get(), "lhs");
  ASSERT_NE(lhs, nullptr);
  EXPECT_THAT(lhs, op::Sharding("{devices=[2,2]0,2,1,3}"));
  auto* rhs = FindInstruction(module.get(), "rhs");
  ASSERT_NE(rhs, nullptr);
  EXPECT_THAT(rhs, op::Sharding("{devices=[2,2]0,2,1,3}"));
  auto* add = FindInstruction(module.get(), "add");
  ASSERT_NE(add, nullptr);
  EXPECT_THAT(add, op::Sharding("{devices=[2,2]0,2,1,3}"));
  for (HloInstruction* instruction : {lhs, rhs, add}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("b"), CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, PartialShardingOnElementwise2) {
  const char* const hlo_string = R"(
HloModule module

ENTRY entry {
  %p0 = f32[2,9] parameter(0),
    sharding={devices=[1,2,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate metadata={op_name="a"}}
  %p1 = f32[2,9] parameter(1),
    sharding={devices=[2,1,4]0,1,4,5,2,3,6,7 last_tile_dim_replicate metadata={op_name="b"}}
  %lhs = f32[2,9] copy(%p0)
  %rhs = f32[2,9] copy(%p1)
  %add = f32[2,9] add(%lhs, %rhs)
  ROOT %copy = f32[2,9] copy(%add)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* lhs = FindInstruction(module.get(), "lhs");
  ASSERT_NE(lhs, nullptr);
  EXPECT_THAT(
      lhs,
      op::Sharding("{devices=[2,2,2]0,1,4,5,2,3,6,7 last_tile_dim_replicate}"));
  auto* rhs = FindInstruction(module.get(), "rhs");
  ASSERT_NE(rhs, nullptr);
  EXPECT_THAT(
      rhs,
      op::Sharding("{devices=[2,2,2]0,1,4,5,2,3,6,7 last_tile_dim_replicate}"));
  auto* add = FindInstruction(module.get(), "add");
  ASSERT_NE(add, nullptr);
  EXPECT_THAT(
      add,
      op::Sharding("{devices=[2,2,2]0,1,4,5,2,3,6,7 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(lhs->sharding(),
                ShardingMetadata({CreateMetadata("b"), CreateMetadata("a")}));
    EXPECT_THAT(rhs->sharding(),
                ShardingMetadata({CreateMetadata("b"), CreateMetadata("a")}));
    EXPECT_THAT(add->sharding(),
                ShardingMetadata({CreateMetadata("b"), CreateMetadata("a")}));

  } else {
    for (HloInstruction* instruction : {lhs, rhs}) {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, PartialShardingTransposeForwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %transpose {
  %param = f32[7,11,13]{2,1,0} parameter(0),
    sharding={devices=[2,1,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate metadata={op_name="a"}}
  %transpose = f32[11,13,7]{2,1,0} transpose(%param), dimensions={1,2,0}
  ROOT %copy = f32[11,13,7]{2,1,0} copy(%transpose)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "transpose");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(
      instruction,
      op::Sharding(
          "{devices=[1,2,2,2]0,1,4,5,2,3,6,7 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, PartialShardingTransposeBackwardPass) {
  const char* const hlo_string = R"(
HloModule module
ENTRY %transpose {
  %param = f32[7,11,13]{2,1,0} parameter(0)
  %copy = f32[7,11,13]{2,1,0} copy(%param)
  ROOT %transpose = f32[11,13,7]{2,1,0} transpose(%copy), dimensions={1,2,0},
    sharding={devices=[1,2,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate metadata={op_name="a"}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "copy");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(
      instruction,
      op::Sharding(
          "{devices=[2,1,2,2]0,1,4,5,2,3,6,7 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ParallelGatherFromOperandForwardPass) {
  const char* const hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0),
    sharding={devices=[8,1,1,1]0,1,4,5,2,3,6,7 metadata={op_name="a"}}
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2
  %concatenate.19 = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0}
  %gather = s32[8,4,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %parameter.0,
    s32[2,8,4]{2,1,0} %concatenate.19), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2}
  ROOT %copy = s32[8,4,2,2]{3,2,1,0} copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "gather");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[8,1,1,1]0,1,4,5,2,3,6,7}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ParallelGatherFromIndexForwardPass) {
  const char* const hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0)
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1,
    sharding={devices=[1,8,1]0,1,4,5,2,3,6,7 metadata={op_name="a"}}
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2
  %concatenate.19 = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0}
  %gather = s32[8,4,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %parameter.0,
    s32[2,8,4]{2,1,0} %concatenate.19), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2}
  ROOT %copy = s32[8,4,2,2]{3,2,1,0} copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "gather");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(instruction, op::Sharding("{devices=[8,1,1,1]0,1,4,5,2,3,6,7}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, ParallelGatherBackwardPass) {
  const char* const hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0)
  %copy.p = s32[8,4,2,2]{3,2,1,0} copy(%parameter.0)
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2
  %concatenate = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0}
  %gather = s32[8,4,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %copy.p,
    s32[2,8,4]{2,1,0} %concatenate), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2},
    sharding={devices=[8,1,1,1]0,1,4,5,2,3,6,7 metadata={op_name="a"}}
  ROOT %copy = s32[8,4,2,2]{3,2,1,0} copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* concatenate = FindInstruction(module.get(), "concatenate");
  ASSERT_NE(concatenate, nullptr);
  EXPECT_THAT(concatenate, op::Sharding("{devices=[1,8,1]0,1,4,5,2,3,6,7}"));
  auto* copy_p = FindInstruction(module.get(), "copy.p");
  ASSERT_NE(copy_p, nullptr);
  EXPECT_THAT(copy_p, op::Sharding("{devices=[8,1,1,1]0,1,4,5,2,3,6,7}"));
  for (HloInstruction* instruction : {concatenate, copy_p}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, ParallelGatherBackwardPass2) {
  const char* const hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[4,8,2,2]{3,2,1,0} parameter(0)
  %copy.p = s32[4,8,2,2]{3,2,1,0} copy(%parameter.0)
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2
  %concatenate = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0}
  %gather = s32[8,4,2,2]{3,2,1,0} gather(
    s32[4,8,2,2]{3,2,1,0} %copy.p,
    s32[2,8,4]{2,1,0} %concatenate), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={1,0}, index_vector_dim=0,
    slice_sizes={1,1,2,2},
    sharding={devices=[1,4,1,1]0,1,4,5 metadata={op_name="a"}}
  ROOT %copy = s32[8,4,2,2]{3,2,1,0} copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* concatenate = FindInstruction(module.get(), "concatenate");
  ASSERT_NE(concatenate, nullptr);
  EXPECT_THAT(concatenate, op::Sharding("{devices=[1,1,4]0,1,4,5}"));
  auto* copy_p = FindInstruction(module.get(), "copy.p");
  ASSERT_NE(copy_p, nullptr);
  EXPECT_THAT(copy_p, op::Sharding("{devices=[4,1,1,1]0,1,4,5}"));
  for (HloInstruction* instruction : {concatenate, copy_p}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest,
       PartialShardingParallelGatherFromOperandForwardPass) {
  const char* const hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0),
    sharding={devices=[4,1,1,1,2]0,1,4,5,2,3,6,7 last_tile_dim_replicate metadata={op_name="a"}}
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2
  %concatenate.19 = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0}
  %gather = s32[8,4,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %parameter.0,
    s32[2,8,4]{2,1,0} %concatenate.19), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2}
  ROOT %copy = s32[8,4,2,2]{3,2,1,0} copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "gather");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(
      instruction,
      op::Sharding(
          "{devices=[4,1,1,1,2]0,1,4,5,2,3,6,7 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest,
       PartialShardingParallelGatherFromIndexForwardPass) {
  const char* const hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0)
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1,
    sharding={devices=[1,4,1,2]0,1,4,5,2,3,6,7 last_tile_dim_replicate metadata={op_name="a"}}
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2
  %concatenate.19 = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0}
  %gather = s32[8,4,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %parameter.0,
    s32[2,8,4]{2,1,0} %concatenate.19), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2}
  ROOT %copy = s32[8,4,2,2]{3,2,1,0} copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* instruction = FindInstruction(module.get(), "gather");
  ASSERT_NE(instruction, nullptr);
  EXPECT_THAT(
      instruction,
      op::Sharding(
          "{devices=[4,1,1,1,2]0,1,4,5,2,3,6,7 last_tile_dim_replicate}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(instruction->sharding(),
                ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, PartialShardingParallelGatherBackwardPass) {
  const char* const hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0)
  %copy.p = s32[8,4,2,2]{3,2,1,0} copy(%parameter.0)
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2
  %concatenate = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0}
  %gather = s32[8,4,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %copy.p,
    s32[2,8,4]{2,1,0} %concatenate), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2},
    sharding={devices=[4,1,1,1,2]0,1,4,5,2,3,6,7 last_tile_dim_replicate metadata={op_name="a"}}
  ROOT %copy = s32[8,4,2,2]{3,2,1,0} copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* concatenate = FindInstruction(module.get(), "concatenate");
  ASSERT_NE(concatenate, nullptr);
  EXPECT_THAT(
      concatenate,
      op::Sharding(
          "{devices=[1,4,1,2]0,1,4,5,2,3,6,7 last_tile_dim_replicate}"));
  auto* copy_p = FindInstruction(module.get(), "copy.p");
  ASSERT_NE(copy_p, nullptr);
  EXPECT_THAT(
      copy_p,
      op::Sharding(
          "{devices=[4,1,1,1,2]0,1,4,5,2,3,6,7 last_tile_dim_replicate}"));
  for (HloInstruction* instruction : {concatenate, copy_p}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, PartialShardingParallelGatherBackwardPass2) {
  const char* const hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[4,8,2,2]{3,2,1,0} parameter(0)
  %copy.p = s32[4,8,2,2]{3,2,1,0} copy(%parameter.0)
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2
  %concatenate = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0}
  %gather = s32[8,4,2,2]{3,2,1,0} gather(
    s32[4,8,2,2]{3,2,1,0} %copy.p,
    s32[2,8,4]{2,1,0} %concatenate), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={1,0}, index_vector_dim=0,
    slice_sizes={1,1,2,2},
    sharding={devices=[1,2,1,1,2]0,1,4,5 last_tile_dim_replicate metadata={op_name="a"}}
  ROOT %copy = s32[8,4,2,2]{3,2,1,0} copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  auto* concatenate = FindInstruction(module.get(), "concatenate");
  ASSERT_NE(concatenate, nullptr);
  EXPECT_THAT(
      concatenate,
      op::Sharding("{devices=[1,1,2,2]0,1,4,5 last_tile_dim_replicate}"));
  auto* copy_p = FindInstruction(module.get(), "copy.p");
  ASSERT_NE(copy_p, nullptr);
  EXPECT_THAT(
      copy_p,
      op::Sharding("{devices=[2,1,1,1,2]0,1,4,5 last_tile_dim_replicate}"));
  for (HloInstruction* instruction : {concatenate, copy_p}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, GatherParallelAndPassthroughMerged) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %arg0 = s32[4,8,2,2]{3,2,1,0} parameter(0)
  %arg1 =  s32[4]{0} parameter(1)
  %input = s32[4,8,2,2]{3,2,1,0} copy(%arg0),
    sharding={devices=[2,1,2,1]0,1,4,5 metadata={op_name="a"}}
  %seq_size = s32[4]{0} copy(s32[4]{0} %arg1)
  %seq_b = s32[1,4,8]{2,1,0} broadcast(s32[4]{0} %seq_size
  ), dimensions={1}
  %iota.11 = s32[1,4,8]{2,1,0} iota(), iota_dimension=1
  %concatenate = s32[2,4,8]{2,1,0} concatenate(s32[1,4,8]{2,1,0} %iota.11,
    s32[1,4,8]{2,1,0} %seq_b), dimensions={0}
  %gather = s32[4,8,2,2]{3,2,1,0} gather(s32[4,8,2,2]{3,2,1,0} %input,
    s32[2,4,8]{2,1,0} %concatenate), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2}
  ROOT %copy = s32[4,8,2,2]{3,2,1,0} copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  const HloInstruction* input = FindInstruction(module.get(), "input");
  ASSERT_NE(input, nullptr);
  EXPECT_THAT(input, op::Sharding("{devices=[2,1,2,1]0,1,4,5 }"));
  const HloInstruction* concatenate =
      FindInstruction(module.get(), "concatenate");
  ASSERT_NE(concatenate, nullptr);
  EXPECT_THAT(
      concatenate,
      op::Sharding("{devices=[1,2,1,2]0,1,4,5 last_tile_dim_replicate}"));
  const HloInstruction* gather = FindInstruction(module.get(), "gather");
  ASSERT_NE(gather, nullptr);
  EXPECT_THAT(gather, op::Sharding("{devices=[2,1,2,1]0,1,4,5}"));

  for (const HloInstruction* instruction : {input, gather}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest, GatherParallelAndTrivialMerged) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %arg0 = s32[4,8,2,2]{3,2,1,0} parameter(0)
  %arg1 =  s32[4]{0} parameter(1)
  %input = s32[4,8,2,2]{3,2,1,0} copy(%arg0),
    sharding={devices=[2,2,1,1]0,1,4,5 metadata={op_name="a"}}
  %seq_size = s32[4]{0} copy(s32[4]{0} %arg1)
  %seq_b = s32[1,4,1]{2,1,0} broadcast(s32[4]{0} %seq_size), dimensions={1}
  %iota.11 = s32[1,4,1]{2,1,0} iota(), iota_dimension=1
  %concatenate = s32[2,4,1]{2,1,0} concatenate(s32[1,4,1]{2,1,0} %iota.11,
    s32[1,4,1]{2,1,0} %seq_b), dimensions={0}
  %gather = s32[4,1,2,2]{3,2,1,0} gather(s32[4,8,2,2]{3,2,1,0} %input,
    s32[2,4,1]{2,1,0} %concatenate), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2}
  ROOT %copy = s32[4,1,2,2]{3,2,1,0} copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  const HloInstruction* input = FindInstruction(module.get(), "input");
  ASSERT_NE(input, nullptr);
  EXPECT_THAT(input, op::Sharding("{devices=[2,2,1,1]0,1,4,5}"));
  const HloInstruction* concatenate =
      FindInstruction(module.get(), "concatenate");
  ASSERT_NE(concatenate, nullptr);
  EXPECT_THAT(
      concatenate,
      op::Sharding("{devices=[1,2,1,2]0,1,4,5 last_tile_dim_replicate}"));
  const HloInstruction* gather = FindInstruction(module.get(), "gather");
  ASSERT_NE(gather, nullptr);
  EXPECT_THAT(
      gather,
      op::Sharding("{devices=[2,1,1,1,2]0,1,4,5 last_tile_dim_replicate}"));
  for (const HloInstruction* instruction : {input, gather}) {
    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
      EXPECT_THAT(instruction->sharding(),
                  ShardingMetadata({CreateMetadata("a")}));
    } else {
      EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
    }
  }
}

TEST_P(ParameterizedMetadataTest,
       GatherParallelAndPassthroughMergedBackwardPass) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %arg0 = s32[4,8,2,2]{3,2,1,0} parameter(0)
  %arg1 =  s32[4]{0} parameter(1)
  %input = s32[4,8,2,2]{3,2,1,0} copy(%arg0)
  %seq_size = s32[4]{0} copy(s32[4]{0} %arg1)
  %seq_b = s32[1,4,8]{2,1,0} broadcast(s32[4]{0} %seq_size
  ), dimensions={1}
  %iota.11 = s32[1,4,8]{2,1,0} iota(), iota_dimension=1
  %concatenate = s32[2,4,8]{2,1,0} concatenate(s32[1,4,8]{2,1,0} %iota.11,
    s32[1,4,8]{2,1,0} %seq_b), dimensions={0}
  %gather = s32[2,2,4,8]{3,2,1,0} gather(s32[4,8,2,2]{3,2,1,0} %input,
    s32[2,4,8]{2,1,0} %concatenate), offset_dims={0,1},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2},
    sharding={devices=[2,1,2,1]0,4,1,5 metadata={op_name="a"}}
  ROOT %copy = s32[2,2,4,8]{3,2,1,0} copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);
  const HloInstruction* input = FindInstruction(module.get(), "input");
  ASSERT_NE(input, nullptr);
  EXPECT_THAT(input, op::Sharding("{devices=[2,1,2,1]0,1,4,5 }"));
  const HloInstruction* concatenate =
      FindInstruction(module.get(), "concatenate");
  ASSERT_NE(concatenate, nullptr);
  EXPECT_THAT(
      concatenate,
      op::Sharding("{devices=[1,2,1,2]0,1,4,5 last_tile_dim_replicate}"));
  const HloInstruction* gather = FindInstruction(module.get(), "gather");
  ASSERT_NE(gather, nullptr);
  EXPECT_THAT(gather, op::Sharding("{devices=[2,1,2,1]0,4,1,5}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(gather->sharding(), ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(gather->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, CorrectlyReplicateGatherIndex) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = bf16[1,2,2,2,8]{4,3,2,1,0} parameter(0)
  %parameter.1 = s32[1,2,2]{2,1,0} parameter(1)
  %index = s32[1,2,2]{2,1,0} copy(%parameter.1)
  %gather = bf16[1,2,2,2,8]{4,3,2,1,0} gather(
    bf16[1,2,2,2,8]{4,3,2,1,0} %parameter.0, s32[1,2,2]{2,1,0} %index),
    offset_dims={2,3,4}, collapsed_slice_dims={0,1}, start_index_map={0,1},
    index_vector_dim=2, slice_sizes={1,1,2,2,8},
    sharding={devices=[1,1,2,1,1]0,1 metadata={op_name="a"}}
  ROOT %copy = bf16[1,2,2,2,8]{4,3,2,1,0} copy(%gather)
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);

  const HloInstruction* index = FindInstruction(module.get(), "index");

  ASSERT_NE(index, nullptr);

  EXPECT_THAT(index, op::Sharding("{replicated}"));
  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
    EXPECT_THAT(index->sharding(), ShardingMetadata({CreateMetadata("a")}));
  } else {
    EXPECT_THAT(index->sharding(), ShardingMetadata({}));
  }
}

TEST_P(ParameterizedMetadataTest, GatherToOperand_ParallelDimIsNotPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[2,1000,1]{2,1,0} parameter(0)
  %parameter.1 = bf16[2,4819,4]{2,1,0} parameter(1)
  %iota = s32[2,1000,1]{1,0,2} iota(), iota_dimension=0
  %operand = bf16[2,4819,4]{2,1,0} copy(%parameter.1)
  %index = s32[2,1000,2]{2,1,0} concatenate(s32[2,1000,1]{1,0,2} %parameter.0,
    s32[2,1000,1]{2,1,0} %iota), dimensions={2},
    sharding={devices=[1,4,1]0,1,2,3}
  ROOT %gather = bf16[2,1000,4]{2,1,0} gather(bf16[2,4819,4]{2,1,0} %operand,
    s32[2,1000,2]{2,1,0} %index), offset_dims={2},
    collapsed_slice_dims={0,1}, start_index_map={0,1},
    index_vector_dim=2, slice_sizes={1,1,4},
    sharding={devices=[1,4,1]0,1,2,3}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(hlo_string));
  if (GetParam().clear_metadata) {
    ClearMetadata(module.get());
  }
  TF_ASSERT_OK_AND_ASSIGN(
      bool changed,
      ShardingPropagation(/*is_spmd=*/true, GetParam().propagate_metadata)
          .Run(module.get()));
  EXPECT_TRUE(changed);

  const HloInstruction* operand = FindInstruction(module.get(), "operand");
  EXPECT_THAT(operand, op::Sharding("{replicated}"));
}

}  // namespace
}  // namespace xla
