ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl → 0.3.0.dev20240926__py3-none-any.whl

ai_edge_torch/hlfb/mark_pattern/passes.py

@@ -12,13 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Passes to clean up the model graph for pattern matching."""
+
 import torch
 
 
 def remove_clone_ops(gm: torch.fx.GraphModule):
-  # torch export adds additional aten.clone nodes to produce contiguous in memory tensors
-  # depending on tensor sizes for runtime efficiency. However, these unpredictable clone
-  # nodes can break the pattern matching. Thus remove all clones in model and pattern graphs.
+  """Removes clone ops from the graph.
+
+  torch export adds additional aten.clone nodes to produce contiguous in memory
+  tensors depending on tensor sizes for runtime efficiency. However, these
+  unpredictable clone nodes can break the pattern matching. Thus remove all
+  clones in model and pattern graphs.
+
+  Args:
+    gm: The graph module to remove clone ops from.
+
+  Returns:
+    The graph module with clone ops removed.
+  """
   for node in gm.graph.nodes:
     if node.op == "call_function" and node.name.startswith("clone"):
       node.replace_all_uses_with(node.args[0])
@@ -30,6 +42,14 @@ def remove_clone_ops(gm: torch.fx.GraphModule):
 
 
 def remove_dangling_args(gm: torch.fx.GraphModule):
+  """Removes dangling args from the graph.
+
+  Args:
+    gm: The graph module to remove dangling args from.
+
+  Returns:
+    The graph module with dangling args removed.
+  """
   nodes_to_erase = []
   for node in gm.graph.nodes:
     if node.op == "placeholder" and len(node.users) == 0:

ai_edge_torch/hlfb/mark_pattern/pattern.py

@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-import copy
+"""Mark pattern."""
+
 import dataclasses
 from typing import Any, Callable, Optional, Union
 
+from ai_edge_torch.hlfb.mark_pattern import passes
 import torch
 from torch.export.graph_signature import TensorArgument
 from torch.fx import Graph
@@ -23,8 +25,6 @@ from torch.fx import GraphModule
 from torch.fx.passes.utils.matcher_utils import InternalMatch
 from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
 
-from ai_edge_torch.hlfb.mark_pattern import passes
-
 
 def _are_equal(x: Any, y: Any) -> bool:
   if type(x) != type(y):
@@ -46,6 +46,7 @@ def _are_equal(x: Any, y: Any) -> bool:
 @dataclasses.dataclass
 class ScalarAttrTracker:
   """ScalarAttrTracker is used to track the occurrence of a pattern's
+
   scalar arg/attr in the pattern decomposed graph. Since a scalar attr
   to the pattern can be transformed and turned into a/some ops' scalar
   arg in the decomposed graph, it would be hard to programmatically get
@@ -58,27 +59,31 @@ class ScalarAttrTracker:
     pattern_arg_pos (int): the index of the attr to track in the pattern's
       export_args.
     transform (Callable): the transform function used when targeting the
-      occurrence of the attr value in the decomposed graph. An attr value
-      may be transformed during the decomposition and appear as a derived
-      value.
-    inverse_transform (Callable): the inverse transform function that maps
-      the transformed value back to the original attr value.
+      occurrence of the attr value in the decomposed graph. An attr value may be
+      transformed during the decomposition and appear as a derived value.
+    inverse_transform (Callable): the inverse transform function that maps the
+      transformed value back to the original attr value.
   """
 
   attr_name: str
   pattern_arg_pos: int
   transform: Callable = lambda x: x
   inverse_transform: Callable = lambda x: x
-  _source_targets: list[tuple[Any, Any]] = dataclasses.field(default_factory=list)
+  _source_targets: list[tuple[Any, Any]] = dataclasses.field(
+      default_factory=list
+  )
 
   def track(self, *sources):
     """Register magic values to track the (transformed) attr values in
