Searched refs:GroupedSharding (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner_util.h | 297 struct GroupedSharding { struct
298 GroupedSharding(std::vector<std::vector<int64>> device_groups, in GroupedSharding() function
315 GroupedSharding GroupShardingOnDims(const HloSharding& sharding,
320 GroupedSharding GroupShardingOnDims(const HloSharding& sharding,
324 HloSharding UngroupSharding(const GroupedSharding& grouped_sharding);
328 GroupedSharding AlignGroupsWith(GroupedSharding grouped_sharding,
329 const GroupedSharding& reference,
342 Shape GetPerGroupBaseShape(const GroupedSharding& grouped_sharding,
|
| D | spmd_partitioner_util.cc | 1378 GroupedSharding GroupShardingOnDims(const HloSharding& sharding, in GroupShardingOnDims()
1384 GroupedSharding GroupShardingOnDims(const HloSharding& sharding, in GroupShardingOnDims()
1409 auto grouped = GroupedSharding( in GroupShardingOnDims()
1435 HloSharding UngroupSharding(const GroupedSharding& grouped_sharding) { in UngroupSharding()
1482 GroupedSharding AlignGroupsWith(GroupedSharding grouped_sharding, in AlignGroupsWith()
1483 const GroupedSharding& reference, in AlignGroupsWith()
1556 Shape GetPerGroupBaseShape(const GroupedSharding& grouped_sharding, in GetPerGroupBaseShape()
|
| D | gather_scatter_handler.cc | 341 GroupedSharding grouped_indices = in PartitionIndexParallelDimensions()
343 GroupedSharding grouped_operand = in PartitionIndexParallelDimensions()
441 GroupedSharding grouped_output = in PartitionIndexParallelDimensions()
|
| D | dot_handler.cc | 1880 GroupedSharding GetNonContractingPartitionGroupedShardingForMatchedOperand( in GetNonContractingPartitionGroupedShardingForMatchedOperand()
1897 GroupedSharding output_grouped = in GetNonContractingPartitionGroupedShardingForMatchedOperand()
1910 absl::optional<GroupedSharding>
1927 GroupedSharding output_grouped = in GetNonContractingPartitionGroupedShardingForOtherOperand()
2008 GroupedSharding output_grouped = in PartitionDotGroupOnNonContracting()
2010 GroupedSharding matching_grouped = in PartitionDotGroupOnNonContracting()
2018 absl::optional<GroupedSharding> other_grouped = in PartitionDotGroupOnNonContracting()
2371 GroupedSharding output_grouped = in LhsIsBestMatchForNonContractingPartitioning()
2373 GroupedSharding matching_grouped = in LhsIsBestMatchForNonContractingPartitioning()
2377 absl::optional<GroupedSharding> other_grouped = in LhsIsBestMatchForNonContractingPartitioning()
|