ai-edge-torch-nightly 0.3.0.dev20241224__py3-none-any.whl → 0.3.0.dev20250105__py3-none-any.whl

ai_edge_torch/_convert/conversion.py

@@ -14,9 +14,9 @@
 # ==============================================================================
 
 import logging
-import os
 from typing import Any, Literal, Optional, Union
 
+import ai_edge_torch
 from ai_edge_torch import fx_pass_base
 from ai_edge_torch import lowertools
 from ai_edge_torch import model
@@ -26,8 +26,6 @@ from ai_edge_torch.generative import fx_passes as generative_fx_passes
 from ai_edge_torch.quantize import quant_config as qcfg
 import torch
 
-os.environ["EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM"] = "1"
-
 
 def _run_convert_passes(
     exported_program: torch.export.ExportedProgram,
@@ -35,21 +33,27 @@ def _run_convert_passes(
   exported_program = generative_fx_passes.run_generative_passes(
       exported_program
   )
-  exported_program = fx_pass_base.run_passes(
-      exported_program,
-      [
-          fx_passes.BuildInterpolateCompositePass(),
-          fx_passes.CanonicalizePass(),
-          fx_passes.OptimizeLayoutTransposesPass(),
-          fx_passes.CanonicalizePass(),
-          fx_passes.BuildAtenCompositePass(),
-          fx_passes.CanonicalizePass(),
-          fx_passes.RemoveNonUserOutputsPass(),
-          fx_passes.CanonicalizePass(),
-          fx_passes.InjectMlirDebuginfoPass(),
-          fx_passes.CanonicalizePass(),
-      ],
-  )
+
+  passes = [
+      fx_passes.BuildInterpolateCompositePass(),
+      fx_passes.CanonicalizePass(),
+      fx_passes.OptimizeLayoutTransposesPass(),
+      fx_passes.CanonicalizePass(),
+      fx_passes.BuildAtenCompositePass(),
+      fx_passes.CanonicalizePass(),
+      fx_passes.RemoveNonUserOutputsPass(),
+      fx_passes.CanonicalizePass(),
+  ]
+
+  # Debuginfo is not injected automatically by odml_torch. Only inject
+  # debuginfo via fx pass when using torch_xla.
+  if ai_edge_torch.config.use_torch_xla:
+    passes += [
+        fx_passes.InjectMlirDebuginfoPass(),
+        fx_passes.CanonicalizePass(),
+    ]
+
+  exported_program = fx_pass_base.run_passes(exported_program, passes)
   return exported_program
 
 

ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py

@@ -62,6 +62,7 @@ def _wrap_call_function_node_with_debuginfo_writer(node: torch.fx.GraphModule):
 
 
 class InjectMlirDebuginfoPass(fx_pass_base.PassBase):
+  """DEPRECATED: Debuginfo is injected automatically by odml_torch."""
 
   def call(self, graph_module: torch.fx.GraphModule):
     for node in graph_module.graph.nodes:

ai_edge_torch/lowertools/torch_xla_utils.py

@@ -27,6 +27,9 @@ if "PJRT_DEVICE" not in os.environ:
   # https://github.com/google-ai-edge/ai-edge-torch/issues/326
   os.environ["PJRT_DEVICE"] = "CPU"
 
+os.environ["EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM"] = "1"
+
+
 from ai_edge_torch import model
 from ai_edge_torch._convert import conversion_utils
 from ai_edge_torch._convert import signature as signature_module

ai_edge_torch/odml_torch/export.py

@@ -202,7 +202,7 @@ class MlirLowered:
       target_version = stablehlo.get_minimum_version()
     else:
       target_version = stablehlo.get_version_from_compatibility_requirement(
-          stablehlo.StablehloCompatibilityRequirement.WEEK_4
+          stablehlo.StablehloCompatibilityRequirement.WEEK_12
       )
     module_bytecode = xla_extension.mlir.serialize_portable_artifact(
         self.module_bytecode, target_version
@@ -222,11 +222,6 @@ class MlirLowered:
     # Lazy importing TF when execution is needed.
     return self.tf_function(*args)
 
-  def to_flatbuffer(self):
-    from . import tf_integration
-
-    return tf_integration.mlir_to_flatbuffer(self)
-
 
 # TODO(b/331481564) Make this a ai_edge_torch FX pass.
 def _convert_i64_to_i32(exported_program: torch.export.ExportedProgram):

ai_edge_torch/odml_torch/tf_integration.py

@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""APIs to convert lowered MLIR from PyTorch to TensorFlow and TFLite artifacts."""
+"""APIs to convert lowered MLIR from PyTorch to TensorFlow artifacts."""
 
 import re
-import tempfile
 
 import tensorflow as tf
 import torch
@@ -104,20 +103,26 @@ def _extract_call_args(
 def _wrap_as_tf_func(lowered, tf_state_dict):
   """Build tf.function from lowered and tf_state_dict."""
 
-  def inner(*args):
+  version = 6
+  if hasattr(tfxla, "call_module_maximum_supported_version"):
+    version = tfxla.call_module_maximum_supported_version()
+
+  def tf_func(*args):
     t_outs = [torch_dtype_to_tf(sig.dtype) for sig in lowered.output_signature]
     s_outs = [_get_shape_with_dynamic(sig) for sig in lowered.output_signature]
     call_args = _extract_call_args(lowered, args, tf_state_dict)
     return tfxla.call_module(
         tuple(call_args),
-        version=5,
+        version=version,
         Tout=t_outs,  # dtype information
-        Sout=s_outs,  # Shape information
+        Sout=s_outs,  # shape information
         function_list=[],
-        module=lowered.module_bytecode,
+        module=lowered.module_bytecode_vhlo,
+        has_token_input_output=False,
+        platforms=["CPU"],
     )
 
-  return inner
+  return tf_func
 
 
 def _make_input_signatures(
@@ -149,46 +154,3 @@ def mlir_to_tf_function(lowered: export.MlirLowered):
       _wrap_as_tf_func(lowered, tf_state_dict),
       input_signature=_make_input_signatures(lowered),
   )
-
-
-def mlir_to_flatbuffer(lowered: export.MlirLowered):
-  """Convert the MLIR lowered to a TFLite flatbuffer binary."""
-  tf_state_dict = _build_tf_state_dict(lowered)
-  signature_names = [tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
-  tf_signatures = [_make_input_signatures(lowered)]
-  tf_functions = [_wrap_as_tf_func(lowered, tf_state_dict)]
-
-  tf_module = tf.Module()
-  tf_module.f = []
-
-  for tf_sig, func in zip(tf_signatures, tf_functions):
-    tf_module.f.append(
-        tf.function(
-            func,
-            input_signature=tf_sig,
-        )
-    )
-
-    tf_module._variables = list(tf_state_dict.values())
-
-  tf_concrete_funcs = [
-      func.get_concrete_function(*tf_sig)
-      for func, tf_sig in zip(tf_module.f, tf_signatures)
-  ]
-
-  # We need to temporarily save since TFLite's from_concrete_functions does not
-  # allow providing names for each of the concrete functions.
-  with tempfile.TemporaryDirectory() as temp_dir_path:
-    tf.saved_model.save(
-        tf_module,
-        temp_dir_path,
-        signatures={
-            sig_name: tf_concrete_funcs[idx]
-            for idx, sig_name in enumerate(signature_names)
-        },
-    )
-
-    converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir_path)
-    tflite_model = converter.convert()
-
-  return tflite_model

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.3.0.dev20241224"
+__version__ = "0.3.0.dev20250105"

{ai_edge_torch_nightly-0.3.0.dev20241224.dist-info → ai_edge_torch_nightly-0.3.0.dev20250105.dist-info}/METADATA

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

@@ -3,9 +3,9 @@ ai_edge_torch/_config.py,sha256=QIrerb6uHMahRvMilmhodJ_6jfiRps3qgLOBeidPnS4,1614
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/fx_pass_base.py,sha256=518ziQ0TUxqum2qZXqlD8qr65pHPh8ZNLnwFC6zvK3k,4253
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=TkfJYt2lJC8A_AcieO1xVmMQ2xdnoTOwF8CZ5dZeaqc,706
+ai_edge_torch/version.py,sha256=rEruohWdKGtxlBLh9SF_NnC4pbAqrOU4MKG598yJRHY,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=SzbR16V2JEfkCjjPwRVAFUbFnzu-_1iHPKgGT9Yz7gQ,5678
+ai_edge_torch/_convert/conversion.py,sha256=_PoH0E1gbbsWhLGwDRwUtW2G_IgNzNF7pKQbn9ct6-4,5778
 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=rGpBNss3Y9FCCCcdBwDo16KqavJi8N5P0M_6Wnmoz6Y,2320
@@ -13,7 +13,7 @@ ai_edge_torch/_convert/to_channel_last_io.py,sha256=_31phf7TYgZY2ftpNbrdlB1RhDiu
 ai_edge_torch/_convert/fx_passes/__init__.py,sha256=NVe-eGcm7j8jZpP2pcMhC8j5dVjgR1pPzyXhHdvKH4E,1267
 ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py,sha256=doaww8KqrgRTD5LotBVAIRFsEqzPn9R5lcGehBJOczA,9098
 ai_edge_torch/_convert/fx_passes/build_interpolate_composite_pass.py,sha256=qb4JBDi4Xca14JJUIcaaZQIJiyqKyHJF49jsRCIFCVA,4335
-ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=IlZuK42kfVcRqAWZp4j2k_81T2uWo9T2558U_GPJAlU,2327
+ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=NxT-iCOHq3r3jeZ8qhNoPXV5w8l2eRMu4yEcBri3NxY,2398
 ai_edge_torch/_convert/fx_passes/remove_non_user_outputs_pass.py,sha256=f1IUVWyhioOClsMiZzLyynoW2R17U83vA-7Q-3pGPM4,2126
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/__init__.py,sha256=lxnoH-WGLeiQIF8XjMGodjiZEFTxucl7g05N7MR9OPk,796
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py,sha256=hDsl9AHzmyuSWsdHOSO114l4nBUgUdAOUWafMTipMgA,7629
@@ -166,14 +166,14 @@ ai_edge_torch/lowertools/_shim.py,sha256=Mbg16tnCVK0YsHowfbpqpNX1qySuMLvpGI_-I5S
 ai_edge_torch/lowertools/common_utils.py,sha256=Z7p-ivOHtddktpnHrlDm_dSoTxJOdEjFXIGQbzjgwQo,4504
 ai_edge_torch/lowertools/odml_torch_utils.py,sha256=dxg2pBuVhSZeY2Ouc0F6nHiZilWZmpSPA7I8kGqSkVI,8282
 ai_edge_torch/lowertools/test_utils.py,sha256=mdxTlhqHABZEQ_GEmPFCL8LIAWtqRtYZUGdSY1ieZjw,1949
-ai_edge_torch/lowertools/torch_xla_utils.py,sha256=tH5BW8-Up1uy5Iq1LdXiJInXBh4-YqNXJpSwwy3kwSg,9460
+ai_edge_torch/lowertools/torch_xla_utils.py,sha256=1EytIw2R6dthhLhf69wN1L9BaQTeybCD0wga-PhHcMI,9518
 ai_edge_torch/lowertools/translate_recipe.py,sha256=ymkBpFqAUiupRWqrPOWiVphKcXR1K5vHK0RjgBFtxlE,5652
 ai_edge_torch/odml_torch/__init__.py,sha256=S8jOzE9nLof-6es3XDiGJRN-9H_XTxsVm9dE7lD3RWo,812
 ai_edge_torch/odml_torch/_torch_future.py,sha256=AJ0klpsbu2ZBTfiZlqSOoaYzBVITt40a1fYN8xKkEPw,3044
 ai_edge_torch/odml_torch/_torch_library.py,sha256=Lw1gqL2HWNRspdTwNhIkYAHDyafHedHtkXyKKxn-Wss,805
-ai_edge_torch/odml_torch/export.py,sha256=QzOPmcNPB7R-KhhPEP0oGVbDRgGPptIxRSoz3S8py9I,13405
+ai_edge_torch/odml_torch/export.py,sha256=sqIMXmxK_qIuVC-_DNJ6wKlIWiXq4_WOCKbSqMRFudg,13293
 ai_edge_torch/odml_torch/export_utils.py,sha256=QeA37Irlty6AiIBuqmHmJgn3lqahBQ5xsh6IKRoKm1g,4774
-ai_edge_torch/odml_torch/tf_integration.py,sha256=lTFJPPEijLPFmn6qq2jbpVTQOo0YaOTK36kK6rCiyIE,5956
+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
@@ -203,8 +203,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.3.0.dev20241224.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241224.dist-info/METADATA,sha256=J8nAtbMNmSIcHuThVv0omkhpldggz91pIIYy-6ATJgM,1966
-ai_edge_torch_nightly-0.3.0.dev20241224.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20241224.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241224.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20250105.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20250105.dist-info/METADATA,sha256=d8fPEhT1HG6ZlbX2joNTeIpEQNqth8LduM_W6aQZQn8,1966
+ai_edge_torch_nightly-0.3.0.dev20250105.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20250105.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20250105.dist-info/RECORD,,