Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc340 std::vector<OffsetCalculation> left_halo_size_functions( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() local
379 left_halo_size_functions[rhs_dimension] = in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
395 left_halo_size_functions[rhs_dimension].MaxInRange(1, shard_counts[i]) + in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
412 left_halo_size_functions[rhs_dimension], in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
467 left_halo_size_functions[dim].Calculate(0); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
481 left_halo_size_functions[dim]; in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
486 rhs_with_halo, rhs.base_shape(), left_halo_size_functions[dim], in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
649 std::vector<OffsetCalculation> left_halo_size_functions( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS() local
679 left_halo_size_functions[lhs_dimension] = in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
713 lhs_with_halo, lhs.base_shape(), left_halo_size_functions[dim], in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
Dspmd_partitioner_util.h228 std::vector<OffsetCalculation> left_halo_size_functions,
Dspmd_partitioner_util.cc901 std::vector<OffsetCalculation> left_halo_size_functions, in ExchangeHalo() argument
906 CHECK(left_halo_size_functions.size() == hlo->shape().rank()); in ExchangeHalo()
911 auto concat = ExchangeHalo(visiting_hlo, left_halo_size_functions[dim], in ExchangeHalo()
Dspmd_partitioner.cc776 std::vector<OffsetCalculation> left_halo_size_functions(base_shape_.rank()); in ReshardAsWindowedInput() local
795 left_halo_size_functions[dim] = in ReshardAsWindowedInput()
805 visiting_hlo, base_shape_, left_halo_size_functions[dim], in ReshardAsWindowedInput()