Searched refs:dst_sharding (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner_util.h | 368 const HloSharding& src_sharding, const HloSharding& dst_sharding,
391 const HloSharding& src_sharding, const HloSharding& dst_sharding,
|
| D | spmd_partitioner_util.cc | 410 const HloSharding& src_sharding, const HloSharding& dst_sharding, in TileToPartialReplicateHaloExchange() argument
419 GetPaddedShapeForUnevenPartitioning(base_shape, dst_sharding); in TileToPartialReplicateHaloExchange()
425 MakeTiledPartitionOrdinals(dst_sharding, partition_id, b); in TileToPartialReplicateHaloExchange()
430 int64 dst_shard_count = dst_sharding.tile_assignment().dim(dim); in TileToPartialReplicateHaloExchange()
454 dst_sharding.tile_assignment().dim(dim); in TileToPartialReplicateHaloExchange()
504 const HloSharding& src_sharding, const HloSharding& dst_sharding, in PadFromPartialReplicateShape() argument
511 GetPaddedShapeForUnevenPartitioning(base_shape, dst_sharding); in PadFromPartialReplicateShape()
|