message: [XLA::GPU] Explicitly use {lhs,rhs}_contracting_dimension from DotDimensionNumbers as the reduction dimension for Dot.
diff: @@ -495,18 +495,10 @@ Status IrEmitter::HandleDot(HloInstruction* dot) {
  TF_RET_CHECK(!ShapeUtil::IsScalar(lhs_shape) &&
  !ShapeUtil::IsScalar(rhs_shape));
  - // Reduce along the last dimension of the LHS and the second-to-last dimension
  - // of the RHS. Vectors are a special case where the reduction dimension is 0
  - // for both LHS and RHS. This results in a vector dot product producing a
  - // scalar.
  - const int64 lhs_reduction_dimension =
  - ShapeUtil::GetDimensionNumber(lhs_shape, -1);
  - const int64 rhs_reduction_dimension =
  - ShapeUtil::Rank(rhs_shape) >= 2 + dnums.lhs_batch_dimensions_size()
  - ? ShapeUtil::GetDimensionNumber(rhs_shape, -2)
  - : dnums.lhs_batch_dimensions_size();
  -
  - // Check that the batch dims don't cover the last two dims.
  + const int64 lhs_reduction_dimension = dnums.lhs_contracting_dimensions(0);
  + const int64 rhs_reduction_dimension = dnums.rhs_contracting_dimensions(0);
  +
  + // Check that the batch dims don't cover the reduction dimensions.
  for (int64 batch_dim : dnums.lhs_batch_dimensions()) {
  CHECK_NE(lhs_reduction_dimension, batch_dim);
  CHECK_NE(rhs_reduction_dimension, batch_dim);
  @@ -514,7 +506,11 @@ Status IrEmitter::HandleDot(HloInstruction* dot) {
  // Verify the reduction dimension in the two operands are the same size.
  TF_RET_CHECK(lhs_shape.dimensions(lhs_reduction_dimension) ==
  - rhs_shape.dimensions(rhs_reduction_dimension));
  + rhs_shape.dimensions(rhs_reduction_dimension))
  + << "lhs_shape.dimensions(" << lhs_reduction_dimension
  + << ") = " << lhs_shape.dimensions(lhs_reduction_dimension)
  + << ", and rhs_shape.dimensions(" << rhs_reduction_dimension
  + << ") = " << rhs_shape.dimensions(rhs_reduction_dimension);
  // Create loop nests which loop through the LHS operand dimensions and the RHS
  // operand dimensions. The reduction dimension of the LHS and RHS are handled
  