# CHECK-NEXT: [[R1:%.*]] = "tf.MatMul"([[ARG0]], [[R0]]) <{{{.*}}}> {{{.*}}} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
# CHECK-NEXT: return [[R1]] : tensor<3x3xf32>


def Test():

  x = tf.constant([[1.0], [1.0], [1.0]])
  y = tf.compat.v1.get_variable(
      name='y',
      shape=(1, 3),

  func.func @matmul(%arg0: tensor<2x12xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<12x2xf32>} : () -> tensor<12x2xf32>
    %1 = "tf.PartitionedCall"(%arg0, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn_1} : (tensor<2x12xf32>, tensor<12x2xf32>) -> tensor<*xf32>
    func.return %1: tensor<*xf32>
  }

  %11 = "mhlo.dot"(%10, %5) : (tensor<32x3920xf32>, tensor<3920x10xf32>) -> tensor<32x10xf32>
  %12 = "mhlo.broadcast_in_dim"(%2) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<10xf32>) -> tensor<32x10xf32>
  %13 = mhlo.add %11, %12 : tensor<32x10xf32>
  %14 = mhlo.maximum %13, %0 : tensor<32x10xf32>
  return %14 : tensor<32x10xf32>
}
// CHECK: ReadVariableOp
// CHECK: mhlo.convolution
// CHECK: mhlo.add

func.func @softmax(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
  %0 = "tf.Softmax"(%arg0) : (tensor<8x16xf32>) -> tensor<8x16xf32>
  func.return %0 : tensor<8x16xf32>
// CHECK: %[[CUSTOM_0:.*]] = "tfl.custom"(%arg0) <{custom_code = "FlexSoftmax", custom_option = #tfl<const_bytes : "0x07536F66746D617800161207536F66746D61781A002A070A0154120230013200000221191414042801">}> : (tensor<8x16xf32>) -> tensor<8x16xf32>

  func.func @matmul(%arg0: tensor<2x12xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<12x2xf32>} : () -> tensor<12x2xf32>
    %1 = "tf.PartitionedCall"(%arg0, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn_1} : (tensor<2x12xf32>, tensor<12x2xf32>) -> tensor<*xf32>
    func.return %1: tensor<*xf32>
  }

// -----

// CHECK-LABEL: transpose_simple_2d
func.func @transpose_simple_2d() -> tensor<3x2xf32> {
  %0 = stablehlo.constant dense<[[0.000000e+0, 1.000000e+0, 2.000000e+0], [3.000000e+0, 4.000000e+0, 5.000000e+0]]> : tensor<2x3xf32>
  %1 = stablehlo.transpose %0, dims = [1, 0] : (tensor<2x3xf32>) -> tensor<3x2xf32>
  return %1 : tensor<3x2xf32>
}

  # Verify that the function defined using function.Defun
  # has a corresponding tf.LegacyCall op.
  # CHECK:      func {{@[a-zA-Z_0-9]+}}(
  # CHECK-SAME: [[ARG0:%.*]]: tensor<3x1xf32> {tf_saved_model.index_path = ["y"]},
  # CHECK-SAME: [[ARG1:%.*]]: tensor<3x1xf32> {tf_saved_model.index_path = ["x"]}
  #
  # CHECK-NEXT: [[R0:%.*]] = "tf.LegacyCall"([[ARG1]], [[ARG0]])
  z = plus(x, y)

  %2 = "tf.MatMul"(%arg0, %1) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
  func.return %2 : tensor<3x3xf32>
}

// CHECK-LABEL: func @string_type
func.func @string_type(%arg: tensor<1x2x!tf_type.string>) -> tensor<?x6x!tf_type.string> {

  // CHECK-NEXT: }
  // CHECK-EMPTY:

  %0 = "tfl.pseudo_const" () {value = dense<[[-5.0, 10000.0], [150.0, 10.0], [5.0, 100.0]]> : tensor<3x2xf32>} : () -> tensor<3x2xf32> loc("Const")
  %1 = "tfl.bucketize"(%0) {boundaries = [0.0 : f32, 10.0 : f32, 100.0 : f32]} : (tensor<3x2xf32>) -> tensor<3x2xi32> loc("bucketize")
  func.return %1 : tensor<3x2xi32>

    %2 = "tf.Reshape"(%1, %cst_0) : (tensor<1x10xf32>, tensor<2xi32>) -> tensor<?x10xf32> loc(fused["Reshape:", "model/reshape"])
    %3 = "tf.AddV2"(%2, %cst) : (tensor<?x10xf32>, tensor<10xf32>) -> tensor<?x10xf32> loc(fused["AddV2:", "model/add"])
    func.return %3 : tensor<?x10xf32>
  }
}

// CHECK-LABEL: func @matmul_with_reshape_disable_quantization
// CHECK: "tf.PartitionedCall"