Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.cc418 auto padded_dst_shape = in TileToPartialReplicateHaloExchange() local
420 if (ShapeUtil::Compatible(padded_dst_shape, hlo->shape())) { in TileToPartialReplicateHaloExchange()
436 padded_dst_shape.dimensions(dim) / dst_shard_count; in TileToPartialReplicateHaloExchange()
510 auto padded_dst_shape = in PadFromPartialReplicateShape() local
512 if (ShapeUtil::Compatible(padded_dst_shape, hlo->shape())) { in PadFromPartialReplicateShape()
532 padded_dst_shape.dimensions(dim) / src_shard_count; in PadFromPartialReplicateShape()
583 int64{0}, padded_dst_shape.dimensions(dim) - in PadFromPartialReplicateShape()
617 padded_dst_shape.dimensions(dim) - padded_src_shape.dimensions(dim)); in PadFromPartialReplicateShape()
619 padded_dst_shape.dimensions(dim) - in PadFromPartialReplicateShape()