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

// CHECK-NEXT: %[[V0:.*]] = "tf.PartitionedCall"(%[[ARG1]], %[[CST_1]], %[[CST_0]]) <{config = "", config_proto = "", executor_type = "", f = @quantize_i8}> : (tensor<2x3x3x2xf32>, tensor<f32>, tensor<i32>) -> tensor<2x3x3x2xi8>

// CHECK-NOT: f = @NoOp

      // CHECK: %[[CALL_OUTPUT_0:[0-9]*]] = "tf.PartitionedCall"(%[[A_OUTPUT]])
      // CHECK-SAME: f = @_func
      // CHECK-SAME: device = "/device:test_device:0"
      %3 = "tf_device.launch_func"(%2) {device = "/device:test_device:0", func = @_func} : (tensor<?xf32>) -> tensor<?xf32>

      // CHECK: %[[CALL_OUTPUT_1:[0-9]*]] = "tf.PartitionedCall"(%[[CALL_OUTPUT_0]])
      // CHECK-SAME: f = @_func

// CHECK-LABEL: @func0
func.func @func0(%arg0 : tensor<i1>) -> tensor<f32> {
  %0 = tf_executor.graph {
// CHECK: island
// CHECK: PartitionedCall
// CHECK-SAME: @_tpu_v1_compat_outlined::@_tpu_v1_compat_outlined_func0
    %1:2 = tf_executor.island {

// 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>

    %cst = "tf.Const"() {device = "", value = dense<[7.72826624, 8.8264122, 3.64885974]> : tensor<3xf32>} : () -> tensor<3xf32>

    %i8_max = "tf.Const"() {value = dense<127> : tensor<i8>} : () -> tensor<i8>

    %clip_min = "tf.PartitionedCall"(%i8_min, %out_scale, %out_zp) {
        config = "", config_proto = "", executor_type = "", f=@internal_dequantize_i8_fn
      } : (tensor<i8>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
    %clip_max = "tf.PartitionedCall"(%i8_max, %out_scale, %out_zp) {

    func.return
  }
}

// -----
// Test stateful and stateless partitioned calls.
// CHECK-LABEL: func @f
func.func @f() {
  // CHECK: "tf.PartitionedCall"() <{config = "",  config_proto = "", executor_type = "", f = @g}> : () -> ()
  "tf.PartitionedCall"() {config = "",  config_proto = "", executor_type = "", f = @g} : () -> ()

    %val = "tf.PartitionedCall"(%handle) {config = "", config_proto = "", executor_type = "", f = @f_callee} : (tensor<!tf_type.resource<tensor<0xf32>>>) -> (tensor<0xf32>)
    func.return %val : tensor<0xf32>
  }

  // CHECK: func private @f_callee() -> tensor<0xf32>
  func.func private @f_callee(%arg0: tensor<*x!tf_type.resource>) -> tensor<0xf32> {