Searched refs:shard_size_with_halo (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | convolution_handler.cc | 468 int64 shard_size_with_halo = new_window.dimensions(i).size(); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() local
488 padded_full_shape_size, shard_size_with_halo, dim, rhs.sharding(), in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
694 int64 shard_size_with_halo = in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS() local
715 padded_full_shape_size, shard_size_with_halo, dim, lhs.sharding(), in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
|
| D | spmd_partitioner_util.h | 257 int64 shard_size_with_halo, int64 dim, const HloSharding& target,
|
| D | spmd_partitioner_util.cc | 927 int64 shard_size_with_halo, int64 dim, const HloSharding& target, in ExchangeHaloAndGetValidData() argument
962 shard_size_with_halo - concat->shape().dimensions(dim) - in ExchangeHaloAndGetValidData()
987 if (shard_size_with_halo != concat->shape().dimensions(dim)) { in ExchangeHaloAndGetValidData()
989 CHECK_LT(shard_size_with_halo, concat->shape().dimensions(dim)); in ExchangeHaloAndGetValidData()
991 slice_shape.set_dimensions(dim, shard_size_with_halo); in ExchangeHaloAndGetValidData()
|