ai-edge-torch-nightly 0.4.0.dev20250309__py3-none-any.whl → 0.4.0.dev20250311__py3-none-any.whl

ai_edge_torch/odml_torch/composite/mark_tensor.py

@@ -82,6 +82,9 @@ _torch_library.ODML_TORCH_LIB.define(
 
 mark_tensor_op = torch.ops.odml_torch.mark_tensor.default
 
+# Prevent composite inputs and outputs from being DCE'd during torch.export.
+torch.fx.node.has_side_effect(mark_tensor_op)
+
 
 @torch.library.impl(
     _torch_library.ODML_TORCH_LIB, "mark_tensor", "CompositeExplicitAutograd"

ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py

@@ -41,8 +41,8 @@ class StableHLOCompositeBuilder:
     self.attr = attr
     self.name = name
     self.id = _get_uuid()
-    self._inputs = []
-    self._outputs = []
+    self._input_cnt = 0
+    self._output_cnt = 0
 
   def _mark_tensor(self, *tensors: torch.Tensor, is_input: bool):
     """Mark the input/output tensors of the StableHLO Composite."""
@@ -53,9 +53,20 @@ class StableHLOCompositeBuilder:
         else None
     )
 
-    for pos, tensor in enumerate(tensors):
+    def _pos() -> int:
+      if is_input:
+        self._input_cnt += 1
+        return self._input_cnt - 1
+      else:
+        self._output_cnt += 1
+        return self._output_cnt - 1
+
+    for tensor in tensors:
       if not isinstance(tensor, torch.Tensor):
         raise ValueError(f"input must be a torch tensor. Got {type(tensor)}.")
+
+      pos = _pos()
+
       marked_tensors.append(
           mark_tensor.mark_tensor_op(
               tensor,

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.4.0.dev20250309.dist-info → ai_edge_torch_nightly-0.4.0.dev20250311.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.4.0.dev20250309
+Version: 0.4.0.dev20250311
 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.dev20250309.dist-info → ai_edge_torch_nightly-0.4.0.dev20250311.dist-info}/RECORD

@@ -2,7 +2,7 @@ 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=1j-Ebny0vsJeVlpK_uQTeonyGRfzjE3sW8IZopMHuDM,706
+ai_edge_torch/version.py,sha256=nA0GT6cYQD_gQdlYKHCiF75nDcTUgY4bU2Zti4BK2gI,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=gpXQnifODU-mWxkUZw_3ov1lEYBw1SPVIcqj5k7pTGo,5550
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -205,8 +205,8 @@ ai_edge_torch/odml_torch/export.py,sha256=7l8R0DEq_vfns8iWpruMlIyaIKZAFzoAy369-7
 ai_edge_torch/odml_torch/export_utils.py,sha256=QeA37Irlty6AiIBuqmHmJgn3lqahBQ5xsh6IKRoKm1g,4774
 ai_edge_torch/odml_torch/tf_integration.py,sha256=NN29WeXmHZ0S1RPDFHUnBi2DEjMvAtwczStPYIsQ1w8,4849
 ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xKA9rknO93JS9kw9w_c,778
-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/composite/mark_tensor.py,sha256=yWMTkE0IESPUQyL20M6b_hv3VTptm9owC_evZcvoEMU,3647
+ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py,sha256=h6DQkYVS4fnKMALIVdU6Q7J6Ehg3hMCV4C406SyIk3k,3513
 ai_edge_torch/odml_torch/debuginfo/__init__.py,sha256=3A_lMyj-B-DOhLJG6WmjKvZK5te2rXje8FrfqOhZsN0,959
 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
@@ -233,8 +233,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.dev20250309.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.4.0.dev20250309.dist-info/METADATA,sha256=U7lJyT7P05KA4eOEZ_nkLRWjp_iihoX2G9QZKBs8H7I,1966
-ai_edge_torch_nightly-0.4.0.dev20250309.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.4.0.dev20250309.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.4.0.dev20250309.dist-info/RECORD,,
+ai_edge_torch_nightly-0.4.0.dev20250311.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.4.0.dev20250311.dist-info/METADATA,sha256=jfoiraSEOCTVbZMoOK-4snxTq9ejUm43i5z1EuWKTLU,1966
+ai_edge_torch_nightly-0.4.0.dev20250311.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.4.0.dev20250311.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.4.0.dev20250311.dist-info/RECORD,,