ai-edge-torch-nightly 0.2.0.dev20240725__py3-none-any.whl → 0.2.0.dev20240726__py3-none-any.whl

ai_edge_torch/debug/culprit.py

@@ -354,7 +354,7 @@ def _search_model(
     max_granularity: Optional[int] = None,
     enable_fx_minifier_logging: bool = False,
 ) -> Generator[SearchResult, None, None]:
-  """Finds subgraphs in the torch model that satify a certain predicate function provided by the users.
+  """Finds subgraphs in the torch model that satisfy a certain predicate function provided by the users.
 
   Args:
     predicate_f: a predicate function the users specify.
@@ -382,26 +382,33 @@ def _search_model(
   fx_gm, fx_inputs = utils.exported_program_to_fx_graph_module_and_inputs(ep)
   fx_gm = _normalize_getitem_nodes(fx_gm)
 
-  # HACK: temporarily disable XLA_HLO_DEBUG so that fx_minifier won't dump
-  # intermediate stablehlo files to storage.
+  # HACK: temporarily disable XLA_HLO_DEBUG and create_minified_hlo_graph so that
+  # fx_minifier won't dump intermediate stablehlo files to storage.
   # https://github.com/pytorch/pytorch/blob/main/torch/_functorch/fx_minifier.py#L440
   @contextlib.contextmanager
-  def disable_xla_hlo_debug():
+  def disable_minifier_xla_debug():
     xla_hlo_debug_value = None
     if "XLA_HLO_DEBUG" in os.environ:
       xla_hlo_debug_value = os.environ["XLA_HLO_DEBUG"]
       del os.environ["XLA_HLO_DEBUG"]
 
+    create_minified_hlo_graph = torch._functorch.fx_minifier.create_minified_hlo_graph
+    torch._functorch.fx_minifier.create_minified_hlo_graph = (
+        lambda *args, **kwargs: None
+    )
+
     try:
-      yield None
+      yield
     finally:
       if xla_hlo_debug_value is not None:
         os.environ["XLA_HLO_DEBUG"] = xla_hlo_debug_value
 
+      torch._functorch.fx_minifier.create_minified_hlo_graph = create_minified_hlo_graph
+
   found_culprits_num = 0
   while True:
     try:
-      with disable_xla_hlo_debug(), open(os.devnull, "w") as devnull:
+      with disable_minifier_xla_debug(), open(os.devnull, "w") as devnull:
         with contextlib.nullcontext() if enable_fx_minifier_logging else utils.redirect_stdio(
             stdout=devnull,
             stderr=devnull,

{ai_edge_torch_nightly-0.2.0.dev20240725.dist-info → ai_edge_torch_nightly-0.2.0.dev20240726.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.2.0.dev20240725
+Version: 0.2.0.dev20240726
 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
@@ -27,7 +27,10 @@ Requires-Dist: numpy
 Requires-Dist: scipy
 Requires-Dist: safetensors
 Requires-Dist: tabulate
-Requires-Dist: torch >2.3
+Requires-Dist: torch >=2.4.0
+Requires-Dist: torch-xla >=2.4.0
+Requires-Dist: tf-nightly >=2.18.0.dev20240722
+Requires-Dist: ai-edge-quantizer-nightly ==0.0.1.dev20240718
 
 Library that supports converting PyTorch models into a .tflite format, which can
 then be run with TensorFlow Lite and MediaPipe.  This enables applications for

{ai_edge_torch_nightly-0.2.0.dev20240725.dist-info → ai_edge_torch_nightly-0.2.0.dev20240726.dist-info}/RECORD

@@ -27,7 +27,7 @@ ai_edge_torch/convert/test/test_convert_composites.py,sha256=8UkdPtGkjgSVLCzB_rp
 ai_edge_torch/convert/test/test_convert_multisig.py,sha256=kMaGnHe9ylfyU68qCifYcaGwJqyejKz--QQt9jS2oUA,4537
 ai_edge_torch/convert/test/test_to_channel_last_io.py,sha256=I8c4ZG3v1vo0yxQYzLK_BTId4AOL9vadHGDtfCUZ4UI,2930
 ai_edge_torch/debug/__init__.py,sha256=N05Mmvi41KgSuK0JhuMejERESgP8QekiGdp9_PEyuKU,742
-ai_edge_torch/debug/culprit.py,sha256=urtCKPXORPvn6oyDxDSCSjgvngUnjjcsUMwAOeIl15E,14236
+ai_edge_torch/debug/culprit.py,sha256=lN2N_J3EJTPcEH42xK6kvs4yOqZtyvvecSc4aidlog4,14556
 ai_edge_torch/debug/utils.py,sha256=hjVmQVVl1dKxEF0D6KB4a3ouQ3wBkTsebOX2YsUObZM,1430
 ai_edge_torch/debug/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/debug/test/test_culprit.py,sha256=9An_n9p_RWTAYdHYTCO-__EJlbnjclCDo8tDhOzMlwk,3731
@@ -125,8 +125,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=eO9Ra160ITjQSyRBEGy6nNIVH3gYacSWDd
 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=kzIulTldq8R9E-lAZsvfSTvLu3FYEX7b9DyYM3qisXM,4485
-ai_edge_torch_nightly-0.2.0.dev20240725.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240725.dist-info/METADATA,sha256=fwnI0u8Q4MJNeahUsCYe34o4YaeoRBJrUieJUk5ogY8,1745
-ai_edge_torch_nightly-0.2.0.dev20240725.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240725.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240725.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240726.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240726.dist-info/METADATA,sha256=QstgV90JdvfDL5-OAb_Z4dzWItNf6qT-gNfDtzdoTok,1889
+ai_edge_torch_nightly-0.2.0.dev20240726.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240726.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240726.dist-info/RECORD,,