Searched refs:right_halo_size_function (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner_util.h | 219 const OffsetCalculation& right_halo_size_function, int64 dim,
255 const OffsetCalculation& right_halo_size_function,
|
| D | spmd_partitioner_util.cc | 461 OffsetCalculation right_halo_size_function = in TileToPartialReplicateHaloExchange() local
467 result, left_halo_size_function, right_halo_size_function, dim, in TileToPartialReplicateHaloExchange()
559 OffsetCalculation right_halo_size_function = in PadFromPartialReplicateShape() local
567 result, left_halo_size_function, right_halo_size_function, dim, in PadFromPartialReplicateShape()
581 right_halo_size_function.MaxInRange(0, src_shard_count - 1); in PadFromPartialReplicateShape()
800 const OffsetCalculation& right_halo_size_function, int64 dim, in ExchangeHalo() argument
811 right_halo_size_function.MaxInRange(0, shard_count - 1); in ExchangeHalo()
925 const OffsetCalculation& right_halo_size_function, in ExchangeHaloAndGetValidData() argument
933 ExchangeHalo(hlo, left_halo_size_function, right_halo_size_function, dim, in ExchangeHaloAndGetValidData()
|
| D | fft_handler.cc | 72 OffsetCalculation right_halo_size_function = in PadEachPartitionWithHaloExchange() local
80 ExchangeHalo(hlo, left_halo_size_function, right_halo_size_function, in PadEachPartitionWithHaloExchange()
|