+
     the pattern decomposed graph.
     """
     for source in sources:
       target = self.transform(source)
       if not _are_equal(self.inverse_transform(target), source):
-        raise Exception(f"Invalid transform/inverse_transform for {self.attr_name}")
+        raise Exception(
+            f"Invalid transform/inverse_transform for {self.attr_name}"
+        )
       self._source_targets.append([source, target])
     return self
 
@@ -155,24 +160,22 @@ class Pattern:
     """The PyTorch computation pattern to match against a model.
 
     Args:
-      name (str): the name of the pattern. It would be propagated to
-        the `name` attr in StableHLO composite ops for the matched
-        model subgraphs in the lowering.
+      name (str): the name of the pattern. It would be propagated to the `name`
+        attr in StableHLO composite ops for the matched model subgraphs in the
+        lowering.
       module (torch.nn.Module or Callable): the PyTorch computation.
-      export_args (tuple[Any]): the args used to export the pattern module
-        with torch.export.export. If export_args contains non-tensor
-        Python scalars, there must be a corresponding attr tracker
-        in `scalar_attr_trackers` for each scalar arg.
-      attr_builder (Callable[[Pattern, GraphModule, InternalMatch], Optional[dict[str, Any]]]):
-        the callable that produces the a scalar attrs dict, which would be
-        propagated to `attr` in StableHLO composite ops for the matched
-        model subgraphs in the lowering.
-      scalar_attr_trackers (list[ScalarAttrTracker]): the trackers
-        for scalar args in `export_args`, which are used to track
-        the attr occurrence(s) and retrieve their values from the
-        matched subgraph.
-      decomp_table (Optional[dict[torch._ops.OperatorBase, Callable]]):
-        The decomposition table to be run on the pattern's exported program.
+      export_args (tuple[Any]): the args used to export the pattern module with
+        torch.export.export. If export_args contains non-tensor Python scalars,
+        there must be a corresponding attr tracker in `scalar_attr_trackers` for
+        each scalar arg. attr_builder (Callable[[Pattern, GraphModule,
+        InternalMatch], Optional[dict[str, Any]]]): the callable that produces
+        the a scalar attrs dict, which would be propagated to `attr` in
+        StableHLO composite ops for the matched model subgraphs in the lowering.
+      scalar_attr_trackers (list[ScalarAttrTracker]): the trackers for scalar
+        args in `export_args`, which are used to track the attr occurrence(s)
+        and retrieve their values from the matched subgraph.
+      decomp_table (Optional[dict[torch._ops.OperatorBase, Callable]]): The
+        decomposition table to be run on the pattern's exported program.
     """
     if not isinstance(module, torch.nn.Module):
 
@@ -189,7 +192,9 @@ class Pattern:
 
     self.name = name
     self.attr_builder = attr_builder
-    self._scalar_attr_trackers = scalar_attr_trackers if scalar_attr_trackers else []
+    self._scalar_attr_trackers = (
+        scalar_attr_trackers if scalar_attr_trackers else []
+    )
 
     exported_program = torch.export.export(module, export_args)
     if decomp_table is not None:
@@ -201,7 +206,9 @@ class Pattern:
     self._scalar_attr_locations = []
     for tracker in self._scalar_attr_trackers:
       self._scalar_attr_locations.append(
-          _find_scalar_attr(module, export_args, tracker, decomp_table=decomp_table)
+          _find_scalar_attr(
+              module, export_args, tracker, decomp_table=decomp_table
+          )
       )
 
     # Sanitize graph_module for more precise pattern matching.
@@ -251,7 +258,9 @@ class Pattern:
         attrs = {}
 
       for loc in self._scalar_attr_locations:
-        attrs[loc.attr_name] = self._get_attr_value_from_pattern_match(match, loc)
+        attrs[loc.attr_name] = self._get_attr_value_from_pattern_match(
+            match, loc
+        )
 
       attrs = attrs if attrs else None
       match_with_attrs.append((match, attrs))

ai_edge_torch/hlfb/test/test_mark_pattern.py

@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Tests for mark_pattern."""
 
-import unittest
-
+from ai_edge_torch import lowertools
+from ai_edge_torch.hlfb import mark_pattern
+from ai_edge_torch.hlfb.mark_pattern import pattern as pattern_module
 import torch
