Home
last modified time | relevance | path

Searched defs:shard_shape (Results 1 – 4 of 4) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner.cc470 auto shard_shape = MakePartitionedShape(shape, target); in ReshardNoCache() local
568 auto shard_shape = base_shape_; in ReshardAsWindowedInput() local
868 const Shape& shard_shape = hlo()->shape(); in ReplicatePartial() local
1024 auto shard_shape = MakePartitionedShape(base_shape_, temp_target_sharding); in ReshardFromPartialReplicateWithDynamicSlice() local
1491 const Shape shard_shape = MakePartitionedShape(hlo->shape(), hlo->sharding()); in HandleConcatenate() local
1627 auto shard_shape = MakePartitionedShape(hlo->shape(), sharding); in HandleSlice() local
1691 auto shard_shape = MakePartitionedShape(replicated_shape, topk_sharding); in HandleSort() local
1845 auto shard_shape = in HandleCustomCall() local
2274 auto shard_shape = MakePartitionedShape(hlo->shape(), hlo->sharding()); in HandleConstant() local
2507 auto shard_shape = MakePartitionedShape(shape, sharding); in HandleInfeed() local
[all …]
Dspmd_partitioner_util.cc214 auto shard_shape = MakePartitionedShape(shape, sharding); in MakePartitionOffsets() local
284 auto shard_shape = MakePartitionedShape(base_shape, sharding); in GetPaddedShapeForUnevenPartitioning() local
1726 auto shard_shape = in PerGroupSliceFromReplicated() local
Dconvolution_handler.cc803 auto shard_shape = MakePartitionedShape(output_base_shape, output_sharding); in PartitionConvolutionTiledOutput() local
/external/tensorflow/tensorflow/core/tpu/kernels/xla/
Dinfeed_op.cc58 auto shard_shape = in UpdateInfeedLayout() local