Searched refs:new_tile_assignment (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_sharding_util.cc | 400 Array<int64> new_tile_assignment = sharding.tile_assignment(); in ReshapeSharding() local
405 new_tile_assignment.Reshape(target_tile_assignment_dimensions); in ReshapeSharding()
407 ? HloSharding::PartialTile(new_tile_assignment, in ReshapeSharding()
409 : HloSharding::Tile(new_tile_assignment, sharding.metadata()); in ReshapeSharding()
418 Array<int64> new_tile_assignment(sharding.tile_assignment().dimensions()); in ReverseSharding() local
419 new_tile_assignment.Each([&](absl::Span<const int64> indices, int64* device) { in ReverseSharding()
423 new_tile_assignment.dim(d) - 1 - original_indices[d]; in ReverseSharding()
428 ? HloSharding::PartialTile(new_tile_assignment, in ReverseSharding()
430 : HloSharding::Tile(new_tile_assignment, sharding.metadata()); in ReverseSharding()
524 Array<int64> new_tile_assignment = index_sharding.tile_assignment(); in GatherOutputSharding() local
[all …]
|
| D | sharding_propagation.cc | 765 Array<int64> new_tile_assignment = op->sharding().tile_assignment(); in InferShardingFromOperands() local
766 new_tile_assignment.Reshape(target_tile_assignment_dimensions); in InferShardingFromOperands()
769 ? HloSharding::PartialTile(new_tile_assignment, in InferShardingFromOperands()
771 : HloSharding::Tile(new_tile_assignment, in InferShardingFromOperands()
1194 auto new_tile_assignment = in GetShardingFromUser() local
1196 if (new_tile_assignment.num_elements() == 1) { in GetShardingFromUser()
1197 return HloSharding::AssignDevice(*new_tile_assignment.begin(), in GetShardingFromUser()
1200 return HloSharding::Tile(new_tile_assignment, user.sharding().metadata()); in GetShardingFromUser()
|
| D | hlo_sharding.cc | 57 auto new_tile_assignment = Array<int64>(new_tile_dims); in PartialTile() local
58 new_tile_assignment.Each([&](absl::Span<const int64> indices, int64* device) { in PartialTile()
64 return PartialTile(new_tile_assignment, metadata); in PartialTile()
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner_util.cc | 1906 auto new_tile_assignment = to_adjust->tile_assignment(); in GatherOperandsShardedAcrossParallelDims() local
1907 new_tile_assignment.Reshape(new_tile_assignment_dims); in GatherOperandsShardedAcrossParallelDims()
1910 AlignShardingOnDims(HloSharding::PartialTile(new_tile_assignment), in GatherOperandsShardedAcrossParallelDims()
1913 *to_adjust = AlignShardingOnDims(HloSharding::Tile(new_tile_assignment), in GatherOperandsShardedAcrossParallelDims()
|