Searched refs:lhs_dimension (Results 1 – 1 of 1) sorted by relevance
323 int64 lhs_dimension = dnums.input_spatial_dimensions(i); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() local332 CeilOfRatio(lhs.base_shape().dimensions(lhs_dimension), shard_count); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()360 int64 lhs_dimension = dnums.input_spatial_dimensions(i); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() local375 int64 base = lhs.base_shape().dimensions(lhs_dimension); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()430 pad_config.mutable_dimensions(lhs_dimension) in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()432 pad_config.mutable_dimensions(lhs_dimension) in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()434 pad_shape.set_dimensions(lhs_dimension, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()436 dynamic_slice_start_indices[lhs_dimension] = in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()440 .Calculate(partition_ordinals[lhs_dimension], b); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()441 dynamic_slice_shape.set_dimensions(lhs_dimension, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()[all …]