Searched refs:dims_mapping (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | dot_handler.cc | 340 const DotConvDimsMapping& dims_mapping, int64 lhs_rank, int64 rhs_rank, in ComputeDimensionIndexMapping() argument
363 for (const auto& mapping : dims_mapping.batch_dims) { in ComputeDimensionIndexMapping()
366 for (const auto& mapping : dims_mapping.contracting_dims) { in ComputeDimensionIndexMapping()
369 for (const auto& mapping : dims_mapping.lhs_non_contracting_dims) { in ComputeDimensionIndexMapping()
372 for (const auto& mapping : dims_mapping.rhs_non_contracting_dims) { in ComputeDimensionIndexMapping()
375 for (const auto& mapping : dims_mapping.conv_spatial_dims) { in ComputeDimensionIndexMapping()
468 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionBaseCase() argument
492 dims_mapping, lhs.base_shape().rank(), rhs.base_shape().rank(), in PartitionBaseCase()
1515 for (const auto& cd : dims_mapping.contracting_dims) { in PartitionBaseCase()
1625 for (const auto& cd : dims_mapping.contracting_dims) { in PartitionBaseCase()
[all …]
|
| D | convolution_handler.h | 31 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping,
|
| D | spmd_partitioner.h | 443 const DotConvDimsMapping& dims_mapping,
|
| D | convolution_handler.cc | 966 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionConvolution() argument
|