// Tests that the dot_general accepts i8 tensors and outputs an i32 tensor.
// Note: Argument quantization sequence omitted.
// CHECK: stablehlo.dot_general %{{.+}}, %{{.+}}, contracting_dims = [1] x [0] : (tensor<1x1024xi8>, tensor<1024x3xi8>) -> tensor<1x3xi32>

// Note: Result dequantization sequence omitted.
// CHECK: return %{{.+}} : tensor<1x3xf32>

// -----

    %4 = "tf.Cast"(%3) {Truncate = false, device = ""} : (tensor<1x1024xf32>) -> tensor<1x1024xi8>
    %5 = "tf.Cast"(%4) {Truncate = false, device = ""} : (tensor<1x1024xi8>) -> tensor<1x1024xi32>
    %6 = "tf.Sub"(%5, %cst_4) {device = ""} : (tensor<1x1024xi32>, tensor<i32>) -> tensor<1x1024xi32>
    %7 = "tf.Identity"(%cst_2) {device = ""} : (tensor<1024x3xi8>) -> tensor<1024x3xi8>

    %cst = "tf.Const"() {value = dense<127> : tensor<2x1024xi8>} : () -> tensor<2x1024xi8>
    %cst_0 = "tf.Const"() {value = dense<0.0157480314> : tensor<f32>} : () -> tensor<f32>
    %0 = "tf.Identity"(%cst) : (tensor<2x1024xi8>) -> tensor<2x1024xi8>
    %1 = "tf.Cast"(%0) {Truncate = false} : (tensor<2x1024xi8>) -> tensor<2x1024xf32>

// CHECK-LABEL: func @matmul
// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<2x1024xi8>
// CHECK: %[[PRESERVE_W:.*]] = "tf.Identity"(%[[W]]) : (tensor<2x1024xi8>) -> tensor<2x1024xi8>
// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<2x1024xi8>) -> tensor<2x1024xf32>

// CHECK: %[[quant:.*]] = "tfr.cast"(%arg0) : (tensor<1x10xi8>) -> !tfr.tensor
// CHECK: %[[scale_cast:.*]] = "tfr.cast"(%[[scale]])
// CHECK: %[[zp_cast:.*]] = "tfr.cast"(%[[zp]])
// CHECK: %[[requant:.*]] = tfr.call @tf__requantize(%[[quant]], %[[scale_cast]], %[[zp_cast]])
// CHECK: %[[result:.*]] = "tfr.cast"(%[[requant]])
// CHECK-NOT: quant.uniform
// CHECK: return %[[result]], %[[float_resunt:.*]] : tensor<1x10xi8>, tensor<1x5xf32>
}

}
// CHECK-LABEL: dot_general_srq
// CHECK-SAME: (%[[ARG_1:.+]]: tensor<1x1024x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>
// CHECK-NOT: stablehlo.dot_general

// CHECK-SAME:                         %[[ARG_1:.*]]: tensor<?x?x1xi32>) -> tensor<?x?x1024xi8> {
// CHECK:            %[[VAL_0:.*]] = "tf.GatherNd"(%[[ARG_0]], %[[ARG_1]]) <{bad_indices_policy = ""}> : (tensor<256000x1024xi8>, tensor<?x?x1xi32>) -> tensor<?x?x1024xi8>
// CHECK:            return %[[VAL_0]] : tensor<?x?x1024xi8>
// CHECK:         }

  // Move binary op batched RHS before reshape:
  // binary(reshape(lhs), rhs) => reshape(binary(lhs, flatten(rhs)))
  // Pattern targetted here is as follows-
  // [input, lhr, rhs] == [<1x1024x128>, <1x1024x8x16>, <1x1x8x16xf32>]
  // This is valid only when the-
  // 1.last dimension of lhs is equal to the number of elements in constant rhs.
  // 2.Reduded shape of rhs, here <8x16> is equal to last dimensions of lhs.