Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc342 std::vector<OffsetCalculation> right_halo_size_functions( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() local
385 right_halo_size_functions[rhs_dimension] = in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
396 right_halo_size_functions[rhs_dimension].MaxInRange( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
487 right_halo_size_functions[dim], explicit_left_padding_on_full_shape, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
651 std::vector<OffsetCalculation> right_halo_size_functions( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS() local
683 right_halo_size_functions[lhs_dimension] = in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
714 right_halo_size_functions[dim], explicit_left_padding_on_full_shape, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
Dspmd_partitioner_util.h229 std::vector<OffsetCalculation> right_halo_size_functions,
Dspmd_partitioner_util.cc902 std::vector<OffsetCalculation> right_halo_size_functions, in ExchangeHalo() argument
907 CHECK(right_halo_size_functions.size() == hlo->shape().rank()); in ExchangeHalo()
912 right_halo_size_functions[dim], dim, target, in ExchangeHalo()
Dspmd_partitioner.cc777 std::vector<OffsetCalculation> right_halo_size_functions(base_shape_.rank()); in ReshardAsWindowedInput() local
801 right_halo_size_functions[dim] = in ReshardAsWindowedInput()
806 right_halo_size_functions[dim], explicit_left_padding[dim], in ReshardAsWindowedInput()