Home
last modified time | relevance | path

Searched refs:ShardCountAtDim (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc102 (ShardCountAtDim(lhs.sharding(), dnums.input_batch_dimension()) == in PartitionConvolutionWithBatchGroupCount()
105 (ShardCountAtDim(rhs.sharding(), in PartitionConvolutionWithBatchGroupCount()
193 (ShardCountAtDim(lhs.sharding(), dnums.input_feature_dimension()) == in PartitionConvolutionWithFeatureGroupCount()
196 (ShardCountAtDim(rhs.sharding(), in PartitionConvolutionWithFeatureGroupCount()
Dspmd_partitioner_util.h281 int64 ShardCountAtDim(const HloSharding& sharding, int64 dim);
Dspmd_partitioner_util.cc1276 int64 ShardCountAtDim(const HloSharding& sharding, int64 dim) { in ShardCountAtDim() function