// Test that VarHandleOp is not marked as "potentially written" if it is not
// assigned inside the function called by "tf.BatchFunction".

module {
// CHECK-LABEL: @serving_default
  func.func @serving_default() -> (tensor<*xi32>) {
    // expected-remark@below {{device: "", container: "", shared_name: "var_0", is_potentially_written: false}}

  func.func @serving_default(%arg0: tensor<1x4xf32> {tf_saved_model.index_path = ["x"]}) -> (tensor<1x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_x:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
    %cst = "tf.Const"() <{value = dense<1.000000e+01> : tensor<f32>}> {device = ""} : () -> tensor<f32>

// RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-insert-calibration-statistics-saver='aggregator-ops-to-ignore=skipping_id' | FileCheck %s

func.func @serving_default(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x2x2x2xf32>) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}} {

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

          converted_model2.signatures._signatures.keys(), {'serving_default'}
      )

      # Expect two models to produce the same results.
      test_data = ops.convert_to_tensor(
          np.random.uniform(low=0, high=64, size=(32)).astype(
              dtypes.int32.as_numpy_dtype
          )
      )
      new_outputs_1 = converted_model1.signatures['serving_default'](
          input_tensor=test_data

// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-rewrite-ifrt-load-variable %s | FileCheck %s

// Variable is used by both CPU and TPU
//
// CHECK-LABEL: func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    [[HANDLE:%.*]] = "tf.VarHandleOp"()
// CHECK-NEXT:    [[ARRAYKEY:%.*]], [[FURTURE:%.*]] = "tf_mlrt.tf_ifrt_load_variable"([[HANDLE]])

  func.func @serving_default(%arg0: tensor<1x20x20x28xf32> {tf_saved_model.index_path = ["a"]}) -> (tensor<1x40x40x28xf32> {tf_saved_model.index_path = ["b"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "c:0", outputs = "d:0"}, tf_saved_model.exported_names = ["serving_default"]} {
      %c = stablehlo.constant dense<40> : tensor<2xi32>

  func.func @serving_default(%arg0: tensor<1x32x32x128xf32> {tf_saved_model.index_path = ["a"]}) -> (tensor<1x64x64x128xf32> {tf_saved_model.index_path = ["b"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "c:0", outputs = "d:0"}, tf_saved_model.exported_names = ["serving_default"]} {
      %0  = "tf.Const"() {value = dense<[56, 904]> : tensor<2xi32>} : () -> tensor<2xi32>

  return %0 : tensor<1xf32>
}
// The contents of `@serving_default` should have been inlined to `@batch_func`.
// CHECK: func.func @serving_default(%[[ARG0:.*]]: tensor<1xf32>, %[[ARG1:.*]]: tensor<1xf32>) -> tensor<1xf32>
// CHECK-NOT: tf.BatchFunction
// CHECK: %[[ADD0:.*]] = "tf.AddV2"(%[[ARG0]], %[[ARG1]])
// CHECK: return %[[ADD0]] : tensor<1xf32>

    // If the signature_key has not been explicitly specified, use the default
    // value of "serving_default".
    const std::string signature_key = dataset_config.has_signature_key()
                                          ? dataset_config.signature_key()
                                          : "serving_default";
    if (repr_dataset_file_map.contains(signature_key)) {
      return absl::InvalidArgumentError(