-import torch_xla
 
-from ai_edge_torch.hlfb import mark_pattern
+from absl.testing import absltest as googletest
 
 
 def _export_stablehlo_mlir(model, args=None):
@@ -26,11 +27,10 @@ def _export_stablehlo_mlir(model, args=None):
     ep = torch.export.export(model, args)
   else:
     ep = model
-  stablehlo_gm = torch_xla.stablehlo.exported_program_to_stablehlo(ep)
-  return stablehlo_gm.get_stablehlo_text()
+  return lowertools.exported_program_to_mlir_text(ep)
 
 
-class TestMarkPattern(unittest.TestCase):
+class TestMarkPattern(googletest.TestCase):
 
   def test_mark_pattern(self):
 
@@ -39,7 +39,7 @@ class TestMarkPattern(unittest.TestCase):
       def forward(self, x):
         return x * x + x + x
 
-    pattern = mark_pattern.Pattern(
+    pattern = pattern_module.Pattern(
         "test.add",
         lambda a, b: a + b,
         export_args=(torch.rand(2, 2), torch.rand(2, 2)),
@@ -51,7 +51,12 @@ class TestMarkPattern(unittest.TestCase):
     mark_pattern.mark_pattern(exported_program.graph_module, pattern)
     mlir = _export_stablehlo_mlir(exported_program)
 
-    self.assertEqual(mlir.count('stablehlo.composite "test.add"'), 2)
+    lowertools.assert_string_count(
+        self,
+        mlir,
+        {'stablehlo.composite "test.add"': 2},
+        {"stablehlo.custom_call @mark_tensor": 6},
+    )
 
   def test_mark_pattern_with_attr_builder(self):
     class TestModel(torch.nn.Module):
@@ -59,7 +64,7 @@ class TestMarkPattern(unittest.TestCase):
       def forward(self, x):
         return x * x * x + x - x * x + x
 
-    pattern = mark_pattern.Pattern(
+    pattern = pattern_module.Pattern(
         "test.add",
         lambda a, b: a + b,
         export_args=(torch.rand(2, 2), torch.rand(2, 2)),
@@ -72,8 +77,16 @@ class TestMarkPattern(unittest.TestCase):
     mark_pattern.mark_pattern(exported_program.graph_module, pattern)
     mlir = _export_stablehlo_mlir(exported_program)
 
-    self.assertEqual(mlir.count('stablehlo.composite "test.add"'), 2)
-    self.assertEqual(mlir.count('composite_attributes = {alias = "test.test_add"}'), 2)
+    lowertools.assert_string_count(
+        self,
+        mlir,
+        {
+            'stablehlo.composite "test.add"': 2,
+            'composite_attributes = {alias = "test.test_add"}': 2,
+        },
+        {"stablehlo.custom_call @mark_tensor": 6},
+        {'{"alias": "test.test_add"}': 2},
+    )
 
   def test_mark_pattern_with_scalar_attr_tracker(self):
     class TestModel(torch.nn.Module):
@@ -84,12 +97,12 @@ class TestMarkPattern(unittest.TestCase):
           r = torch.nn.LogSoftmax(dim=idx % 2)(r) * x
         return r
 
-    pattern = mark_pattern.Pattern(
+    pattern = pattern_module.Pattern(
         "test.log_softmax",
         lambda x, dim: torch.nn.functional.log_softmax(x, dim=dim),
         export_args=(torch.rand(10, 10, 10), 1),
         scalar_attr_trackers=[
-            mark_pattern.ScalarAttrTracker("dim", pattern_arg_pos=1)
+            pattern_module.ScalarAttrTracker("dim", pattern_arg_pos=1)
             .track(0)
             .track(1)
             .track(2),
@@ -102,9 +115,17 @@ class TestMarkPattern(unittest.TestCase):
     mark_pattern.mark_pattern(exported_program.graph_module, pattern)
     mlir = _export_stablehlo_mlir(exported_program)
 
-    self.assertEqual(mlir.count('stablehlo.composite "test.log_softmax"'), 5)
-    self.assertEqual(mlir.count("composite_attributes = {dim = 0 : i64}"), 3)
-    self.assertEqual(mlir.count("composite_attributes = {dim = 1 : i64}"), 2)
+    lowertools.assert_string_count(
+        self,
+        mlir,
+        {
+            'stablehlo.composite "test.log_softmax"': 5,
+            "composite_attributes = {dim = 0 : i64}": 3,
+            "composite_attributes = {dim = 1 : i64}": 2,
+        },
+        {"stablehlo.custom_call @mark_tensor": 10},
+        {'{"dim": 0}': 3, '{"dim": 1}': 2},
+    )
 
   def test_mark_tangent_model_and_pattern_input(self):
     class TestModel(torch.nn.Module):
@@ -114,7 +135,7 @@ class TestMarkPattern(unittest.TestCase):
         z = z + y
         return z
 
-    pattern = mark_pattern.Pattern(
+    pattern = pattern_module.Pattern(
         "test.relu",
         lambda x: torch.ops.aten.relu(x),
         export_args=(torch.rand(2, 2),),
@@ -126,8 +147,13 @@ class TestMarkPattern(unittest.TestCase):
     mark_pattern.mark_pattern(exported_program.graph_module, pattern)
     mlir = _export_stablehlo_mlir(exported_program)
 
-    self.assertEqual(mlir.count('stablehlo.composite "test.relu'), 1)
+    lowertools.assert_string_count(
+        self,
+        mlir,
+        {'stablehlo.composite "test.relu"': 1},
+        {"stablehlo.custom_call @mark_tensor": 2},
+    )
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py

@@ -12,23 +12,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Tests for StableHLOCompositeBuilder."""
+
 import math
-import unittest
 
+from ai_edge_torch import config
+from ai_edge_torch import lowertools
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
 import torch
 import torch.nn.functional as F
-import torch_xla
 
-from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+from absl.testing import absltest as googletest
 
 
 def _export_stablehlo_mlir(model, args):
   ep = torch.export.export(model, args)
-  stablehlo_gm = torch_xla.stablehlo.exported_program_to_stablehlo(ep)
-  return stablehlo_gm.get_stablehlo_text()
+  return lowertools.exported_program_to_mlir_text(ep)
 
 
-class TestStableHLOCompositeBuilder(unittest.TestCase):
+@googletest.skipIf(
+    not config.Config.use_torch_xla,
+    reason="The odml_torch counter part is in odml_torch.",
+)
+class TestStableHLOCompositeBuilder(googletest.TestCase):
 
   def test_build_composite(self):
     class SampleModel(torch.nn.Module):
@@ -80,7 +86,9 @@ class TestStableHLOCompositeBuilder(unittest.TestCase):
         super().__init__()
 
       def log_softmax(self, x: torch.Tensor, dim: int):
-        builder = StableHLOCompositeBuilder(name="test.log_softmax", attr={"dim": dim})
+        builder = StableHLOCompositeBuilder(
+            name="test.log_softmax", attr={"dim": dim}
+        )
         x = builder.mark_inputs(x)
         y = torch.nn.functional.log_softmax(x, dim=dim)
         y = builder.mark_outputs(y)
@@ -126,7 +134,8 @@ class TestStableHLOCompositeBuilder(unittest.TestCase):
     self.assertEqual(mlir.count('stablehlo.composite "test.log_softmax"'), 1)
     self.assertEqual(
         mlir.count(
-            'composite_attributes = {dim = 0 : i64, source = "torch.nn", version = 1.000000e+00 : f32}'
+            'composite_attributes = {dim = 0 : i64, source = "torch.nn",'
+            " version = 1.000000e+00 : f32}"
         ),
         1,
     )
@@ -236,8 +245,12 @@ class TestStableHLOCompositeBuilder(unittest.TestCase):
     self.assertEqual(
         mlir.count('stablehlo.composite "test.scaled_dot_product_attention"'), 2
     )
-    self.assertEqual(mlir.count("composite_attributes = {include_captanh = true}"), 1)
-    self.assertEqual(mlir.count("composite_attributes = {include_captanh = false}"), 1)
+    self.assertEqual(
+        mlir.count("composite_attributes = {include_captanh = true}"), 1
+    )
+    self.assertEqual(
+        mlir.count("composite_attributes = {include_captanh = false}"), 1
+    )
 
   def test_build_composite_with_multiple_inputs_outputs(self):
     class SampleModel(torch.nn.Module):
@@ -267,4 +280,4 @@ class TestStableHLOCompositeBuilder(unittest.TestCase):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/lowertools/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================
+
+from ._shim import *
+from .common_utils import flat_dict_names
+from .test_utils import *

ai_edge_torch/lowertools/_shim.py

@@ -0,0 +1,80 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================
+
+from typing import Any, Optional
+
+from ai_edge_torch import config
+from ai_edge_torch._convert import signature
+from ai_edge_torch.quantize import quant_config as qcfg
+import torch
+
+# isort: off
+if config.Config.use_torch_xla:
+  from ai_edge_torch.lowertools import torch_xla_utils as utils
+  from ai_edge_torch.lowertools.torch_xla_utils import exported_program_to_mlir_text
+  from torch_xla.experimental.mark_pattern_utils import StableHLOCompositeBuilder
+  from torch_xla.experimental.xla_marker import serialize_composite_attr
+  # The following imports are needed to register the needed torch_xla ops.
+  import torch_xla.experimental.xla_marker
+  import torch_xla.experimental.xla_mlir_debuginfo
+
+  mark_tensor_op = torch.ops.xla.mark_tensor.default
+  write_mlir_debuginfo_op = torch.ops.xla.write_mlir_debuginfo.default
+else:
+  from ai_edge_torch.lowertools import odml_torch_utils as utils
+  from ai_edge_torch.lowertools.odml_torch_utils import exported_program_to_mlir_text
+  from ai_edge_torch.odml_torch.composite import StableHLOCompositeBuilder
+  from ai_edge_torch.odml_torch.composite.mark_tensor import serialize_composite_attr
+  from ai_edge_torch.odml_torch.composite.mark_tensor import mark_tensor_op
+  from ai_edge_torch.odml_torch.debuginfo import write_mlir_debuginfo_op
+# isort: on
+
+
+def exported_programs_to_tflite(
+    exported_programs: list[torch.export.ExportedProgram],
+    signatures: list[signature.Signature],
+    *,
+    quant_config: Optional[qcfg.QuantConfig] = None,
+    _tfl_converter_flags: Optional[dict[str, Any]] = None,
+):
+  """Converts a list of ExportedProgram to a TFLite model.
+
+  Args:
+    exported_programs: A list of ExportedProgram.
+    signatures: A list of Signature.
+    quant_config: A QuantConfig.
+    _tfl_converter_flags: A dict of flags for TFLiteConverter.
+
+  Returns:
+    A TFLite model.
+  """
+  if _tfl_converter_flags is None:
+    _tfl_converter_flags = {}
+
+  bundles: list[utils.MlirBundle] = [
+      utils.exported_program_to_mlir(exported, sig.flat_args)
+      for exported, sig in zip(exported_programs, signatures)
+  ]
+
+  merged_bundle: utils.MergedBundle = utils.merge_mlir_bundles(
+      bundles, signatures, exported_programs
+  )
+
+  return utils.merged_bundle_to_tfl_model(
+      merged_bundle,
+      signatures,
+      quant_config=quant_config,
+      _tfl_converter_flags=_tfl_converter_flags,
+  )

ai_edge_torch/lowertools/common_utils.py

@@ -0,0 +1,142 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================
+
+import logging
+from typing import List
+
+from ai_edge_torch._convert import signature as signature_module
+import tensorflow as tf
+import torch
+import torch.utils._pytree as pytree
+
+
+def _flatten_list(l: List) -> List:
+  flattened = []
+  for item in l:
+    if isinstance(item, list):
+      flattened.extend(_flatten_list(item))
+    else:
+      flattened.append(item)
+  return flattened
+
+
+def flat_dict_names(
+    tree_spec: pytree.TreeSpec, context: pytree.Context
+) -> List[str]:
+  """Given a TreeSpec, this produces a list of names for the leaves.
+
+  The list of names embeddeds the structure of the tree_spec. A nesting level is
+  indicated by an `_` and elements in a list are indicated by `_<index>`.
+
+  TODO b/361601485: The flattening of names is not collision-free and needs to
+  be revised.
+
+  Args:
+    tree_spec: The TreeSpec to extract the names from.
+    context: The context used to check if the provided spec belongs to a
+      dictionary or a list.
+
+  Returns:
+    A list of flattened names.
+  """
+  flat_names = []
+  if context is None:
+    for i, spec in enumerate(tree_spec):
+      if spec.children_specs:
+        flat_names.extend([
+            f"{i}_{name}"
+            for name in flat_dict_names(spec.children_specs, spec.context)
+        ])
+      else:
+        flat_names.append(f"{i}")
+  else:
+    flat_ctx = _flatten_list(context)
+    for prefix, spec in zip(flat_ctx, tree_spec):
+      leaf_flat_names = flat_dict_names(spec.children_specs, spec.context)
+      if leaf_flat_names:
+        flat_names.extend([f"{prefix}_{name}" for name in leaf_flat_names])
+      else:
+        flat_names.append(prefix)
+
+  return flat_names
+
+
+def _torch_to_tf_variable(torch_tensor: torch.Tensor):
+  if not torch_tensor.is_contiguous():
+    torch_tensor = torch_tensor.contiguous()
+
+  try:
+    dlpack_capsule = torch.utils.dlpack.to_dlpack(torch_tensor)
+    tf_tensor = tf.experimental.dlpack.from_dlpack(dlpack_capsule)
+  except Exception:
+    logging.info(
+        "Can not use dlpack to convert torch tensors. Falling back to numpy."
+    )
+    nparray = torch_tensor.cpu().detach().numpy()
+    tf_tensor = tf.convert_to_tensor(nparray)
+
+  return tf.Variable(tf_tensor, trainable=False)
+
+
+def _get_states(
+    exported_programs: list[torch.export.ExportedProgram],
+    signatures: list[signature_module.Signature],
+):
+  for exported_program, signature in zip(exported_programs, signatures):
+    args, _ = exported_program.example_inputs
+    # Calling this to get **all** the state including model buffers.
+    _flat_input_args = exported_program._graph_module_flat_inputs(args, {})
+    for tensor, input_spec in zip(
+        _flat_input_args, exported_program.graph_signature.input_specs
+    ):
+      # Only interested in Tensors that are part of the state (and not user input).
+      if (
+          not isinstance(tensor, torch.Tensor)
+          or input_spec.kind
+          == torch.export.graph_signature.InputKind.USER_INPUT
+      ):
+        continue
+      yield signature, tensor, input_spec
+
+
+def _tensor_unique_id(tensor: torch.Tensor):
+  return (
+      str(tensor.device),
+      tensor.shape,
+      tensor.stride(),
+      tensor.untyped_storage().data_ptr(),
+  )
+
+
+def gather_state_dict(
+    exported_programs: list[torch.export.ExportedProgram],
+    signatures: list[signature_module.Signature],
+):
+  deduped_tensor_map = {}
+
+  for _, tensor, _ in _get_states(exported_programs, signatures):
+    unique_id = _tensor_unique_id(tensor)
+    deduped_tensor_map[unique_id] = _torch_to_tf_variable(tensor)
+
+  state_dict = {}
+  for signature, tensor, input_spec in _get_states(
+      exported_programs, signatures
+  ):
+    unique_id = _tensor_unique_id(tensor)
+    state_dict[signature.name + "_" + input_spec.target] = deduped_tensor_map[
+        unique_id
+    ]
+
+  return state_dict, list(deduped_tensor_map.values())