Christian Sigg <******@****.***> 1714640622 -0700

      "/job:localhost/replica:0/task:0/device:CPU:0");
  Output a = ops::Placeholder(root.WithOpName("A"), DT_FLOAT);
  for (int i = 0; i < 5; i++) {
    a = ops::MatMul(root.WithOpName(absl::StrCat("matmul_", i)), a, a);
    a = ops::Add(root.WithOpName(absl::StrCat("add_", i)), a, a);
  }

  GraphDef graph_def;
  root.graph()->ToGraphDef(&graph_def);
  return graph_def;
}

# RUN: tf_tfl_translate -unfold_batchmatmul=false -tf-input-arrays=Placeholder,Placeholder_1 -tf-input-shapes=2,5,3:3,7 -tf-input-data-types=DT_FLOAT,DT_FLOAT -tf-output-arrays=MatMul -output-mlir %s -o - 2>&1 | FileCheck %s

node {
  name: "Placeholder"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: 2

// RUN: tf-quant-opt %s -split-input-file -quant-lift-quantizable-spots-as-functions -quant-prepare-quantize-drq -quant-quantize='weight-quantization=true' -verify-each=false | FileCheck %s

// -----

module {
  func.func @matmul(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2x1024xf32>} : () -> tensor<2x1024xf32>

# RUN: tf_tfl_translate -tf-input-arrays=Placeholder,Placeholder_1 -tf-input-shapes=2,5,3:3,7 -tf-input-data-types=DT_FLOAT,DT_FLOAT -tf-output-arrays=MatMul -unfold_batchmatmul=true -output-mlir %s -o - 2>&1 | FileCheck %s

node {
  name: "Placeholder"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: 2
        }

  // CHECK: tfrt_fallback_async.executeop key(2) cost({{.*}}) device("/device:CPU:0") "tf.MatMul"
  %0 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
  %1 = "tf.MatMul"(%arg0, %0) {T = f32, device = "/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>

// CHECK-NEXT:    [[TENSOR:%.*]] = "tf_mlrt.tf_await"([[FURTURE]]) : (!mlrt.future) -> tensor<3x1xf32>
// CHECK-NEXT:    "tf.MatMul"(%arg0, [[TENSOR]]) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>

// CHECK:  "tf.VarHandleOp"
// CHECK-NOT:  [[VARIABLE:%.*]] = "tf.ReadVariableOp"
// CHECK-NEXT:  [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"
// CHECK-SAME:    used_by_host = true
// CHECK-NEXT:  [[MATRES:%.*]] = "tf.MatMul"(%arg0, [[FUTURE]])
// CHECK-NEXT:   [[RES:%.*]] = "tf.IfrtCall"(%arg0, [[KEY]]) <{program_id = 6515870160938153680 : i64, variable_arg_indices = [1 : i32]}>
// CHECK-NEXT:    return [[RES]], [[MATRES]] : tensor<1x1xf32>, tensor<1x1xf32>

// CHECK:    return
//
// CHECK:  func.func @_ifrt_program__func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> tensor<1x1xf32>
// CHECK-SAME:      tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
// CHECK-NEXT:     %0 = "tf.MatMul"(%arg0, %arg1)
// CHECK:          return

    %arg1: tensor<1x3xf32> {tf_saved_model.index_path = [0]})
      -> (tensor<3x3xf32> {tf_saved_model.index_path = []}) {
  // CHECK: {{%.*}} = corert.get_op_handler %arg0 "/device:GPU:0"
  %2 = "tf.MatMul"(%arg0, %arg1) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "/device:GPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>