Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/
Dcollective_decomposer_utils.cc36 absl::Span<const ReplicaGroup> replica_groups, const Shape &shard_shape, in CreateStartIndicesForCollectiveDecomposition()
/external/tensorflow/tensorflow/core/tpu/kernels/xla/
Dinfeed_op.cc58 auto shard_shape = in UpdateInfeedLayout() local
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner.cc558 auto shard_shape = MakePartitionedShape(shape, target); in ReshardNoCache() local
673 auto shard_shape = base_shape_; in ReshardAsWindowedInput() local
1181 const Shape& shard_shape = hlo()->shape(); in ReplicatePartial() local
1420 auto shard_shape = MakePartitionedShape(base_shape_, temp_target_sharding); in ReshardFromPartialReplicateWithDynamicSlice() local
2295 const Shape shard_shape = MakePartitionedShape(hlo->shape(), hlo->sharding()); in HandleConcatenate() local
2434 auto shard_shape = MakePartitionedShape(hlo->shape(), sharding); in HandleSlice() local
2500 auto shard_shape = MakePartitionedShape(replicated_shape, topk_sharding); in HandleSort() local
2977 auto shard_shape = MakePartitionedShape(hlo->shape(), hlo->sharding()); in HandleConstant() local
3262 auto shard_shape = MakePartitionedShape(shape, sharding); in HandleInfeed() local
3447 auto shard_shape = MakePartitionedShape(hlo->shape(), hlo->sharding()); in HandlePad() local
[all …]
Dcustom_call_handler.cc131 auto shard_shape = in HandleCustomCallTopK() local
Dspmd_partitioner_util.cc165 auto shard_shape = MakePartitionedShape(shape, sharding); in MakePartitionOffsets() local
209 auto shard_shape = MakePartitionedShape(base_shape, sharding); in GetPaddedShapeForUnevenPartitioning() local
1613 auto shard_shape = in PerGroupSliceFromReplicated() local
Dconvolution_handler.cc805 auto shard_shape = MakePartitionedShape(output_base_shape, output_sharding); in PartitionConvolutionTiledOutput() local