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

  # CHECK:      func {{@[a-zA-Z_0-9]+}}(
  # CHECK-SAME:   %arg0: tensor<f32> {tf._user_specified_name = "x", tf_saved_model.index_path = [0]},
  # CHECK-SAME:   %arg1: tensor<!tf_type.resource<tensor<f32>>>
  # CHECK-SAME:   %arg2: tensor<!tf_type.resource<tensor<f32>>>
  # CHECK-SAME:   tensor<f32> {tf_saved_model.index_path = []})
  # CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["some_function"]

// TODO(b/316226111): the printer may not guarantee the same order of fields. Rewrite the checks to be less sensitive to proto serialization formats.
// -----
// Non-SPMD: one input and one output
//
// CHECK-LABEL: func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
// CHECK-NEXT:  "tf.IfrtCall"(%arg0)
// CHECK-SAME:       {program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []}
// CHECK-SAME:       (tensor<1x3xf32>) -> tensor<1x3xf32>

// CHECK-SAME: tf_quant.quantized_ops = ["DepthwiseConv2D", "BiasAdd", "Relu"]
// CHECK: func private @quantized_matmul_with_bias_fn
// CHECK: func private @quantized_matmul_with_bias_and_relu_fn
// CHECK: func private @quantized_matmul_with_bias_and_relu6_fn
// CHECK: func private @quantized_matmul_fn
// CHECK-SAME: tf_quant.quantized_ops = ["MatMul"]

# 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: "tf_saved_model.global_tensor"() <{
  // CHECK-SAME: is_mutable
  // CHECK-SAME: }> : () -> ()
  "tf_saved_model.global_tensor"() { is_mutable, sym_name = "v", type = tensor<f32>, value = dense<42.> : tensor<f32> } : () -> ()

  // CHECK: "tf_saved_model.global_tensor"() <{
  // CHECK-SAME: is_mutable
  // CHECK-SAME: }> : () -> ()

// tf_device.Cluster.

// CHECK-LABEL: func @testReadVariableOpColocated
// CHECK-SAME: (%[[ARG0:.*]]: tensor<*x!tf_type.resource<tensor<4xf32>>>)
func.func @testReadVariableOpColocated(%arg0: tensor<*x!tf_type.resource<tensor<4xf32>>>) {
  // CHECK:      tf_device.replicate
  // CHECK-SAME: (%[[ARG0]] as %[[RI_0:[a-z0-9]*]]: tensor<*x!tf_type.resource<tensor<4xf32>>>)

	netns1 := p.UpsertPodCacheWithNetns(string(pod.UID), wl)
	if !reflect.DeepEqual(netns1, ns) {
		t.Fatalf("Expected the same Netns for the same uid, got %v and %v", netns1, ns)
	}

	ns2 := newFakeNsInode(inc(), 1)
	wl2 := WorkloadInfo{
		Workload: podToWorkload(pod),
		Netns:    ns2,
	}
	// when using same uid, the original netns should be returned
	netns2 := p.UpsertPodCacheWithNetns(string(pod.UID), wl2)
	if netns2 != ns {

// CHECK: return %[[CALL]]

// CHECK: quantized_dot_general_fn
// CHECK-SAME: (%[[ARG1:.+]]: tensor<1x2xf32>,  %[[ARG2:.+]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<1x3xf32>
// CHECK: %[[DOT:.+]] = stablehlo.dot_general %[[ARG1]], %[[ARG2]]
// CHECK-SAME: (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<1x3xf32>

// 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".*}}