dot_dimension_numbers = #mhlo.dot<
      lhs_contracting_dimensions = [2],
      rhs_contracting_dimensions = [0]
  >} : (tensor<1x1x512xf32>, tensor<512x13xf32>) -> tensor<1x1x13xf32>
  %1 = "mhlo.dot_general"(%arg1, %arg3) {
    dot_dimension_numbers = #mhlo.dot<
      lhs_contracting_dimensions = [2],
      rhs_contracting_dimensions = [0]
  >} : (tensor<2x1x512xf32>, tensor<512x13xf32>) -> tensor<2x1x13xf32>

// CHECK:           "mhlo.dot_general"({{.*}}, {{.*}}) <{
// CHECK-SAME:        lhs_batching_dimensions = [0]
// CHECK-SAME:        rhs_batching_dimensions = [0]
// CHECK-SAME:        lhs_contracting_dimensions = [2]
// CHECK-SAME:        rhs_contracting_dimensions = [1]
  %0 = "tf.BatchMatMulV2"(%arg0, %arg1) {T = f32, adj_x = false, adj_y = false, device = ""} : (tensor<3x4x2xf32>, tensor<2x4xf32>) -> tensor<3x4x4xf32>
  func.return %0 : tensor<3x4x4xf32>

  // CHECK: return %[[result]]
  %0 = "mhlo.dot_general" (%lhs, %rhs) {
    dot_dimension_numbers = #mhlo.dot<
      lhs_batching_dimensions = [],
      rhs_batching_dimensions = [],
      lhs_contracting_dimensions = [2],
      rhs_contracting_dimensions = [0]
    >} : (
      tensor<?x?x6xi8>, tensor<6x8xi8>
    ) -> tensor<?x?x8xi32>
  %1 = chlo.broadcast_add %0, %zp_offset : (
      tensor<?x?x8xi32>, tensor<8xi32>) -> tensor<?x?x8xi32>

  %0 = "mhlo.dot_general"(%arg0, %arg1) {
    dot_dimension_numbers = #mhlo.dot<
      lhs_batching_dimensions = [0],
      lhs_contracting_dimensions = [1, 2],
      rhs_batching_dimensions = [0],
      rhs_contracting_dimensions = [1, 3]
    >,
    precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]
  } : (tensor<3x2x6x5x1xf32>, tensor<3x2x4x6xf32>) -> tensor<3x5x1x4xf32>

    %0 = "stablehlo.dot_general"(%arg0, %arg1) {
    dot_dimension_numbers = #stablehlo.dot<
      lhs_batching_dimensions = [],
      rhs_batching_dimensions = [],
      lhs_contracting_dimensions = [1],
      rhs_contracting_dimensions = [0]
      >
    } : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
    return %0 : tensor<2x2xf32>
  }
}

// -----

  //   - rank(lhs) is 1 or 2
  //   - rank(rhs) = 2
  //   - size(lhs_contracting_dimensions) = 1
  //   - size(rhs_contracting_dimensions) = 1
  //   - lhs_contracting_dimension = last dimension of lhs.
  //   - `stablehlo.dot_general` should not have `lhs_batching_dim`.
  //   - quantization_dimension(rhs) should not be in
  //     `rhs_contracting_dimensions`.
  // https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dot_general

  lhs_contract_idx.insert(
      dot_dimension_numbers.lhs_contracting_dimensions().begin(),
      dot_dimension_numbers.lhs_contracting_dimensions().end());
  rhs_batch_idx.insert(dot_dimension_numbers.rhs_batch_dimensions().begin(),
                       dot_dimension_numbers.rhs_batch_dimensions().end());
  rhs_contract_idx.insert(
      dot_dimension_numbers.rhs_contracting_dimensions().begin(),

      d2.add_lhs_batch_dimensions(v);
    }
    for (auto v : d1.lhs_contracting_dimensions()) {
      d2.add_lhs_contracting_dimensions(v);
    }
  } else {
    for (auto v : d1.rhs_batch_dimensions()) {
      d2.add_rhs_batch_dimensions(v);
    }
    for (auto v : d1.rhs_contracting_dimensions()) {
      d2.add_rhs_contracting_dimensions(v);
    }
  }
}

  %1 = "stablehlo.dot_general"(%arg0, %0) {
    dot_dimension_numbers = #stablehlo.dot<
      lhs_batching_dimensions = [0, 1],
      rhs_batching_dimensions = [0, 1],
      lhs_contracting_dimensions = [3],
      rhs_contracting_dimensions = [2]>,
      precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]

    attributes.emplace_back(builder.getNamedAttr(
        "lhs_contracting_dimensions",
        BuildVhloTensorV1Attr(
            {static_cast<int64_t>(op->lhs_contracting_dimensions.size())},
            op->lhs_contracting_dimensions, builder)));
    attributes.emplace_back(builder.getNamedAttr(
        "rhs_contracting_dimensions",
        BuildVhloTensorV1Attr(