//   contracting_axes)
//   3. batch_dimensions = Gather(operand_shape, batch_axes)
//   4. flattened_shape = Concat(batch_dimensions, flattened_out_dim,
//   flattened_contracting_dim)
// The flattened shape for LHS
// is like [batch_dimensions, flattened_out_dimension,
// flattened_contracting_dimension] and [batch_dimensions,
// flattened_contracting_dimension, flattened_out_dimension] for RHS.

//   contracting_axes)
//   3. batch_dimensions = Gather(operand_shape, batch_axes)
//   4. flattened_shape = Concat(batch_dimensions, flattened_out_dim,
//   flattened_contracting_dim)
// The flattened shape for LHS
// is like [batch_dimensions, flattened_out_dimension,
// flattened_contracting_dimension] and [batch_dimensions,
// flattened_contracting_dimension, flattened_out_dimension] for RHS.

        llvm::ArrayRef({op.getAdjY() ? rank - 1 : rank - 2}));
    auto dimension_numbers = DotDimensionNumbersAttr::get(
        rewriter.getContext(),
        /*lhs_batching_dimensions=*/batch_dimensions,
        /*rhs_batching_dimensions=*/batch_dimensions,
        /*lhs_contracting_dimensions=*/lhs_contracting_dimensions,
        /*rhs_contracting_dimensions=*/rhs_contracting_dimensions);

  (CreateEinsumOpFromXlaDotV2Op $lhs, $rhs, $dot, $dot_dimension_numbers),
  [(IsPrecisionEmpty $precision_config)]>;

// Only handles the case where batch_dimension is empty.
def IsXlaGatherWithoutBatch :
  Constraint<CPred<"IsXlaGatherWithoutBatch($0, $1)">>;

// Create Slice op from XlaGather op without batch dimension.