// CHECK: %[[CALL:.+]] = "tf.XlaCallModule"(%[[ARG0]], %[[ARG1]])
    // CHECK-SAME: _entry_function = @composite_dot_general_fn_1
    // CHECK-SAME: _original_entry_function = "composite_dot_general_fn_1"
    // CHECK: return %[[CALL]]
  }

  // CHECK: @composite_dot_general_fn_1
  // CHECK-SAME: %[[ARG2:[^:[:space:]]+]]
  // CHECK-SAME: %[[ARG3:[^:[:space:]]+]]

  }
}

// CHECK-LABEL: func @conv2d_unmatching_unit
// CHECK: "tf.PartitionedCall"
// CHECK-SAME: f = @composite_conv2d_fn_1
// Check that the `_tfl_quant_trait` attribute exists since the unit is not in `unit_wise_quantization_specs`.
// CHECK-SAME: _tfl_quant_trait = "fully_quantizable"
// CHECK-SAME: loc(callsite("Model/conv2d@conv2d_unmatching_unit"("Conv2D") at "QuantizationUnit({{.*}})"))

// -----

// disabled by having `_quantization_method = "no_quantization { }"`.
// DISABLE-ALL-DOT-GENERAL-SAME: _entry_function = @composite_dot_general_fn_1
// DISABLE-ALL-DOT-GENERAL-SAME: _original_entry_function
// DISABLE-ALL-DOT-GENERAL-SAME: _quantization_method = "no_quantization { }"
// DISABLE-ALL-DOT-GENERAL-SAME: _tfl_quant_trait = "fully_quantizable"

// DISABLE-ALL-DOT-GENERAL: return %[[XLA_CALL_MODULE:.+]] : tensor<1x1x64xf32>

// CHECK: "tf.PartitionedCall"
// CHECK-SAME: f = @composite_depthwise_conv2d_with_bias_and_relu6_fn_1
// Check that the `_tfl_quant_trait` attribute has been removed.
// CHECK-NOT: _tfl_quant_trait = "fully_quantizable"

// CHECK-LABEL: private @composite_depthwise_conv2d_with_bias_and_relu6_fn_1
// CHECK: %[[DEPTHWISECONV2D_0:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %arg1)

# it is being invoked.
# CHECK: module
# CHECK-SAME: tf.versions
# CHECK-SAME: bad_consumers
# CHECK-SAME: min_consumer
# CHECK-SAME: producer

# CHECK: "tf_saved_model.global_tensor"() <{is_mutable, sym_name = "[[VAR:[a-zA-Z_0-9]+]]", type = tensor<1x3xf32>, value = {{.*}} : tensor<1x3xf32>}> : () -> ()
# CHECK-NOT: session_initializer

# CHECK:      func {{@[a-zA-Z_0-9]+}}(

// CHECK: func.func @init_func_restore_op
// CHECK-SAME: tf_saved_model.exported_names = ["tf_saved_model.session_initializer_restore_op"]
// CHECK-SAME: tf_saved_model.initializer_type = "restore_op"

// Test that a new save function that wraps the SaveV2 op is created.
// CHECK: func.func private @tf_quant__save(%[[ARG:.*]]: tensor<!tf_type.string>)
// CHECK: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"()
// CHECK-SAME: {{.*shared_name = "var_0".*}}

  // CHECK: "tf_saved_model.session_initializer"()
  // CHECK-SAME: initializers = [@init_func_restore_op]

  func.func @init_func_restore_op() -> () attributes {
      tf_saved_model.initializer_type = "restore_op",
      tf_saved_model.exported_names = ["tf_saved_model.session_initializer_restore_op"]} {
    return
  }
  // CHECK: @init_func_restore_op
  // CHECK-SAME: tf_saved_model.initializer_type = "restore_op"

  func.return %1 : tensor<4x128x1xf32>
  // CHECK-NEXT: %[[CONST_WEIGHT:.*]] = arith.constant
  // CHECK-SAME: [1.000000e+00, 2.000000e+00]
  // CHECK-SAME: tensor<1x2xf32>
  // CHECK: %[[FC_RES:.*]] = "tfl.fully_connected"(%arg0, %[[CONST_WEIGHT]]
  // CHECK-SAME: <{fused_activation_function = "NONE", keep_num_dims = true, weights_format = "DEFAULT"}> : (tensor<4x128x2xf32>, tensor<1x2xf32>, none) -> tensor<4x128x1xf32>

  }
}

// CHECK-LABEL: quantize_dot_general_fn
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x2xf32>
// CHECK: %[[CST:.+]] = stablehlo.constant dense<3.000000e-01> : tensor<2x3xf32>
// CHECK: %[[Q:.+]] = "quantfork.qcast"(%[[CST]]) : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>
// CHECK: %[[CALL:.+]] = call @quantized_dot_general_fn(%[[ARG0]], %[[Q]])

// CHECK-NEXT:   %[[RESULT_0:.*]] = "tf.While"(%[[ARG_0]])
// CHECK-SAME:   body = @while_body
// CHECK-SAME:   cond = @while_cond
// CHECK-SAME:   device = "/job:localhost/replica:0/task:0/device:CPU:0"
// CHECK-NEXT:   %[[RESULT_1:.*]] = tf_device.remote_run "/job:worker/replica:0/task:1" @[[MAIN_PARTITION_0:.*]](%[[ARG_1]])