chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tfrt/tests/ifrt/BUILD

@@ -0,0 +1,21 @@
+load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
+
+glob_lit_tests(
+    name = "all_tests",
+    data = [":test_utilities"],
+    driver = "//tensorflow/compiler/mlir:run_lit.sh",
+    test_file_exts = ["mlir"],
+)
+
+# Bundle together all of the test utilities that are used by tests.
+filegroup(
+    name = "test_utilities",
+    testonly = True,
+    data = [
+        "//tensorflow/compiler/mlir/tfrt:tf-tfrt-opt",
+        "@llvm-project//llvm:FileCheck",
+        "@llvm-project//mlir:run_lit.sh",
+    ],
+)

tensorflow/compiler/mlir/tfrt/tests/ifrt/lower_to_ifrt_restore_variable.mlir

@@ -0,0 +1,322 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -verify-diagnostics -lower-to-ifrt-restore-variable %s | FileCheck %s
+
+
+// -----
+// single variable
+
+// CHECK-LABEL:   func.func @restore_single() {
+// CHECK-NEXT:     [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+// CHECK-NEXT:     [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:     [[NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:     [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-NEXT:     "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLEY]])
+// CHECK-SAME:        {restored_dtypes = [f32], returned_tensor_names = [], truncate_in_cast = array<i1: false>}
+// CHECK-NOT:       "tf.RestoreV2"
+// CHECK-NEXT:     return
+
+module {
+  func.func @restore_single() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+}
+
+// -----
+// single variable: VarHandleOp is before RestoreV2
+
+// CHECK-LABEL:   func.func @varhandle_before_restore() {
+// CHECK-NEXT:     [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+// CHECK-NEXT:     [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:     [[NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:     [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-NEXT:     "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLEY]])
+// CHECK-SAME:        {restored_dtypes = [f32], returned_tensor_names = [], truncate_in_cast = array<i1: false>}
+// CHECK-NOT:       "tf.RestoreV2"
+// CHECK-NEXT:     return
+
+module {
+  func.func @varhandle_before_restore() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+}
+
+
+// -----
+// multiple variables
+
+// CHECK-LABEL:   func.func @restore_multiple() {
+// CHECK-NEXT:     [[PREFIX:%.*]] = "tf.Const"()
+// CHECK-NEXT:     [[SLICE:%.*]] = "tf.Const"()
+// CHECK-NEXT:     [[NAME:%.*]] = "tf.Const"()
+// CHECK-NEXT:     [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-NEXT:     [[HANDLEZ:%.*]] = "tf.VarHandleOp"() <{container = "x", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
+// CHECK-NEXT:     "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLEY]], [[HANDLEZ]])
+// CHECK-SAME:        {restored_dtypes = [f32, f32], returned_tensor_names = [], truncate_in_cast = array<i1: false, false>}
+// CHECK-NOT:       "tf.RestoreV2"
+// CHECK-NEXT:     return
+
+module {
+  func.func @restore_multiple() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<["", ""]> : tensor<2x!tf_type.string>}> : () -> tensor<2x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<["y", "z"]> : tensor<2x!tf_type.string>}> : () -> tensor<2x!tf_type.string>
+    %0:2 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<3x1xf32>, tensor<1x3xf32>)
+    %1 = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %0#0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    %2 = "tf.VarHandleOp"() <{container = "x", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
+    "tf.AssignVariableOp"(%2, %0#1) : (tensor<!tf_type.resource<tensor<1x3xf32>>>, tensor<1x3xf32>) -> ()
+    return
+  }
+}
+
+// -----
+// Restored variable is not assigned with a name is an error.
+
+module {
+  func.func @unassigned_restore_return_error() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<["", ""]> : tensor<2x!tf_type.string>}> : () -> tensor<2x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<["y", "z"]> : tensor<2x!tf_type.string>}> : () -> tensor<2x!tf_type.string>
+    //expected-error@below {{'tf.RestoreV2' op expects 2 valid users, but got 1}}
+    %0:2 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<3x1xf32>, tensor<1x3xf32>)
+    %1 = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %0#0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+}
+
+// -----
+//  Restored tensor is consumed by an op other than AssignVariableOp, it is returned as an output by IfrtRestoreVariableOp.
+
+// CHECK-LABEL: func.func @restore_with_consumer() {
+// CHECK-NEXT:   %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+// CHECK-NEXT:   %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:   %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:   %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-NEXT:   %1 = "tf.IfrtRestoreVariableOp"(%cst, %cst_1, %cst_0, %0) <{restored_dtypes = [f32], returned_tensor_names = ["y"], truncate_in_cast = array<i1: false>}> : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+// CHECK-NEXT:   %2 = "tf.ReluOp"(%1) : (tensor<3x1xf32>) -> tensor<3x1xf32>
+// CHECK-NEXT:   %3 = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-NEXT:   "tf.AssignVariableOp"(%3, %2) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+// CHECK-NOT: "tf.RestoreV2"
+// CHECK-NEXT: return
+
+
+module {
+  func.func @restore_with_consumer() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    %2 = "tf.ReluOp"(%0) : (tensor<3x1xf32>) -> tensor<3x1xf32>
+    %1 = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %2) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+}
+
+
+
+// -----
+// variable with cast
+// CHECK-LABEL:   func.func @restore_with_cast() {
+// CHECK-NEXT:     [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+// CHECK-NEXT:     [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:     [[NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:     [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
+// CHECK-NEXT:     "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLEY]])
+// CHECK-SAME:        {restored_dtypes = [f32], returned_tensor_names = [], truncate_in_cast = array<i1: false>}
+// CHECK-NOT:       "tf.RestoreV2"
+// CHECK-NEXT:     return
+
+module {
+  func.func @restore_with_cast() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    %1 = "tf.Cast"(%0) <{Truncate = false}> : (tensor<3x1xf32>) -> tensor<3x1xbf16>
+    %2 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
+    "tf.AssignVariableOp"(%2, %1) : (tensor<!tf_type.resource<tensor<3x1xbf16>>>, tensor<3x1xbf16>) -> ()
+    return
+  }
+}
+
+// -----
+// variable and table lookup
+// CHECK-LABEL: func.func @restore_var_and_table()
+// CHECK-NEXT:  [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"model/foo"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+// CHECK-NEXT:  [[NAMES:%.*]] = "tf.Const"() <{value = dense<["var1", "table1_keys", "table1_vals"]> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
+// CHECK-NEXT:  [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
+// CHECK-NEXT:  [[VAR1:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "var1"}> : () -> tensor<!tf_type.resource<tensor<*xi32>>>
+// CHECK-NEXT:  [[TABLE1_KEYS:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "table1_keys"}> : () -> tensor<!tf_type.resource<tensor<*xi64>>>
+// CHECK-NEXT:  [[TABLE1_VALS:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "table1_vals"}> : () -> tensor<!tf_type.resource<tensor<*xf32>>>
+// CHECK-NEXT:  [[RETURNED_VALUES:%.*]]:2 = "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAMES]], [[SLICE]], [[VAR1]], [[TABLE1_KEYS]], [[TABLE1_VALS]])
+// CHECK-SAME:     <{restored_dtypes = [i32, i64, f32], returned_tensor_names = ["table1_keys", "table1_vals"], truncate_in_cast = array<i1: false, false, false>}> : (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>, tensor<!tf_type.resource<tensor<*xi32>>>, tensor<!tf_type.resource<tensor<*xi64>>>, tensor<!tf_type.resource<tensor<*xf32>>>) -> (tensor<*xi64>, tensor<*xf32>)
+// CHECK-NEXT:  [[TABLE1:%.*]] = "tf.HashTableV2"() <{container = "", key_dtype = i64, shared_name = "table1", value_dtype = f32}> : () -> tensor<!tf_type.resource>
+// CHECK-NEXT:  "tf.LookupTableImportV2"([[TABLE1]], [[RETURNED_VALUES]]#0, [[RETURNED_VALUES]]#1) : (tensor<!tf_type.resource>, tensor<*xi64>, tensor<*xf32>) -> ()
+// CHECK-NOT:   "tf.RestoreV2"
+// CHECK-NEXT:  return
+module {
+  func.func @restore_var_and_table() {
+    %cst = "tf.Const"() {value = dense<"model/foo"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() {value = dense<["var1", "table1_keys", "table1_vals"]> : tensor<3x!tf_type.string>} : () -> tensor<3x!tf_type.string>
+    %cst_1 = "tf.Const"() {value = dense<["", "", ""]> : tensor<3x!tf_type.string>} : () -> tensor<3x!tf_type.string>
+    %0:3 = "tf.RestoreV2"(%cst, %cst_0, %cst_1) {dtypes = [i32, i64, f32]} : (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>) -> (tensor<*xi32>, tensor<*xi64>, tensor<*xf32>)
+    %1 = "tf.VarHandleOp"() {container = "", shared_name = "var1"} : () -> tensor<!tf_type.resource<tensor<*xi32>>>
+    "tf.AssignVariableOp"(%1, %0#0) : (tensor<!tf_type.resource<tensor<*xi32>>>, tensor<*xi32>) -> ()
+    %2 = "tf.HashTableV2"() {container = "", key_dtype = i64, shared_name = "table1", value_dtype = f32} : () -> tensor<!tf_type.resource>
+    "tf.LookupTableImportV2"(%2, %0#1, %0#2) : (tensor<!tf_type.resource>, tensor<*xi64>, tensor<*xf32>) -> ()
+    return
+  }
+}
+
+// -----
+// variable and dense table lookup
+// CHECK-LABEL: func.func @restore_var_and_dense_table()
+// CHECK-NEXT:  [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"model/foo"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+// CHECK-NEXT:  [[NAMES:%.*]] = "tf.Const"() <{value = dense<["var1", "table1_keys", "table1_vals"]> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
+// CHECK-NEXT:  [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
+// CHECK-NEXT:  [[VAR1:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "var1"}> : () -> tensor<!tf_type.resource<tensor<*xi32>>>
+// CHECK-NEXT:  [[TABLE1_KEYS:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "table1_keys"}> : () -> tensor<!tf_type.resource<tensor<*xi64>>>
+// CHECK-NEXT:  [[TABLE1_VALS:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "table1_vals"}> : () -> tensor<!tf_type.resource<tensor<*xf32>>>
+// CHECK-NEXT:  [[RETURNED_VALUES:%.*]]:2 = "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAMES]], [[SLICE]], [[VAR1]], [[TABLE1_KEYS]], [[TABLE1_VALS]])
+// CHECK-SAME:     <{restored_dtypes = [i32, i64, f32], returned_tensor_names = ["table1_keys", "table1_vals"], truncate_in_cast = array<i1: false, false, false>}> : (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>, tensor<!tf_type.resource<tensor<*xi32>>>, tensor<!tf_type.resource<tensor<*xi64>>>, tensor<!tf_type.resource<tensor<*xf32>>>) -> (tensor<*xi64>, tensor<*xf32>)
+// CHECK-NEXT:  [[EMPTY_KEY:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<i64>}> : () -> tensor<i64>
+// CHECK-NEXT:  [[DEFAULT_VALUE:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<f32>}> : () -> tensor<f32>
+// CHECK-NEXT:  [[TABLE1:%.*]] = "tf.MutableDenseHashTableV2"(%cst_2, %cst_3) <{container = "", shared_name = "table1", value_dtype = f32}> {key_dtype = i64} : (tensor<i64>, tensor<f32>) -> tensor<!tf_type.resource>
+// CHECK-NEXT:  "tf.LookupTableImportV2"([[TABLE1]], [[RETURNED_VALUES]]#0, [[RETURNED_VALUES]]#1) : (tensor<!tf_type.resource>, tensor<*xi64>, tensor<*xf32>) -> ()
+// CHECK-NOT:   "tf.RestoreV2"
+// CHECK-NEXT: return
+module {
+  func.func @restore_var_and_dense_table() {
+    %cst = "tf.Const"() {value = dense<"model/foo"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() {value = dense<["var1", "table1_keys", "table1_vals"]> : tensor<3x!tf_type.string>} : () -> tensor<3x!tf_type.string>
+    %cst_1 = "tf.Const"() {value = dense<["", "", ""]> : tensor<3x!tf_type.string>} : () -> tensor<3x!tf_type.string>
+    %0:3 = "tf.RestoreV2"(%cst, %cst_0, %cst_1) {dtypes = [i32, i64, f32]} : (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>) -> (tensor<*xi32>, tensor<*xi64>, tensor<*xf32>)
+    %1 = "tf.VarHandleOp"() {container = "", shared_name = "var1"} : () -> tensor<!tf_type.resource<tensor<*xi32>>>
+    "tf.AssignVariableOp"(%1, %0#0) : (tensor<!tf_type.resource<tensor<*xi32>>>, tensor<*xi32>) -> ()
+    %empty_key = "tf.Const"() {value = dense<-1> : tensor<i64>} : () -> tensor<i64>
+    %default_value = "tf.Const"() {value = dense<0.0> : tensor<f32>} : () -> tensor<f32>
+    %2 = "tf.MutableDenseHashTableV2"(%empty_key, %default_value) {container = "", key_dtype = i64, shared_name = "table1", value_dtype = f32} : (tensor<i64>, tensor<f32>) -> tensor<!tf_type.resource>
+    "tf.LookupTableImportV2"(%2, %0#1, %0#2) : (tensor<!tf_type.resource>, tensor<*xi64>, tensor<*xf32>) -> ()
+    return
+  }
+}
+
+
+// -----
+// restored variable assigned to one var handle and then read and assigned to another
+// CHECK-LABEL: func.func @restore_and_copy() {
+// CHECK-DAG:   [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}>
+// CHECK-DAG:   [[VAR_Y:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-DAG:   [[VAR_Z:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-DAG:   [[NEW_NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<2x!tf_type.string>}>
+// CHECK-DAG:   [[NEW_SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<2x!tf_type.string>}>
+// CHECK:       "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NEW_NAME]], [[NEW_SLICE]], [[VAR_Y]], [[VAR_Z]])
+// CHECK-SAME:  <{restored_dtypes = [f32, f32], returned_tensor_names = [], truncate_in_cast = array<i1: false, false>}>
+// CHECK-NOT:   "tf.RestoreV2"
+// CHECK-NOT:   "tf.ReadVariableOp"
+// CHECK-NEXT:  return
+
+module {
+  func.func @restore_and_copy() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    %read = "tf.ReadVariableOp"(%1) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %2 = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%2, %read) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+}
+
+// -----
+// restored variable assigned to one var handle and then read and assigned to another in a different block/function
+// CHECK-LABEL: func.func @restore_derived() {
+// CHECK-DAG:   [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_variables"> : tensor<!tf_type.string>}>
+// CHECK-DAG:   [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-DAG:   [[HANDLEZ:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-DAG:   [[NEW_NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<2x!tf_type.string>}>
+// CHECK-DAG:   [[NEW_SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<2x!tf_type.string>}>
+// CHECK:       "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NEW_NAME]], [[NEW_SLICE]], [[HANDLEY]], [[HANDLEZ]])
+// CHECK-SAME:  <{restored_dtypes = [f32, f32], returned_tensor_names = [], truncate_in_cast = array<i1: false, false>}>
+// CHECK-NOT:   "tf.RestoreV2"
+// CHECK-NEXT:  return
+
+module {
+  func.func @restore_derived() {
+    %cst = "tf.Const"() <{value = dense<"restore_variables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+
+  func.func @derived_init() {
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %1 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %2 = "tf.Identity"(%1) : (tensor<3x1xf32>) -> tensor<3x1xf32>
+    %3 = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%3, %2) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+}
+
+// -----
+// restored variable assigned to one var handle and then read, cast, and assigned to another
+// CHECK-LABEL: func.func @restore_and_cast_copy() {
+// CHECK-DAG:   [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}>
+// CHECK-DAG:   [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-DAG:   [[HANDLEZ:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
+// CHECK-DAG:   [[NEW_NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<2x!tf_type.string>}>
+// CHECK-DAG:   [[NEW_SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<2x!tf_type.string>}>
+// CHECK:       "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NEW_NAME]], [[NEW_SLICE]], [[HANDLEY]], [[HANDLEZ]])
+// CHECK-SAME:  <{restored_dtypes = [f32, f32], returned_tensor_names = [], truncate_in_cast = array<i1: false, true>}>
+// CHECK-NOT:   "tf.RestoreV2"
+// CHECK-NOT:   "tf.ReadVariableOp"
+// CHECK-NEXT:  return
+module {
+  func.func @restore_and_cast_copy() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    %read = "tf.ReadVariableOp"(%1) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %cast = "tf.Cast"(%read) <{Truncate = true}> : (tensor<3x1xf32>) -> tensor<3x1xbf16>
+    %2 = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
+    "tf.AssignVariableOp"(%2, %cast) : (tensor<!tf_type.resource<tensor<3x1xbf16>>>, tensor<3x1xbf16>) -> ()
+    return
+  }
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/propagate_static_shapes.mlir

@@ -0,0 +1,104 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -propagate-static-shapes %s | FileCheck %s
+
+// -----
+// CHECK-LABEL: func.func @callee(%arg0: tensor<?x?xi32> {tf._static_shape_arg_idx = 1 : i32}, %arg1: tensor<2xi64>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 123 : i64}
+// CHECK: return %arg0
+// CHECK-LABEL: func.func @main
+// CHECK-NEXT: %[[C0:.*]] = "tf.Const"
+// CHECK-NEXT: %[[C1:.*]] = "tf.IfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 123 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+// CHECK-NEXT: return %[[C1]]
+
+module {
+  func.func @callee(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 123 : i64} {
+    func.return %arg0 : tensor<?x?xi32>
+  }
+  func.func @main(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {
+    %0 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64>
+    %1 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+    %2 = "tf.IfrtCall"(%1) {program_id = 123 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
+    func.return %2 : tensor<?x?xi32>
+  }
+}
+
+// -----
+// CHECK-LABEL: func.func @callee(%arg0: tensor<?x?xi32> {tf._static_shape_arg_idx = 3 : i32}, %arg1: tensor<?xi32> {tf._static_shape_arg_idx = 4 : i32}, %arg2: tensor<?x?xf32>, %arg3: tensor<2xi64>, %arg4: tensor<1xi64>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>) attributes {tfrt_ifrt_serving.program_id = 456 : i64}
+// CHECK: return %arg0, %arg1, %arg2
+// CHECK-LABEL: func.func @main
+// CHECK-NEXT: %[[C0:.*]] = "tf.Const"
+// CHECK-NEXT: %[[C1:.*]] = "tf.Const"
+// CHECK-NEXT: %[[R:.*]]:3 = "tf.IfrtCall"(%arg0, %arg1, %arg2, %[[C0]], %[[C1]]) <{operandSegmentSizes = array<i32: 3, 2>, program_id = 456 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>, tensor<2xi64>, tensor<1xi64>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>)
+// CHECK-NEXT: return %[[R]]#0, %[[R]]#1, %[[R]]#2
+
+module {
+  func.func @callee(%arg0: tensor<?x?xi32>, %arg1: tensor<?xi32>, %arg2: tensor<?x?xf32>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>) attributes {tfrt_ifrt_serving.program_id = 456 : i64} {
+    func.return %arg0, %arg1, %arg2 : tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>
+  }
+  func.func @main(%arg0: tensor<?x?xi32>, %arg1: tensor<?xi32>, %arg2: tensor<?x?xf32>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>) {
+    %c0 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64>
+    %c1 = "tf.Const"() {value = dense<4> : tensor<1xi64>} : () -> tensor<1xi64>
+    %0 = "tf.SetStaticDimensionBounds"(%arg0, %c0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+    %1 = "tf.SetStaticDimensionBounds"(%arg1, %c1) : (tensor<?xi32>, tensor<1xi64>) -> tensor<?xi32>
+    %2:3 = "tf.IfrtCall"(%0, %1, %arg2) {program_id = 456 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 3, 0>} : (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>)
+    func.return %2#0, %2#1, %2#2 : tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>
+  }
+}
+
+// -----
+// CHECK-LABEL: func.func @callee(%arg0: tensor<?x?xi32> {tf._static_shape_arg_idx = 1 : i32}, %arg1: tensor<2xi64>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 789 : i64}
+// CHECK: return %arg0
+// CHECK-LABEL: func.func @main
+// CHECK-NEXT: %[[C0:.*]] = "tf.Const"
+// CHECK-NEXT: %[[C1:.*]] = "tf.IfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 789 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+// CHECK-NEXT: %[[C2:.*]] = "tf.IfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 789 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+// CHECK-NEXT: return %[[C2]]
+
+module {
+  func.func @callee(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 789 : i64} {
+    func.return %arg0 : tensor<?x?xi32>
+  }
+  func.func @main(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {
+    %0 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64>
+    %1 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+    %2 = "tf.IfrtCall"(%1) {program_id = 789 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
+    %3 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+    %4 = "tf.IfrtCall"(%3) {program_id = 789 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
+    func.return %4 : tensor<?x?xi32>
+  }
+}
+
+// -----
+// CHECK-LABEL: func.func @callee(%arg0: tensor<?x?xi32> {tf._static_shape_arg_idx = 1 : i32}, %arg1: tensor<2xi64>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 999 : i64}
+// CHECK: return %arg0
+// CHECK-LABEL: func.func @main
+// CHECK-NEXT: %[[C0:.*]] = "tf.Const"
+// CHECK-NEXT: %[[C1:.*]] = "tf.IfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 999 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+// CHECK-NEXT: %[[C2:.*]] = "tf.AsyncIfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 999 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+// CHECK-NEXT: return %[[C2]]
+
+module {
+  func.func @callee(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 999 : i64} {
+    func.return %arg0 : tensor<?x?xi32>
+  }
+  func.func @main(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {
+    %0 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64>
+    %1 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+    %2 = "tf.IfrtCall"(%1) {program_id = 999 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
+    %3 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+    %4 = "tf.AsyncIfrtCall"(%3) {program_id = 999 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
+    func.return %4 : tensor<?x?xi32>
+  }
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/rewrite_cluster_to_ifrt_call.mlir

@@ -0,0 +1,140 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -rewrite-cluster-to-ifrt-call %s | FileCheck %s
+// 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:       {operandSegmentSizes = array<i32: 1, 0>, program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []}
+// CHECK-SAME:       (tensor<1x3xf32>) -> tensor<1x3xf32>
+// CHECK:    return
+//
+// CHECK:  func.func @_ifrt_program__func(%arg0: tensor<1x3xf32>)
+// CHECK-SAME: __tpu_compile_metadata_text = "args { dtype: DT_FLOAT shape { dim { size: 1 } dim { size: 3 } } kind: PARAMETER sharding { } is_bounded_dynamic_dim: false } retvals { sharding { } } num_replicas: 1 num_cores_per_replica: 1 "
+// CHECK-SAME: device_assignment = []
+// CHECK-SAME: tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
+// CHECK:      return
+
+module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1"], tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1704 : i32}} {
+func.func @serving_default(%arg0: tensor<1x3xf32>) -> (tensor<1x3xf32>) {
+  %outputs  =  "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>
+  %0 = "tf_device.cluster_func"(%arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = [""], num_cores_per_replica = 1 : i64, device_assignment = [], topology = "", output_sharding_configuration = [""], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = false, use_tpu = true} : (tensor<1x3xf32>) -> (tensor<1x3xf32>)
+  return %0 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: @_func
+func.func private @_func(%arg0: tensor<1x3xf32>) -> (tensor<1x3xf32>) {
+  return %arg0 : tensor<1x3xf32>
+}
+}
+
+
+// -----
+// SPMD: one input and no return
+//
+// CHECK-LABEL: func.func @serving_default(%arg0: tensor<1x3xf32>) {
+// CHECK-NEXT:  "tf.IfrtCall"(%arg0)
+// CHECK-SAME:       {operandSegmentSizes = array<i32: 1, 0>, program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []}
+// CHECK-SAME:       (tensor<1x3xf32>) -> ()
+// CHECK:    return
+//
+// CHECK:  func.func @_ifrt_program__func(%arg0: tensor<1x3xf32>)
+// CHECK-SAME: __tpu_compile_metadata_text = "args { dtype: DT_FLOAT shape { dim { size: 1 } dim { size: 3 } } kind: PARAMETER sharding { type: OTHER tile_assignment_dimensions: 2 tile_assignment_dimensions: 1 tile_assignment_devices: 0 tile_assignment_devices: 1 } is_bounded_dynamic_dim: false } num_replicas: 1 num_cores_per_replica: 2 use_spmd_for_xla_partitioning: true "
+// CHECK-SAME: device_assignment = [0, 0, 0, 0, 0, 0, 0, 1]
+// CHECK-SAME: tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
+// CHECK:      return
+
+module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1"], tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1704 : i32}} {
+func.func @serving_default(%arg0: tensor<1x3xf32>) -> () {
+  %outputs  =  "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>
+  "tf_device.cluster_func"(%arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = ["{devices=[2,1]0,1}"], num_cores_per_replica = 2 : i64, device_assignment = [0, 0, 0, 0, 0, 0, 0, 1], topology = "\0A\04\01\01\01\02\10\01\18\02\22\08\00\00\00\00\00\00\00\01", output_sharding_configuration = [], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = true, use_tpu = true} : (tensor<1x3xf32>) -> ()
+  return
+}
+
+// CHECK-LABEL: @_func
+func.func private @_func(%arg0: tensor<1x3xf32>) -> () {
+  return
+}
+}
+
+// -----
+// Multiple ifrt calls and have two sharded arguments
+
+// CHECK-LABEL: func.func @serving_default(%arg0: tensor<3x1xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x1xf32> {
+// CHECK-NEXT:  %0 = "tf.IfrtCall"(%arg1, %arg0)
+// CHECK-SAME:       {operandSegmentSizes = array<i32: 2, 0>, program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []
+// CHECK-SAME:       (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+// CHECK-NEXT:    %1 = "tf.Identity"(%arg1) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
+// CHECK-NEXT:    %2 = "tf.IfrtCall"(%1, %arg0)
+// CHECK-SAME:       {operandSegmentSizes = array<i32: 2, 0>, program_id = [[PROGRAM_ID]] : i64, variable_arg_indices = []
+// CHECK-SAME:       (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+// CHECK-NEXT:    %3 = "tf.add"(%0, %2) : (tensor<1x1xf32>, tensor<1x1xf32>) -> tensor<1x1xf32>
+// CHECK:    return
+//
+// CHECK:  func.func @_ifrt_program__func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> tensor<1x1xf32>
+// CHECK-SAME:      device_assignment = [0, 0, 0, 0, 0, 0, 0, 1]
+// CHECK-SAME:      tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
+// CHECK-NEXT:     %0 = "tf.MatMul"(%arg0, %arg1)
+// CHECK:          return
+
+module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1"], tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1704 : i32}} {
+func.func @serving_default(%arg0: tensor<3x1xf32>,  %arg1: tensor<1x3xf32>) -> (tensor<1x1xf32>) {
+  %outputs  =  "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>
+  %outputs_0 = "tf_device.cluster_func"(%arg1, %arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = ["{devices=[2,1]0,1}", ""], num_cores_per_replica = 2 : i64, device_assignment = [0, 0, 0, 0, 0, 0, 0, 1], topology = "\0A\04\01\01\01\02\10\01\18\02\22\08\00\00\00\00\00\00\00\01", output_sharding_configuration = [""], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = true, use_tpu = true} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  %duplicate_arg =  "tf.Identity"(%arg1) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
+  %outputs_1 = "tf_device.cluster_func"(%duplicate_arg, %arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = ["{devices=[2,1]0,1}", ""], num_cores_per_replica = 2 : i64, device_assignment = [0, 0, 0, 0, 0, 0, 0, 1], topology = "\0A\04\01\01\01\02\10\01\18\02\22\08\00\00\00\00\00\00\00\01", output_sharding_configuration = [""], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = true, use_tpu = true} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  %outputs_2 = "tf.add"(%outputs_0, %outputs_1): (tensor<1x1xf32>, tensor<1x1xf32>) -> tensor<1x1xf32>
+  return %outputs_2 : tensor<1x1xf32>
+}
+
+// CHECK-LABEL: @_func
+func.func private @_func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> (tensor<1x1xf32>) {
+  %outputs_0 =  "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  return %outputs_0 : tensor<1x1xf32>
+}
+}
+
+
+// -----
+// Missing topology and device assignment attribute in spmd is ok
+
+// CHECK-LABEL: func.func @serving_default(%arg0: tensor<3x1xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x1xf32> {
+// CHECK-NEXT:  %0 = "tf.IfrtCall"(%arg1, %arg0)
+// CHECK-SAME:       {operandSegmentSizes = array<i32: 2, 0>, program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []
+// CHECK-SAME:       (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+// CHECK:    return
+//
+// CHECK:  func.func @_ifrt_program__func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> tensor<1x1xf32>
+// CHECK-SAME:       device_assignment = []
+// CHECK-SAME:      tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
+// CHECK-NEXT:     %0 = "tf.MatMul"(%arg0, %arg1)
+// CHECK:          return
+
+module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1"], tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1704 : i32}} {
+func.func @serving_default(%arg0: tensor<3x1xf32>,  %arg1: tensor<1x3xf32>) -> (tensor<1x1xf32>) {
+  %outputs  =  "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>
+  %outputs_0 = "tf_device.cluster_func"(%arg1, %arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = ["{devices=[2,1]0,1}", ""], num_cores_per_replica = 2 : i64, output_sharding_configuration = [""], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = true, use_tpu = true} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  return %outputs_0 : tensor<1x1xf32>
+}
+
+// CHECK-LABEL: @_func
+func.func private @_func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> (tensor<1x1xf32>) {
+  %outputs_0 =  "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  return %outputs_0 : tensor<1x1xf32>
+}
+}
+

tensorflow/compiler/mlir/tfrt/tests/ifrt/sink_variable_as_named_array.mlir

@@ -0,0 +1,144 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tf-device-decompose-resource-ops -sink-variable-as-named-array %s | FileCheck %s
+
+// -----
+// Basic test: all variables tensors are for devices and sinked as named ifrt arrays
+//
+//
+// CHECK-LABEL:  func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
+// CHECK-NEXT:   [[HANDLE2:%.*]] = "tf.VarHandleOp"
+// CHECK-NEXT:   [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"([[HANDLE2]])
+// CHECK-SAME:       used_by_host = false 
+// CHECK-NEXT:   [[RES:%.*]] = "tf.IfrtCall"([[KEY]], %arg0) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [0 : i32]}>
+// CHECK-SAME:    : (tensor<!tf_type.string>, tensor<1x3xf32>) -> tensor<1x1xf32>
+// CHECK-NEXT:    return [[RES]] : tensor<1x1xf32>
+//
+module {
+  func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %result = "tf.IfrtCall"(%2, %arg0) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = []}> : (tensor<3x1xf32>, tensor<1x3xf32>) -> (tensor<1x1xf32>)
+    return %result : tensor<1x1xf32>
+  }
+}
+
+// -----
+// Variable tensor for host can still be used.
+//
+// CHECK-LABEL:  func.func @serving_default(%arg0: tensor<1x3xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>) {
+// CHECK:  "tf.VarHandleOp"
+// CHECK-NOT:  [[VARIABLE:%.*]] = "tf.ReadVariableOp"
+// CHECK-NEXT:  [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"
+// CHECK-SAME:    used_by_host = true
+// CHECK-NEXT:  [[MATRES:%.*]] = "tf.MatMul"(%arg0, [[FUTURE]])
+// CHECK-NEXT:   [[RES:%.*]] = "tf.IfrtCall"(%arg0, [[KEY]]) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [1 : i32]}>
+// CHECK-NEXT:    return [[RES]], [[MATRES]] : tensor<1x1xf32>, tensor<1x1xf32>
+//
+module {
+  func.func @serving_default(%arg0: tensor<1x3xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>) {
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %3 = "tf.MatMul"(%arg0, %2) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+    %result = "tf.IfrtCall"(%arg0, %2) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = []}> : (tensor<1x3xf32>, tensor<3x1xf32>) -> (tensor<1x1xf32>)
+    return %result, %3 : tensor<1x1xf32>, tensor<1x1xf32>
+  }
+}
+
+// -----
+// Variable tensor is only for host
+//
+// CHECK-LABEL:  func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
+// CHECK:  "tf.VarHandleOp"
+// CHECK-NOT:   tf.ReadVariableOp
+// CHECK-NEXT:  [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"
+// CHECK-SAME:    used_by_host = true
+// CHECK-NEXT:  [[RES:%.*]] = "tf.MatMul"(%arg0, [[FUTURE]])
+// CHECK-NEXT:    return [[RES]] : tensor<1x1xf32>
+//
+module {
+  func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %3 = "tf.MatMul"(%arg0, %2) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+    return %3: tensor<1x1xf32>
+  }
+}
+
+// -----
+//  Resources that are created in the same module are not sinked.
+//
+// CHECK-LABEL:  func.func @serving_default
+// CHECK-NOT:  IfrtLoadVariable
+// CHECK:      "tf.VarHandleOp"
+// CHECK-NEXT: "tf.AssignVariableOp"
+// CHECK-NEXT: "tf.ReadVariableOp"
+// CHECK-NEXT: "tf.StatefulPartitionedCall"
+// CHECK-NEXT:  return 
+//
+module {
+  func.func @serving_default() -> tensor<*xi32> {
+    %cst = "tf.Const"() <{value = dense<"some_test.txt"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "Variable"}> : () -> tensor<!tf_type.resource<tensor<!tf_type.string>>>
+    "tf.AssignVariableOp"(%0, %cst) <{validate_shape = false}> : (tensor<!tf_type.resource<tensor<!tf_type.string>>>, tensor<!tf_type.string>) -> ()
+    %2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<!tf_type.string>>>) -> tensor<*x!tf_type.string>
+    %4 = "tf.StatefulPartitionedCall"(%2) <{config = "", config_proto = "", executor_type = "", f = @__initializer}> : (tensor<*x!tf_type.string>) -> tensor<*xi32>
+    return %4: tensor<*xi32>
+  }
+  func.func @__initializer(%arg0: tensor<*x!tf_type.string>) -> tensor<i32> {
+    %0 = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
+    return %0 : tensor<i32>
+  }
+}
+
+
+// -----
+//  Decomposable Resource Ops usage
+//
+// CHECK-LABEL:  func.func @serving_default
+// CHECK:      "tf.VarHandleOp"
+// CHECK-NEXT: "tf.IfrtLoadVariable"
+// CHECK-NEXT: "tf.GatherV2"
+// CHECK-NEXT:  return 
+//
+module {
+  func.func @serving_default() -> tensor<1x3xbf16> {
+    %cst = "tf.Const"() <{value = dense<[1]> : tensor<1xi32>}> : () -> tensor<1xi32>
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "Variable"}> : () -> tensor<!tf_type.resource<tensor<2x3xbf16>>>
+    %1 = "tf.ResourceGather"(%0, %cst) <{batch_dims = 0 : i64, validate_indices = true}> : (tensor<!tf_type.resource<tensor<2x3xbf16>>>, tensor<1xi32>) -> tensor<1x3xbf16>
+    return %1: tensor<1x3xbf16>
+  }
+}
+
+// -----
+// AsyncIfrtCall test: all variables tensors are for devices and sinked as named ifrt arrays
+//
+//
+// CHECK-LABEL:  func.func @serving_default_async(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
+// CHECK-NEXT:   [[HANDLE2:%.*]] = "tf.VarHandleOp"
+// CHECK-NEXT:   [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"([[HANDLE2]])
+// CHECK-SAME:       used_by_host = false
+// CHECK-NEXT:   [[RES:%.*]] = "tf.AsyncIfrtCall"([[KEY]], %arg0) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [0 : i32]}>
+// CHECK-SAME:    : (tensor<!tf_type.string>, tensor<1x3xf32>) -> tensor<1x1xf32>
+// CHECK-NEXT:    return [[RES]] : tensor<1x1xf32>
+//
+module {
+  func.func @serving_default_async(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %result = "tf.AsyncIfrtCall"(%2, %arg0) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = []}> : (tensor<3x1xf32>, tensor<1x3xf32>) -> (tensor<1x1xf32>)
+    return %result : tensor<1x1xf32>
+  }
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_device_cleanup.mlir

@@ -0,0 +1,22 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: tf-tfrt-opt %s -tf-device-cleanup | FileCheck %s
+
+// CHECK-LABEL: func @ops_with_device
+func.func @ops_with_device() {
+  %0 = "tf.VarHandleOp"() {container = "", shared_name = "var", device = "/device/..."} : () -> tensor<!tf_type.resource<tensor<1xf32>>>
+  // CHECK-NOT: device = "/device/..."
+  func.return
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_identity_propagation.mlir

@@ -0,0 +1,52 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: tf-tfrt-opt %s -tf-identity-propagation -canonicalize | FileCheck %s
+
+// CHECK-LABEL: func @identity
+// CHECK-SAME:    (%[[ARG0:.*]]: tensor<i32>)
+func.func @identity(%arg0: tensor<i32>) -> tensor<i32> {
+  // CHECK-NOT: "tf.Identity"
+  %0 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>
+  // CHECK: return %[[ARG0]]
+  func.return %0 : tensor<i32>
+}
+
+// CHECK-LABEL: func @identity_terminator
+// CHECK-SAME:    (%[[ARG0:.*]]: tensor<i32>)
+func.func @identity_terminator(%arg0: tensor<i32>) -> (tensor<*xi32>, tensor<i32>) {
+  // CHECK: %[[IDENTITY:.*]] = "tf.Identity"
+  %0 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<*xi32>
+  // CHECK-NOT: "tf.Identity"
+  %1 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>
+  // CHECK: return %[[IDENTITY]], %[[ARG0]]
+  func.return %0, %1 : tensor<*xi32>, tensor<i32>
+}
+
+// CHECK-LABEL: func @xla_sharding
+func.func @xla_sharding(%arg0: tensor<i32>) -> tensor<i32> {
+  // CHECK: %[[OUTPUT:.*]] = "tf.Identity"
+  %0 = "tf.Identity"(%arg0) {_XlaSharding = ""} : (tensor<i32>) -> tensor<i32>
+  // CHECK: return %[[OUTPUT]]
+  func.return %0 : tensor<i32>
+}
+
+// CHECK-LABEL: func @identity_n
+// CHECK-SAME:    (%[[ARG0:.*]]: tensor<i32>, %[[ARG1:.*]]: tensor<f32>)
+func.func @identity_n(%arg0: tensor<i32>, %arg1: tensor<f32>) -> (tensor<i32>, tensor<f32>) {
+  // CHECK-NOT: "tf.IdentityN"
+  %0:2 = "tf.IdentityN"(%arg0, %arg1) : (tensor<i32>, tensor<f32>) -> (tensor<i32>, tensor<f32>)
+  // CHECK: return %[[ARG0]], %[[ARG1]]
+  func.return %0#0, %0#1 : tensor<i32>, tensor<f32>
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_restore_merging.mlir

@@ -0,0 +1,35 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: tf-tfrt-opt %s -tf-restore-merging | FileCheck %s
+
+// CHECK-LABEL: func @single_restore_group
+// CHECK-SAME:    (%[[ARG0:.*]]: {{.*}})
+func.func @single_restore_group(%arg0: tensor<!tf_type.string>) -> (tensor<*xf32>, tensor<*xi32>) {
+  %0 = "tf.Const"() {value = dense<"foo"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
+  %1 = "tf.Const"() {value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
+  %2 = "tf.RestoreV2"(%arg0, %0, %1) : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<*xf32>
+
+  %3 = "tf.Const"() {value = dense<"bar"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
+  %4 = "tf.Const"() {value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
+  %5 = "tf.RestoreV2"(%arg0, %3, %4) : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<*xi32>
+
+  // CHECK:      %[[NAMES:.*]] = "tf.Const"() <{value = dense<["foo", "bar"]> : tensor<2x!tf_type.string>}>
+  // CHECK-NEXT:      %[[SHAPES:.*]] = "tf.Const"() <{value = dense<""> : tensor<2x!tf_type.string>}>
+  // CHECK-NEXT:      %[[TENSORS:.*]]:2 = "tf.RestoreV2"(%[[ARG0]], %[[NAMES]], %[[SHAPES]])
+  // CHECK-SAME:   -> (tensor<*xf32>, tensor<*xi32>)
+
+  // CHECK:      return %[[TENSORS]]#0, %[[TENSORS]]#1
+  func.return %2, %5 : tensor<*xf32>, tensor<*xi32>
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_restore_pruning.mlir

@@ -0,0 +1,39 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: tf-tfrt-opt -tf-restore-pruning %s | FileCheck %s
+
+// CHECK-LABEL:   func.func @prune_unused_restore
+func.func @prune_unused_restore() {
+  %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+  %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+  %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+  // CHECK-NOT: tf.RestoreV2
+  %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+  %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+  return
+}
+
+
+// CHECK-LABEL: func.func @used_restore_remains
+func.func @used_restore_remains() {
+  %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+  %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+  %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+  // CHECK: tf.RestoreV2
+  %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+  %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+  "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+  return
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_restore_splitting.mlir

@@ -0,0 +1,32 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: tf-tfrt-opt %s -tf-restore-splitting | FileCheck %s
+
+// CHECK-LABEL: func @single_restore
+// CHECK-SAME:    (%[[ARG0:.*]]: {{.*}})
+func.func @single_restore(%arg0: tensor<!tf_type.string>) -> (tensor<*xf32>, tensor<*xi32>) {
+  %0 = "tf.Const"() {value = dense<["foo", "bar"]> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
+  %1 = "tf.Const"() {value = dense<""> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
+  %2:2 = "tf.RestoreV2"(%arg0, %0, %1) : (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<*xf32>, tensor<*xi32>)
+
+  // CHECK: %[[FOO_NAME:.*]] = "tf.Const"() <{value = dense<"foo"> : tensor<1x!tf_type.string>}>
+  // CHECK: %[[FOO:.*]] = "tf.RestoreV2"(%[[ARG0]], %[[FOO_NAME]], {{.*}})
+
+  // CHECK: %[[BAR_NAME:.*]] = "tf.Const"() <{value = dense<"bar"> : tensor<1x!tf_type.string>}>
+  // CHECK: %[[BAR:.*]] = "tf.RestoreV2"(%[[ARG0]], %[[BAR_NAME]], {{.*}})
+
+  // CHECK: return %[[FOO]], %[[BAR]]
+  func.return %2#0, %2#1 : tensor<*xf32>, tensor<*xi32>
+}