Searched refs:src_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
416 GetPaddedShapeForUnevenPartitioning(base_shape, src_sharding); in TileToPartialReplicateHaloExchange()
453 int64 replicate_factor = src_sharding.tile_assignment().dim(dim) / in TileToPartialReplicateHaloExchange()
468 src_sharding, collective_ops_creator, next_channel_id, b); in TileToPartialReplicateHaloExchange()
491 src_sharding.tile_assignment().dim(dim)); in TileToPartialReplicateHaloExchange()
504 const HloSharding& src_sharding, const HloSharding& dst_sharding, in PadFromPartialReplicateShape() argument
509 GetPaddedShapeForUnevenPartitioning(base_shape, src_sharding); in PadFromPartialReplicateShape()
517 MakeTiledPartitionOrdinals(src_sharding, partition_id, b); in PadFromPartialReplicateShape()
526 int64 src_shard_count = src_sharding.tile_assignment().dim(dim); in PadFromPartialReplicateShape()
568 src_sharding, collective_ops_creator, next_channel_id, b); in PadFromPartialReplicateShape()
|