Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.cc276 Shape GetPaddedShapeForUnevenPartitioning(const Shape& base_shape, in GetPaddedShapeForUnevenPartitioning() function
296 GetPaddedShapeForUnevenPartitioning(hlo->shape(), sharding); in PadBaseShapeBeforeUnevenTiledSharding()
416 GetPaddedShapeForUnevenPartitioning(base_shape, src_sharding); in TileToPartialReplicateHaloExchange()
419 GetPaddedShapeForUnevenPartitioning(base_shape, dst_sharding); in TileToPartialReplicateHaloExchange()
509 GetPaddedShapeForUnevenPartitioning(base_shape, src_sharding); in PadFromPartialReplicateShape()
511 GetPaddedShapeForUnevenPartitioning(base_shape, dst_sharding); in PadFromPartialReplicateShape()
Dspmd_partitioner_util.h121 Shape GetPaddedShapeForUnevenPartitioning(const Shape& base_shape,
Ddot_handler.cc612 padded_result_buffer_shape = GetPaddedShapeForUnevenPartitioning( in PartitionBaseCase()