Home
last modified time | relevance | path

Searched refs:dst_per_shard_size (Results 1 – 1 of 1) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.cc435 int64 dst_per_shard_size = in TileToPartialReplicateHaloExchange() local
439 if (src_per_shard_size <= dst_per_shard_size || dst_shard_count == 1) { in TileToPartialReplicateHaloExchange()
457 src_per_shard_size - dst_per_shard_size, 0, replicate_factor)); in TileToPartialReplicateHaloExchange()
485 dst_per_shard_size - src_per_shard_size, in TileToPartialReplicateHaloExchange()
486 (src_per_shard_size - dst_per_shard_size) * (dst_shard_count - 1), in TileToPartialReplicateHaloExchange()
531 int64 dst_per_shard_size = in PadFromPartialReplicateShape() local
535 if (src_per_shard_size >= dst_per_shard_size) { in PadFromPartialReplicateShape()
561 dst_per_shard_size - src_per_shard_size, in PadFromPartialReplicateShape()
562 dst_per_shard_size - src_per_shard_size, 1)); in PadFromPartialReplicateShape()
597 dst_per_shard_size - src_per_shard_size, 0, 1)); in PadFromPartialReplicateShape()
[all …]