ai-edge-torch-nightly 0.3.0.dev20241216__py3-none-any.whl → 0.3.0.dev20241217__py3-none-any.whl

ai_edge_torch/_convert/conversion.py

@@ -78,7 +78,8 @@ def convert_signatures(
     *,
     strict_export: Union[Literal["auto"], bool] = True,
     quant_config: Optional[qcfg.QuantConfig] = None,
-    _tfl_converter_flags: Optional[dict[str, Any]],
+    _tfl_converter_flags: Optional[dict[str, Any]] = None,
+    _saved_model_dir: Optional[str] = None,
 ) -> model.TfLiteModel:
   """Converts a list of `signature.Signature`s and embeds them into one `model.TfLiteModel`.
 
@@ -93,6 +94,8 @@ def convert_signatures(
       quant_config: User-defined quantization method and scheme of the model.
       _tfl_converter_flags: A nested dictionary allowing setting flags for the
         underlying tflite converter.
+      _saved_model_dir: Directory for the intermediate saved model. If not
+        specified, a random temporary directory would be used.
 
   Returns:
     The converted `model.TfLiteModel` object.
@@ -140,6 +143,7 @@ def convert_signatures(
       signatures,
       quant_config=quant_config,
       _tfl_converter_flags=_tfl_converter_flags,
+      _saved_model_dir=_saved_model_dir,
   )
 
   return model.TfLiteModel(tflite_model)

ai_edge_torch/_convert/converter.py

@@ -106,6 +106,7 @@ class Converter:
       quant_config: Optional[qcfg.QuantConfig] = None,
       dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
       _ai_edge_converter_flags: Optional[dict[str, Any]] = None,
+      _saved_model_dir: Optional[str] = None,
   ) -> model.TfLiteModel:
     """Finalizes the conversion and produces an edge model.
 
@@ -139,6 +140,8 @@ class Converter:
         of this function and so needs to be treated as such. Please do not rely
         on this parameter except for local debugging as this can be removed in a
         future release.
+      _saved_model_dir: Directory for the intermediate saved model. If not
+        specified, a random temporary directory would be used.
 
     Returns:
       The converted edge model.
@@ -171,6 +174,7 @@ class Converter:
         strict_export=strict_export,
         quant_config=quant_config,
         _tfl_converter_flags=_ai_edge_converter_flags,
+        _saved_model_dir=_saved_model_dir,
     )
 
 
@@ -216,6 +220,7 @@ def convert(
     quant_config: Optional[qcfg.QuantConfig] = None,
     dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
     _ai_edge_converter_flags: Optional[dict[str, Any]] = None,
+    _saved_model_dir: Optional[str] = None,
 ) -> model.TfLiteModel:
   """Converts a PyTorch model to an edge model with a default signature.
 
@@ -240,6 +245,8 @@ def convert(
       this function and so needs to be treated as such. Please do not rely on
       this parameter except for local debugging as this can be removed in a
       future release.
+    _saved_model_dir: Directory for the intermediate saved model. If not
+      specified, a random temporary directory would be used.
 
   Returns:
     The converted edge model.
@@ -259,4 +266,5 @@ def convert(
       quant_config=quant_config,
       dynamic_shapes=dynamic_shapes,
       _ai_edge_converter_flags=_ai_edge_converter_flags,
+      _saved_model_dir=_saved_model_dir,
   )

ai_edge_torch/lowertools/_shim.py

@@ -50,6 +50,7 @@ def exported_programs_to_tflite(
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
     _tfl_converter_flags: Optional[dict[str, Any]] = None,
+    _saved_model_dir: Optional[str] = None
 ):
   """Converts a list of ExportedProgram to a TFLite model.
 
