// CHECK-LABEL: private @composite_matmul_with_bias_and_relu6_fn_1
// CHECK-NEXT: %[[matmul:.*]] = "tf.MatMul"(%arg0, %arg1)
// CHECK-SAME: attr_map = "0:transpose_a,1:transpose_b"
// CHECK-NEXT: tf.BiasAdd
// CHECK-NEXT: tf.Relu6
// CHECK-NEXT: return

// CHECK-LABEL: private @composite_matmul_with_bias_and_relu_fn_1
// CHECK-NEXT: tf.MatMul"(%arg0, %arg1)
// CHECK-SAME: attr_map = "0:transpose_a,1:transpose_b"

module {
  // TODO(b/260020937): Support transpose_a, transpose_b for matmul.
  func.func @matmul(%arg0: tensor<2x12xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<12x2xf32>} : () -> tensor<12x2xf32>

  %cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<512x512xf32>} : () -> tensor<512x512xf32>
  %out_1 = "tf.MatMul"(%arg0, %cst) {
    device = "", transpose_a = false, transpose_b = false
  } : (tensor<1x12x12x512xf32>, tensor<512x512xf32>) -> tensor<*xf32>
  %out_2 = "tf.MatMul"(%arg0, %arg0) {
    device = "", transpose_a = false, transpose_b = true
  } : (tensor<1x12x12x512xf32>, tensor<1x12x12x512xf32>) -> tensor<*xf32>

    (TF_SubOp $beta, (TF_MulOp $m, $mul)))>;

class TFi32<int v> : ConstantAttr<I32ElementsAttr, !cast<string>(v)>;

// Matmul without transpose on b to matmul with explicit transpose op and
// transposed b.
def ConvertMatmulWithoutTransposeToWithTranspose :
      Pat<(TF_MatMulOp $a, $b, ConstBoolAttrFalse:$at, ConstBoolAttrFalse, $grad_a, $grad_b),

  %3 = "tf.Identity"(%2) : (tensor<*xf32>) -> tensor<*xf32>
  func.return %3 : tensor<*xf32>
}

//===----------------------------------------------------------------------===//
// MatMul + BiasAdd + <Activation> fusions.
//===----------------------------------------------------------------------===//

// CHECK-LABEL: matmulBiasAdd

  auto matmul = rewriter.create<TFL::BatchMatMulOp>(
      loc, RankedTensorType::get(matmul_shape, result_type.getElementType()),
      lhs_flattend, rhs_flattend, /*adj_x*/ false_attr, /*adj_y*/ false_attr,
      /*asym_quant_input*/ false_attr);
  if (result_type.hasStaticShape()) {
    auto reshaped =
        rewriter.create<mhlo::ReshapeOp>(loc, result_type, matmul.getResult());
    return reshaped.getResult();

  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  TFE_Op* matmul = MatMulOp(ctx, h0_task1, h1_task1);
  TFE_OpSetDevice(matmul, remote_device_name, status);
  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  TFE_TensorHandle* retvals[1];
  int num_retvals = 1;
  TFE_Execute(matmul, &retvals[0], &num_retvals, status);
  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

///     int idx = 3;
///     auto b = Variable(linear.WithOpName("b_", idx),
///                       {2}, DT_FLOAT);
///     auto x = Const(linear, {...});  // name: "linear/Const"
///     auto m = MatMul(linear, x, W);  // name: "linear/MatMul"
///     auto r = BiasAdd(linear, m, b); // name: "linear/BiasAdd"
///
/// Scope lifetime:
///
/// A new scope is created by calling Scope::NewRootScope. This creates some

// RUN: tf-quant-opt %s -split-input-file -quant-add-dump-tensor-op='debugger_type=float_per_layer' | FileCheck --check-prefix=FloatPerLayer %s

module {
  func.func @matmul2(%arg0: tensor<?x2xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<?x2xf32>) {
    %0 = stablehlo.constant dense<[-0.211145893, -0.708605706]> : tensor<2xf32>

  %dq_weight = "quantfork.dcast"(%q_weight) : (tensor<144x12x!quant.uniform<i8:f32, 0.074855112561992565:-1>>) -> tensor<144x12xf32>
  %9 = "tf.MatMul"(%7, %dq_weight) {transpose_a = false, transpose_b = false} : (tensor<*xf32>, tensor<144x12xf32>) -> tensor<*xf32>
  %10 = "quantfork.qcast"(%9) {volatile} : (tensor<*xf32>) -> tensor<*x!quant.uniform<i8:f32, 4.000000e-03:-12>>