Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Ddot_handler.cc43 DotConvDimsMapping mapping; in HandleDot()
340 const DotConvDimsMapping& dims_mapping, int64 lhs_rank, int64 rhs_rank, in ComputeDimensionIndexMapping()
349 [&](const DotConvDimsMapping::DimsMapping& mapping) { in ComputeDimensionIndexMapping()
468 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionBaseCase()
1638 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping,
1650 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionDotGroupOnBatch()
1883 absl::Span<const DotConvDimsMapping::DimsMapping> partitioned_dims) { in GetNonContractingPartitionGroupedShardingForMatchedOperand()
1917 absl::Span<const DotConvDimsMapping::DimsMapping> matching_partitioned_dims, in GetNonContractingPartitionGroupedShardingForOtherOperand()
1918 absl::Span<const DotConvDimsMapping::DimsMapping> in GetNonContractingPartitionGroupedShardingForOtherOperand()
1920 absl::Span<const DotConvDimsMapping::DimsMapping> other_contracting_dims) { in GetNonContractingPartitionGroupedShardingForOtherOperand()
[all …]
Dconvolution_handler.h31 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping,
Dspmd_partitioner.h381 struct DotConvDimsMapping { struct
443 const DotConvDimsMapping& dims_mapping,
Dconvolution_handler.cc966 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionConvolution()
989 spmd::DotConvDimsMapping mapping; in HandleConvolution()