Searched refs:root_sharding (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_sharding_metadata.cc | 477 auto root_sharding = root->sharding_ptr(); in operator ()() local
479 if (instruction_sharding == nullptr && root_sharding == nullptr) { in operator ()()
483 if (instruction_sharding != nullptr && root_sharding != nullptr && in operator ()()
484 ShardingMatches(*instruction_sharding, *root_sharding)) { in operator ()()
491 if (root_sharding != nullptr) { in operator ()()
492 root_sharding = CloneShardingForDomain(root_sharding); in operator ()()
507 << (root_sharding != nullptr ? root_sharding->ToString() : "None"); in operator ()()
512 absl::make_unique<ShardingMetadata>(root_sharding), in operator ()()
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.h | 199 const HloSharding& root_sharding,
498 const HloSharding& root_sharding,
|
| D | spmd_partitioner.cc | 3382 HloComputation* computation, const HloSharding& root_sharding, in DoPartition() argument
3391 GetPartitionedHlo(computation->root_instruction()).Reshard(root_sharding); in DoPartition()
3633 HloComputation* computation, const HloSharding& root_sharding, in PartitionComputation() argument
3638 return visitor->DoPartition(computation, root_sharding, options_); in PartitionComputation()
3677 HloSharding root_sharding = entry_root->sharding(); in Run() local
3680 PartitionComputation(module->entry_computation(), root_sharding, in Run()
3783 const HloSharding& root_sharding = entry->root_instruction()->sharding(); in PreprocessSharding() local
3784 TF_RET_CHECK(root_sharding.IsReplicated() || in PreprocessSharding()
3785 root_sharding.UniqueDevice().has_value()) in PreprocessSharding()
3786 << "Unsupported entry root sharding: " << root_sharding.ToString(); in PreprocessSharding()
|