@@ -57,6 +58,8 @@ def exported_programs_to_tflite(
     exported_programs: A list of ExportedProgram.
     signatures: A list of Signature.
     quant_config: A QuantConfig.
+    _saved_model_dir: Directory for the intermediate saved model. If not
+      specified, a random temporary directory would be used.
     _tfl_converter_flags: A dict of flags for TFLiteConverter.
 
   Returns:
@@ -79,4 +82,5 @@ def exported_programs_to_tflite(
       signatures,
       quant_config=quant_config,
       _tfl_converter_flags=_tfl_converter_flags,
+      _saved_model_dir=_saved_model_dir,
   )

ai_edge_torch/lowertools/odml_torch_utils.py

@@ -138,6 +138,7 @@ def merged_bundle_to_tfl_model(
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
     _tfl_converter_flags: dict = {},
+    _saved_model_dir: Optional[str] = None,
 ):
   tf_state_dict = merged_bundle.bundles[0].state_dict
 
@@ -173,6 +174,9 @@ def merged_bundle_to_tfl_model(
   # 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:
+    if _saved_model_dir is not None:
+      temp_dir_path = _saved_model_dir
+
     tf.saved_model.save(
         tf_module,
         temp_dir_path,

ai_edge_torch/lowertools/torch_xla_utils.py

@@ -192,6 +192,7 @@ def merged_bundle_to_tfl_model(
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
     _tfl_converter_flags: dict = {},
+    _saved_model_dir: Optional[str] = None,
 ) -> None:
   """Converts a StableHLOGraphModule to a tflite model.
 
@@ -200,6 +201,8 @@ def merged_bundle_to_tfl_model(
     signatures: List of signatures from which names of the signatures is
     extracted.
     quant_config: User-defined quantization method and scheme of the model.
+    _saved_model_dir: Directory for the intermediate saved model. If not
+      specified, a random temporary directory would be used.
     _tfl_converter_flags: A nested dictionary allowing setting flags for the
     underlying tflite converter.
   """
@@ -246,6 +249,9 @@ def merged_bundle_to_tfl_model(
   # 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:
+    if _saved_model_dir is not None:
+      temp_dir_path = _saved_model_dir
+
     tf.saved_model.save(
         tf_module,
         temp_dir_path,

ai_edge_torch/odml_torch/export.py

@@ -304,9 +304,13 @@ def exported_program_to_mlir(
   )
 
   _convert_i64_to_i32(exported_program)
+
   exported_program = _torch_future.safe_run_decompositions(
       exported_program, lowerings.decompositions()
   )
+
+  # Passes below mutate the exported program to a state not executable by torch.
+  # Do not call run_decompositions after applying the passes.
   _convert_q_dq_per_channel_args_to_list(exported_program)
 
   with export_utils.create_ir_context() as context, ir.Location.unknown():

ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py

@@ -52,10 +52,13 @@ def _uniform_quantized_type(
     assert isinstance(scale, (list, tuple))
     assert isinstance(zero_point, (list, tuple))
 
+    scale = list(scale)
+    zero_point = list(zero_point)
+
     if len(scale) == 1:
-      scale *= channel_axis_size
+      scale = scale * channel_axis_size
     if len(zero_point) == 1:
-      zero_point *= channel_axis_size
+      zero_point = zero_point * channel_axis_size
 
     assert len(scale) == len(zero_point) == channel_axis_size
     scale_zp_strs = []

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241216.dist-info → ai_edge_torch_nightly-0.3.0.dev20241217.dist-info}/METADATA

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

@@ -3,11 +3,11 @@ 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=Ghtv10LCkTCFTvDn5IHArtqXgxeWVsVxoZa6YIVIXtA,706
+ai_edge_torch/version.py,sha256=D9ZffEHjGHmublg0LV01j677usvnse7YN7pG7upZoNw,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=HwzfRx_DX5TLtPqwEH1_NOm38_INvHzHl4_mX67KOdQ,5448
+ai_edge_torch/_convert/conversion.py,sha256=SzbR16V2JEfkCjjPwRVAFUbFnzu-_1iHPKgGT9Yz7gQ,5678
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
-ai_edge_torch/_convert/converter.py,sha256=DYbTZMZos8bvm9mLyDv3W1P8ER_iGKVohbFAmLZD4r8,9534
+ai_edge_torch/_convert/converter.py,sha256=075F8LRewk_033Ebsnft7FJr3KgtIbtZ_-8udIPy6ho,9980
 ai_edge_torch/_convert/signature.py,sha256=rGpBNss3Y9FCCCcdBwDo16KqavJi8N5P0M_6Wnmoz6Y,2320
 ai_edge_torch/_convert/to_channel_last_io.py,sha256=_31phf7TYgZY2ftpNbrdlB1RhDium1lz_BXEQ6IsMFc,2893
 ai_edge_torch/_convert/fx_passes/__init__.py,sha256=NVe-eGcm7j8jZpP2pcMhC8j5dVjgR1pPzyXhHdvKH4E,1267
@@ -160,16 +160,16 @@ ai_edge_torch/hlfb/mark_pattern/pattern.py,sha256=NP2mYhe5D2GjtqQfqqldp-ko3xtNgh
 ai_edge_torch/hlfb/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/hlfb/test/test_mark_pattern.py,sha256=ivq0eVjuf31idfNY0E12F4FxdkSI9hwYXapLJBkIf8Q,4831
 ai_edge_torch/lowertools/__init__.py,sha256=A8WBXvWtuFYYWtNTqPD7waVntLaSVAnSMwx5ugjZBIw,761
-ai_edge_torch/lowertools/_shim.py,sha256=xJIHDSWNoF4PkkT0JkjeJxgguQ9JGEwooJf9xZNkVRU,3058
+ai_edge_torch/lowertools/_shim.py,sha256=Mbg16tnCVK0YsHowfbpqpNX1qySuMLvpGI_-I5SIrG0,3276
 ai_edge_torch/lowertools/common_utils.py,sha256=Z7p-ivOHtddktpnHrlDm_dSoTxJOdEjFXIGQbzjgwQo,4504
-ai_edge_torch/lowertools/odml_torch_utils.py,sha256=Smt7p62-lZ_3bBBfnbssAK5GAGxm3U_X7M-1qwsmc68,8161
+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=XGZE0vZG9WSQT-6dFmPlU8W89z8rfXPRGjuZeuhXCIw,9205
+ai_edge_torch/lowertools/torch_xla_utils.py,sha256=tH5BW8-Up1uy5Iq1LdXiJInXBh4-YqNXJpSwwy3kwSg,9460
 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=dgnNGBVkHBz0brlWALX2hGXpQ4YzCKdwbkF4oAfEu4I,13062
+ai_edge_torch/odml_torch/export.py,sha256=Wc_JM7U2IjZeBmXA6t1AZxREGOWjZ6EB-PIhEevWWeU,13207
 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/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xKA9rknO93JS9kw9w_c,778
@@ -187,7 +187,7 @@ ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=Q0aDzyUcZMoSzSbOU-r3LJMgPe6fble0QwdYVIOHHHk,6887
 ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=OVmlPGwyhDXKhmG4SAeEsa6iLpJHEHV_jKqwfjYvetA,11643
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=khJIvDVk2s332Nd2Be-5dM6-wp5DGff61HCV5lskHmQ,3011
-ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=GEs83mtEjh8GOW_OATI_ur11VKujrOL2xdZeZ0l1HtM,6100
+ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQDJnaAAxhNFAd7NfQm5SOVEp_nno_krA,6178
 ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
 ai_edge_torch/odml_torch/lowerings/decomp.py,sha256=UoJeZVcr4zAN_11i-HzfOhxGCxUm-7b1JXPVBxR2hSs,2414
 ai_edge_torch/odml_torch/lowerings/registry.py,sha256=Tp2h11l5uTifO0aIkuUOWAF_ibEjmd65Xx99w3EXuGE,1924
@@ -200,8 +200,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.dev20241216.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241216.dist-info/METADATA,sha256=ofd2acZU87DZojjV3S5U4Uw9Pzm4NRB_2WQL-omfbFY,1966
-ai_edge_torch_nightly-0.3.0.dev20241216.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20241216.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241216.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241217.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241217.dist-info/METADATA,sha256=BG9KpWduInLQP8oEHFqxRogKnZecqzm5E536x6RFmcE,1966
+ai_edge_torch_nightly-0.3.0.dev20241217.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20241217.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241217.dist-info/RECORD,,