Searched refs:lhs_batch_partitions (Results 1 – 1 of 1) sorted by relevance
388 int64 lhs_non_contracting_partitions, int64 lhs_batch_partitions, in GetWindowedEinsumConfiguration() argument436 if (lhs_batch_partitions == num_partitions) { in GetWindowedEinsumConfiguration()474 int64 lhs_batch_partitions, int64 rhs_batch_partitions, in PartitionBaseCase() argument521 if (lhs_batch_partitions == rhs_batch_partitions && in PartitionBaseCase()539 if (lhs_batch_partitions == num_partitions && in PartitionBaseCase()1473 lhs_batch_partitions, output_sharding_dim, in PartitionBaseCase()2295 int64 output_rhs_non_contracting_partitions, int64 lhs_batch_partitions, in LhsIsBestMatchForNonContractingPartitioning() argument2362 assume_lhs_match ? lhs_batch_partitions : rhs_batch_partitions; in LhsIsBestMatchForNonContractingPartitioning()2364 assume_lhs_match ? rhs_batch_partitions : lhs_batch_partitions; in LhsIsBestMatchForNonContractingPartitioning()2472 const int64 lhs_batch_partitions = in PartitionDot() local[all …]