Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc348 auto partition_ordinals = in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() local
440 .Calculate(partition_ordinals[lhs_dimension], b); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
489 offset_on_padded_shape.Calculate(partition_ordinals[dim], b), zero, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
490 partition_ordinals[dim], collective_ops_creator, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
655 auto partition_ordinals = in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS() local
716 offset_on_padded_shape.Calculate(partition_ordinals[dim], b), zero, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
717 partition_ordinals[dim], lhs.state().collective_ops_creator, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
Dspmd_partitioner.cc566 auto partition_ordinals = in ReshardAsWindowedInput() local
633 partition_ordinals[i], state_.b); in ReshardAsWindowedInput()
687 start_on_padded_calculations[i].Calculate(partition_ordinals[i], in ReshardAsWindowedInput()
700 .Calculate(partition_ordinals[i], state_.b); in ReshardAsWindowedInput()
808 offsets_on_padded_shape[dim], pad_value, partition_ordinals[dim], in ReshardAsWindowedInput()
2155 auto partition_ordinals = in HandleIota() local
2161 partition_ordinals[dimension], multiplier)); in HandleIota()
2376 auto partition_ordinals = in HandleDynamicUpdateSlice() local
2401 offset_shape, HloOpcode::kMultiply, partition_ordinals[dim], in HandleDynamicUpdateSlice()
2413 offset_shape, HloOpcode::kAdd, partition_ordinals[dim], in HandleDynamicUpdateSlice()
[all …]
Dfft_handler.cc101 auto partition_ordinals = in PadEachPartitionWithHaloExchange() local
105 partition_ordinals[concat->shape().rank() - 1], b); in PadEachPartitionWithHaloExchange()
Dspmd_partitioner_util.cc424 auto partition_ordinals = in TileToPartialReplicateHaloExchange() local
495 partition_ordinals[dim], b); in TileToPartialReplicateHaloExchange()
516 auto partition_ordinals = in PadFromPartialReplicateShape() local
603 partition_ordinals[dim], b); in PadFromPartialReplicateShape()