Searched refs:ShardingMetadata (Results 1 – 6 of 6) sorted by relevance
154 HloDomainRemover remover(ShardingMetadata::KindName(), in TEST_F()155 ShardingMetadata::NormalizeShardingDomain); in TEST_F()224 HloDomainRemover remover(ShardingMetadata::KindName(), in TEST_F()225 ShardingMetadata::NormalizeShardingDomain); in TEST_F()276 HloDomainRemover remover(ShardingMetadata::KindName(), in TEST_F()277 ShardingMetadata::NormalizeShardingDomain); in TEST_F()324 HloDomainRemover sharding_remover(ShardingMetadata::KindName(), in TEST_F()325 ShardingMetadata::NormalizeShardingDomain); in TEST_F()410 HloDomainRemover remover(ShardingMetadata::KindName(), in TEST_F()411 ShardingMetadata::NormalizeShardingDomain); in TEST_F()[all …]
402 std::unique_ptr<DomainMetadata> ShardingMetadata::Clone() const { in Clone()407 return absl::make_unique<ShardingMetadata>(std::move(sharding)); in Clone()410 bool ShardingMetadata::Matches(const DomainMetadata& other) const { in Matches()411 const ShardingMetadata* other_ptr = in Matches()412 dynamic_cast<const ShardingMetadata*>(&other); in Matches()425 size_t ShardingMetadata::Hash() const { in Hash()432 string ShardingMetadata::ToString() const { in ToString()436 /*static*/ StatusOr<const ShardingMetadata*>437 ShardingMetadata::ToShardingMetadata(const DomainMetadata* metadata) { in ToShardingMetadata()438 if (metadata->Kind() != ShardingMetadata::KindName()) { in ToShardingMetadata()[all …]
28 class ShardingMetadata : public DomainMetadata {30 explicit ShardingMetadata(std::shared_ptr<const HloSharding> sharding) in ShardingMetadata() function47 static StatusOr<const ShardingMetadata*> ToShardingMetadata(
1935 if (*kind == ShardingMetadata::KindName()) { in ParseDomain()1941 absl::make_unique<ShardingMetadata>(std::move(entry_sharding_ptr)); in ParseDomain()1943 absl::make_unique<ShardingMetadata>(std::move(exit_sharding_ptr)); in ParseDomain()
2622 dynamic_cast<const ShardingMetadata*>(operand_side_metadata_.get()); in ToProto()2629 dynamic_cast<const ShardingMetadata*>(user_side_metadata_.get()); in ToProto()
583 absl::make_unique<ShardingMetadata>(entry_hlo_sharding), in CreateFromProto()584 absl::make_unique<ShardingMetadata>(exit_hlo_sharding)); in CreateFromProto()