Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc478 auto offset_on_padded_shape = in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() local
483 offset_on_padded_shape.Calculate(shard_counts[i] - 1) + in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
489 offset_on_padded_shape.Calculate(partition_ordinals[dim], b), zero, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
709 auto offset_on_padded_shape = in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS() local
716 offset_on_padded_shape.Calculate(partition_ordinals[dim], b), zero, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
Dspmd_partitioner_util.h258 HloInstruction* offset_on_padded_shape, HloInstruction* pad_value,
Dspmd_partitioner_util.cc928 HloInstruction* offset_on_padded_shape, HloInstruction* pad_value, in ExchangeHaloAndGetValidData() argument
1025 index_shape, offset_on_padded_shape, {})); in ExchangeHaloAndGetValidData()