ai-edge-torch-nightly 0.3.0.dev20241105__py3-none-any.whl → 0.3.0.dev20241108__py3-none-any.whl

ai_edge_torch/fx_pass_base.py

@@ -76,7 +76,7 @@ def run_passes(
             range_constraints=exported_program.range_constraints,
             module_call_graph=exported_program.module_call_graph,
             example_inputs=exported_program.example_inputs,
-            verifier=exported_program.verifier,
+            verifiers=exported_program.verifiers,
             constants=exported_program.constants,
         )
   return exported_program

ai_edge_torch/generative/layers/kv_cache.py

@@ -161,6 +161,9 @@ def update(
   Returns:
       KVCacheEntry: The updated KVCache entry based on the passed inputs.
   """
+  # Don't enable HLFB for kv cache op for now, since it won't work with LLM
+  # inference engine. Remove this part once we ship a new LLM inference engine.
+  enable_hlfb=False
   update_func = _update_kv_hlfb_impl if enable_hlfb else _update_kv_base_impl
   return update_func(cache, input_pos, k_slice, v_slice)
 

ai_edge_torch/odml_torch/_torch_future.py

@@ -59,3 +59,17 @@ def graph_module_flat_inputs(ep: torch.export.ExportedProgram, args, kwargs):
     ordered_tensor_constants = tuple()
 
   return (*param_buffer_values, *flat_args, *ordered_tensor_constants)
+
+
+# TODO(b/331481564): Replace this with CanonicalizePass + run_decomposition
+def safe_run_decompositions(exported_program, decomp_table=None):
+  for node in exported_program.graph.nodes:
+    if node.target == torch.ops.aten.view.default:
+      # Passes or torch.export may generate aten.view nodes not respecting the
+      # tensor memory format. Changes all the aten.view to torch.reshape
+      # for retracing. If the input memory format is already contiguous,
+      # retracing in run_decomposition below would decompose torch.reshape
+      # back to one aten.view.
+      node.target = lambda self, size: torch.reshape(self.contiguous(), size)
+
+  return exported_program.run_decompositions(decomp_table)

ai_edge_torch/odml_torch/export.py

@@ -262,13 +262,13 @@ def exported_program_to_mlir(
     exported_program: torch.export.ExportedProgram,
 ) -> MlirLowered:
   """Lower the exported program to MLIR."""
-  exported_program = exported_program.run_decompositions(
-      lowerings.decompositions()
+  exported_program = _torch_future.safe_run_decompositions(
+      exported_program, lowerings.decompositions()
   )
 
   _convert_i64_to_i32(exported_program)
-  exported_program = exported_program.run_decompositions(
-      lowerings.decompositions()
+  exported_program = _torch_future.safe_run_decompositions(
+      exported_program, lowerings.decompositions()
   )
 
   with export_utils.create_ir_context() as context, ir.Location.unknown():

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241105.dist-info → ai_edge_torch_nightly-0.3.0.dev20241108.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20241105
+Version: 0.3.0.dev20241108
 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
@@ -29,7 +29,7 @@ Requires-Dist: safetensors
 Requires-Dist: tabulate
 Requires-Dist: torch>=2.4.0
 Requires-Dist: torch-xla>=2.4.0
-Requires-Dist: tf-nightly==2.19.0.dev20241001
+Requires-Dist: tf-nightly>=2.19.0.dev20241001
 Requires-Dist: ai-edge-litert-nightly
 Requires-Dist: ai-edge-quantizer-nightly
 

{ai_edge_torch_nightly-0.3.0.dev20241105.dist-info → ai_edge_torch_nightly-0.3.0.dev20241108.dist-info}/RECORD

@@ -1,9 +1,9 @@
 ai_edge_torch/__init__.py,sha256=48qP37uHT90YPs4eIUQxCiWVwqGEX3idCUs6mQKvX1U,1168
 ai_edge_torch/config.py,sha256=FMWeCH2b7HYILBvaI1iZNnYCO4WAhDOwBZBmIE-xrF0,909
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
-ai_edge_torch/fx_pass_base.py,sha256=SrYveglaiA_DXPoRBqSXClWM1q7853I5ujRorq_MV0M,4251
+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=9qA2TGGeU3t-2-SPdZU_51pP6PVMRAJtkNL4Tg2QEZ4,706
+ai_edge_torch/version.py,sha256=sBOl0mgVPJtokiP8qTbTtY0R_qIaF0KNiALh7P3AJEk,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_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -109,7 +109,7 @@ ai_edge_torch/generative/layers/attention.py,sha256=Z0Y_G8IG0LmvLX2u9D8__Fkr22sz
 ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
 ai_edge_torch/generative/layers/builder.py,sha256=Z5LyzCEThgnYZeyViakaE3yJVzTGHtw13acHsAQR15U,5050
 ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
-ai_edge_torch/generative/layers/kv_cache.py,sha256=2El7kZYnQRCRcVc63xgiAdBh9oVOksDu35p9XggvaGE,6148
+ai_edge_torch/generative/layers/kv_cache.py,sha256=lbm-yJ1jGPtcgWS4C3FmSnB1IlxqDE7g0BLRh3PN4N4,6324
 ai_edge_torch/generative/layers/model_config.py,sha256=DdsdhTP5tZAtyWim-qW2m8HDBsYbs7boqSDb83vwgmE,6998
 ai_edge_torch/generative/layers/normalization.py,sha256=eKAGst9rPuyRFExMcQFJO7R3iHdCtlmjeF_lITjLhwE,6498
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=CZqOoibLcHvUgrgaIIWAlmk3XgE2inzx340MN-npLoU,1347
@@ -155,9 +155,9 @@ ai_edge_torch/lowertools/test_utils.py,sha256=bPgc2iXX16KYtMNvmsRdKfrCY6UJmcfitf
 ai_edge_torch/lowertools/torch_xla_utils.py,sha256=XGZE0vZG9WSQT-6dFmPlU8W89z8rfXPRGjuZeuhXCIw,9205
 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=jSYHf1CMTJzMizPMbu2b39hAt0ZTR6gQLq67GMe9KTo,2336
+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=q6cvB3RAKL7hlQFNGqsE3u-NqWWPSzj-8M38u8loSpk,11544
+ai_edge_torch/odml_torch/export.py,sha256=4xwrsDeOAgzoB9m7EeNsBj6dC5Ajtn5aKDRQkdHxa-o,11584
 ai_edge_torch/odml_torch/export_utils.py,sha256=q84U69ZQ82hLXw-xncJ8IW-K71Xux-NWlzZTs7hdZWA,5127
 ai_edge_torch/odml_torch/tf_integration.py,sha256=lTFJPPEijLPFmn6qq2jbpVTQOo0YaOTK36kK6rCiyIE,5956
 ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xKA9rknO93JS9kw9w_c,778
@@ -186,8 +186,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.dev20241105.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241105.dist-info/METADATA,sha256=Lmr3JaJXqa10_MhDYy3XZAd7u5PTpg8E3mW8IhWRLQ8,1897
-ai_edge_torch_nightly-0.3.0.dev20241105.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20241105.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241105.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/METADATA,sha256=gp2VN_X4YPdK8axZYIhqafgiJhCwfiN_tOWT-yL3lW0,1897
+ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/RECORD,,