}

  SmallVector<int64_t, 4> new_rhs_contracting_dims;
  new_rhs_contracting_dims.reserve(
      dim_nums.getRhsContractingDimensions().size());
  for (int64_t d : dim_nums.getRhsContractingDimensions()) {
    new_rhs_contracting_dims.push_back(d + 1);
  }

  rewriter.replaceOpWithNewOp<mhlo::DotGeneralOp>(
      dot, dot.getType(), dot.getLhs(), reshape.getOperand(),

      const ::mlir::stablehlo::DotDimensionNumbersAttr dimension_numbers =
          dot.getDotDimensionNumbers();
      ArrayRef<int64_t> rhs_contracting_dims =
          dimension_numbers.getRhsContractingDimensions();
      ArrayRef<int64_t> rhs_batching_dims =
          dimension_numbers.getRhsBatchingDimensions();
      int64_t rank = dot.getRhs().getType().cast<TensorType>().getRank();
      for (int i = 0; i < rank; ++i) {

      dot_dimension_numbers.getLhsContractingDimensions());
  DotDimensionsInfo rhs_dot_dimensions_info(
      rhs_type, dot_dimension_numbers.getRhsBatchingDimensions(),
      dot_dimension_numbers.getRhsContractingDimensions());

  // Transposes lhs shape to be in the order of {batch_dimensions,
  // out_dimensions, contracting dimensions}.
  llvm::SmallVector<int64_t, 4> lhs_permutation = Concat<int64_t>(