ai-edge-torch-nightly 0.4.0.dev20250224__py3-none-any.whl → 0.4.0.dev20250226__py3-none-any.whl

ai_edge_torch/_convert/conversion.py

@@ -125,6 +125,10 @@ def convert_signatures(
     else:
       exported_program = torch.export.export(**kwargs, strict=True)
 
+    exported_program = fx_infra.graph_utils.reset_from_node_meta(
+        exported_program
+    )
+
     exported_program = fx_infra.safe_run_decompositions(
         exported_program,
         fx_infra.decomp.pre_convert_decomp(),

ai_edge_torch/fx_infra/graph_utils.py

@@ -13,7 +13,10 @@
 # limitations under the License.
 # ==============================================================================
 """FX graph utilities."""
+
+from packaging import version
 import torch
+from torch.fx import traceback
 
 
 def remove_dangling_args(graph_module: torch.fx.GraphModule):
@@ -40,3 +43,32 @@ def remove_assert_tensor_metadata_nodes(graph_module: torch.fx.GraphModule):
   graph_module.graph.lint()
   graph_module.recompile()
   return graph_module
+
+
+def is_torch_version_under(torch_version: str) -> bool:
+  """Checks if the current torch version is under the given version."""
+  if not torch_version:
+    raise ValueError("torch_version cannot be empty.")
+  current_version = version.parse(torch.__version__)
+  compared_version = version.parse(torch_version)
+  return current_version < compared_version
+
+
+def reset_from_node_meta(ep: torch.export.ExportedProgram):
+  """Resets the "from_node" meta field to fx node name only for the exported program."""
+
+  for node in ep.graph.nodes:
+    if not hasattr(node, "meta") or "from_node" not in node.meta:
+      continue
+    if is_torch_version_under("2.6.0.dev0"):
+      # For torch version under 2.6.0, the history stack is a list of tuple. We
+      # will only keep the current node's name in the history stack.
+      history = [(node.name,)]
+    else:
+      # Clean up the history stack by keeping only the current node info (fx
+      # node name and graph id) in a list of size 1. Clear the "from_node" field
+      # to prevent redundant additions to the history stack.
+      history = [traceback.NodeSource(node)]
+      history[0].from_node = []
+    node.meta["from_node"] = history
+  return ep

ai_edge_torch/odml_torch/debuginfo/_build.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 import re
+from ai_edge_torch.fx_infra import graph_utils
 import torch
 
 
@@ -48,8 +49,6 @@ def _get_canonical_filename(filename: str):
     The canonicalized filename.
   """
 
-  # TODO(yijieyang): We should add a config option to provide a regex to strip
-  # from the debug info. Currently absolute path is used.
   return filename
 
 
@@ -57,9 +56,13 @@ def _get_canoical_nodename(node: torch.fx.Node) -> str:
   """Get the canonical node name from the node's history."""
 
   history = node.meta.get("from_node", [])
+  if not history:
+    return None
 
-  if len(history) > 1:  # Compatible with torch version under 2.6.0
-    return history[1][0]
+  # Compatible with torch version under 2.6.0. The history stack is a list of
+  # tuple. The first element of the first tuple is the node name.
+  if graph_utils.is_torch_version_under("2.6.0.dev0"):
+    return history[0][0]
 
   if not hasattr(history[0], "name"):
     return None
@@ -68,12 +71,10 @@ def _get_canoical_nodename(node: torch.fx.Node) -> str:
     names.append(history[0].name)
     history = history[0].from_node
 
-  # Based on the experiment, the third to last name in the history stack
-  # can be mapped to the original torch node name. The history stack is
-  # generated by tracing the node's transformation history during lowering.
-  if len(names) > 2:
-    return names[-3]
-  return None
+  # The history stack is generated by tracing the node's transformation history
+  # during lowering. The last name in the history stack is used to map to the
+  # original torch fx node name.
+  return names[-1]
 
 
 def build_mlir_file_debuginfo(node: torch.fx.Node):

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.4.0.dev20250224"
+__version__ = "0.4.0.dev20250226"

{ai_edge_torch_nightly-0.4.0.dev20250224.dist-info → ai_edge_torch_nightly-0.4.0.dev20250226.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.4.0.dev20250224
+Version: 0.4.0.dev20250226
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.4.0.dev20250224.dist-info → ai_edge_torch_nightly-0.4.0.dev20250226.dist-info}/RECORD

@@ -2,9 +2,9 @@ ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,120
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=Ga5ZnqgXb2idvQhACflPLWvlRmVZQ38CA2KtcWsGEKU,706
+ai_edge_torch/version.py,sha256=Redqgp3EjtlXSINPZlLb-pjbUH61Ie1ejPLMQ8bl_lE,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=mckvxznKLXdF2HuJg_IxQaT5Ty-iWl_iXElHEugH3VI,5452
+ai_edge_torch/_convert/conversion.py,sha256=gpXQnifODU-mWxkUZw_3ov1lEYBw1SPVIcqj5k7pTGo,5550
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
 ai_edge_torch/_convert/converter.py,sha256=075F8LRewk_033Ebsnft7FJr3KgtIbtZ_-8udIPy6ho,9980
 ai_edge_torch/_convert/signature.py,sha256=-YKJdLk-eNEHfhdPCtcQVtZf915SoVePEFxKXPPf16c,2572
@@ -41,7 +41,7 @@ ai_edge_torch/fx_infra/__init__.py,sha256=APjkSqEfwDxcnI8k53rGi3Ef-G2L-M8fdaPGpx
 ai_edge_torch/fx_infra/_canonicalize_pass.py,sha256=GDRoDdPVQw--QQFTT5J_C3TVuphL31m6K6F1-67SE4s,1097
 ai_edge_torch/fx_infra/_safe_run_decompositions.py,sha256=ZbWheeZ8ydsxCk2aVGUgUynrkEkBOMjBCzPhS5uq4sU,2595
 ai_edge_torch/fx_infra/decomp.py,sha256=S58SCgwMHYVFl_hJwlJxvu2wcI-AGNn82gel3qmTPrU,2500
-ai_edge_torch/fx_infra/graph_utils.py,sha256=3UZAOHWOUh2LCj1E2_AKQn3gRDILi9JCdqSScjyOd4M,1535
+ai_edge_torch/fx_infra/graph_utils.py,sha256=nqGe-xIJ77RamSUh0UYyI2XHOsZqFDWax-vpRAtVR_E,2796
 ai_edge_torch/fx_infra/pass_base.py,sha256=Ic2AlhSoRFscz6l7gJKvWVNMDLQFfAw5kRf84-ZR9qM,2904
 ai_edge_torch/generative/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -205,7 +205,7 @@ ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xK
 ai_edge_torch/odml_torch/composite/mark_tensor.py,sha256=U--rwl-XkWKgkdXCXDn6yySug8FR66o1YFUAIoSaWW4,3523
 ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py,sha256=2Y52E_gLeoXpMcPpV-svXsgN3JbEIjnPVjm0xkpTUdQ,3319
 ai_edge_torch/odml_torch/debuginfo/__init__.py,sha256=3A_lMyj-B-DOhLJG6WmjKvZK5te2rXje8FrfqOhZsN0,959
-ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=k9Kas790kpMS5OrVcLzIr48ejAzcc2smrroKAHHM7TQ,3311
+ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=6Ns2rlfOilLJEk5cUxlkRwm2uxOgEF2-0S2DMcOqr6A,3319
 ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py,sha256=IvOBQyROI9WHS3umHRxsDW-1YElU9BPWzKtJA2eKWOI,1739
 ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNitEeg-IoBUGNfUxsDSA,798
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=LqwZ1vCJTSOzgzvH8LUAN-sAkF-l_pGj1AMEIzAqHCA,6638
@@ -230,8 +230,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.4.0.dev20250224.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.4.0.dev20250224.dist-info/METADATA,sha256=K_9APQQGmsQkmF3tdAhAnLvwHAdFdkvlotoJ2KzrEvA,1966
-ai_edge_torch_nightly-0.4.0.dev20250224.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.4.0.dev20250224.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.4.0.dev20250224.dist-info/RECORD,,
+ai_edge_torch_nightly-0.4.0.dev20250226.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.4.0.dev20250226.dist-info/METADATA,sha256=N6T5-MKa5Ztwx_XE7OJ8wiw2BC00e0dQxgngvI9S6CU,1966
+ai_edge_torch_nightly-0.4.0.dev20250226.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.4.0.dev20250226.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.4.0.dev20250226.dist-info/RECORD,,