ai-edge-torch-nightly 0.2.0.dev20240617__py3-none-any.whl → 0.2.0.dev20240618__py3-none-any.whl

ai_edge_torch/debug/__init__.py

@@ -13,4 +13,5 @@
 # limitations under the License.
 # ==============================================================================
 
+from .culprit import _search_model
 from .culprit import find_culprits

ai_edge_torch/debug/culprit.py

@@ -21,7 +21,7 @@ import io
 import operator
 import os
 import sys
-from typing import Any, Generator, List, Optional, Tuple
+from typing import Any, Callable, Generator, List, Optional, Tuple, Union
 
 from functorch.compile import minifier as fx_minifier
 import torch
@@ -85,10 +85,9 @@ def _tensor_to_buffer(t: torch.Tensor):
 
 
 @dataclasses.dataclass
-class Culprit:
+class SearchResult:
   graph_module: torch.fx.GraphModule
   inputs: Tuple[Any]
-  _runtime_errors: bool
 
   @property
   def graph(self) -> torch.fx.Graph:
@@ -98,6 +97,11 @@ class Culprit:
   def graph(self, fx_g: torch.fx.Graph):
     self.graph_module.graph = fx_g
 
+
+@dataclasses.dataclass
+class Culprit(SearchResult):
+  _runtime_errors: bool
+
   @property
   def stack_traces(self) -> List[str]:
     stack_traces = set()
@@ -342,42 +346,42 @@ def _fx_minifier_checker(fx_gm, inputs, runtime_errors=False):
   return False
 
 
-def find_culprits(
-    torch_model: torch.nn.Module,
-    args: Tuple[Any],
-    max_granularity: Optional[int] = None,
-    runtime_errors: bool = False,
+def _search_model(
+    predicate_f: Callable[[torch.fx.GraphModule, List[Any]], bool],
+    model: Union[torch.export.ExportedProgram, torch.nn.Module],
+    export_args: Tuple[Any] = None,
     *,
+    max_granularity: Optional[int] = None,
     enable_fx_minifier_logging: bool = False,
-) -> Generator[Culprit, None, None]:
-  """Finds culprits in the AI Edge Torch model conversion.
+) -> Generator[SearchResult, None, None]:
+  """Finds subgraphs in the torch model that satify a certain predicate function provided by the users.
 
   Args:
-    torch_model: model to export and save
-    args: A set of args to trace the model with, i.e.
-      torch_model(*args) must run
+    predicate_f: a predicate function the users specify.
+      It takes a FX (sub)graph and the inputs to this graph,
+      return True if the graph satisfies the predicate,
+      return False otherwise.
+    model: model in which to search subgraph.
+    export_args: A set of args to trace the model with,
+      i.e. model(*args) must run.
     max_granularity - FX minifier arg. The maximum granularity (number of nodes)
       in the returned ATen FX subgraph of the culprit.
-    runtime_errors: If true, find culprits for Python runtime errors
-      with converted model.
-    enable_fx_minifier_logging: If true, allows the underlying FX minifier to log
-      the progress.
+    enable_fx_minifier_logging: If true, allows the underlying FX minifier to log the progress.
   """
 
-  try:
-    ep = torch.export.export(torch_model, args)
-  except Exception as err:
-    raise ValueError(
-        "Your model is not exportable by torch.export.export. Please modify your model to be torch-exportable first."
-    ) from err
+  if isinstance(model, torch.nn.Module):
+    try:
+      ep = torch.export.export(model, export_args)
+    except Exception as err:
+      raise ValueError(
+          "Your model is not exportable by torch.export.export. Please modify your model to be torch-exportable first."
+      ) from err
+  else:
+    ep = model
 
   fx_gm, fx_inputs = utils.exported_program_to_fx_graph_module_and_inputs(ep)
   fx_gm = _normalize_getitem_nodes(fx_gm)
 
-  fx_minifier_checker = functools.partial(
-      _fx_minifier_checker, runtime_errors=runtime_errors
-  )
-
   # HACK: temporarily disable XLA_HLO_DEBUG 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
@@ -405,13 +409,13 @@ def find_culprits(
           raw_min_fx_gm, raw_min_inputs = fx_minifier(
               fx_gm,
               fx_inputs,
-              fx_minifier_checker,
+              predicate_f,
               max_granularity=max_granularity,
           )
 
       min_fx_gm, min_inputs = _normalize_minified_fx_gm(raw_min_fx_gm, raw_min_inputs)
       found_culprits_num += 1
-      yield Culprit(min_fx_gm, min_inputs, _runtime_errors=runtime_errors)
+      yield SearchResult(min_fx_gm, min_inputs)
 
       fx_gm, fx_inputs = _erase_sub_gm_from_gm(
           fx_gm, fx_inputs, raw_min_fx_gm, raw_min_inputs
@@ -421,3 +425,40 @@ def find_culprits(
       if str(e) == "Input graph did not fail the tester" and found_culprits_num > 0:
         break
       raise e
+
+
+def find_culprits(
+    torch_model: torch.nn.Module,
+    args: Tuple[Any],
+    max_granularity: Optional[int] = None,
+    runtime_errors: bool = False,
+    *,
+    enable_fx_minifier_logging: bool = False,
+) -> Generator[Culprit, None, None]:
+  """Finds culprits in the AI Edge Torch model conversion.
+
+  Args:
+    torch_model: model to export and save
+    args: A set of args to trace the model with, i.e.
+      torch_model(*args) must run
+    max_granularity - FX minifier arg. The maximum granularity (number of nodes)
+      in the returned ATen FX subgraph of the culprit.
+    runtime_errors: If true, find culprits for Python runtime errors
+      with converted model.
+    enable_fx_minifier_logging: If true, allows the underlying FX minifier to log the progress.
+  """
+
+  fx_minifier_checker = functools.partial(
+      _fx_minifier_checker, runtime_errors=runtime_errors
+  )
+
+  for search_result in _search_model(
+      fx_minifier_checker,
+      torch_model,
+      args,
+      max_granularity=max_granularity,
+      enable_fx_minifier_logging=enable_fx_minifier_logging,
+  ):
+    yield Culprit(
+        search_result.graph_module, search_result.inputs, _runtime_errors=runtime_errors
+    )

ai_edge_torch/debug/test/test_search_model.py

@@ -0,0 +1,50 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+
+import unittest
+
+import torch
+
+from ai_edge_torch.debug import _search_model
+
+
+class TestSearchModel(unittest.TestCase):
+
+  def test_search_model_with_ops(self):
+    class MultipleOpsModel(torch.nn.Module):
+
+      def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        sub_0 = x - 1
+        add_0 = y + 1
+        mul_0 = x * y
+        add_1 = sub_0 + add_0
+        mul_1 = add_0 * mul_0
+        sub_1 = add_1 - mul_1
+        return sub_1
+
+    model = MultipleOpsModel().eval()
+    args = (torch.rand(10), torch.rand(10))
+
+    def find_subgraph_with_sub(fx_gm, inputs):
+      return torch.ops.aten.sub.Tensor in [n.target for n in fx_gm.graph.nodes]
+
+    results = list(_search_model(find_subgraph_with_sub, model, args))
+    self.assertEqual(len(results), 2)
+    self.assertIn(torch.ops.aten.sub.Tensor, [n.target for n in results[0].graph.nodes])
+
+
+if __name__ == "__main__":
+  unittest.main()

{ai_edge_torch_nightly-0.2.0.dev20240617.dist-info → ai_edge_torch_nightly-0.2.0.dev20240618.dist-info}/METADATA

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

@@ -24,11 +24,12 @@ ai_edge_torch/convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrK
 ai_edge_torch/convert/test/test_convert.py,sha256=2qPmmGqnfV_o1gfsSdjGq3-JR1b323ligiy5MdAv9NA,8021
 ai_edge_torch/convert/test/test_convert_composites.py,sha256=_Ojc-H6GOS5s8ek3_8eRBL_AiCs-k3srziPJ2R4Ulrg,7255
 ai_edge_torch/convert/test/test_convert_multisig.py,sha256=kMaGnHe9ylfyU68qCifYcaGwJqyejKz--QQt9jS2oUA,4537
-ai_edge_torch/debug/__init__.py,sha256=TKvmnjVk3asvYcVh6C-LPr6srgAF_nppSAupWEXqwPY,707
-ai_edge_torch/debug/culprit.py,sha256=vklaxBUfINdo44OsH7csILK70N41gEThCGchGEfbTZw,12789
+ai_edge_torch/debug/__init__.py,sha256=N05Mmvi41KgSuK0JhuMejERESgP8QekiGdp9_PEyuKU,742
+ai_edge_torch/debug/culprit.py,sha256=urtCKPXORPvn6oyDxDSCSjgvngUnjjcsUMwAOeIl15E,14236
 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
+ai_edge_torch/debug/test/test_search_model.py,sha256=0guAEon5cvwBpPXk6J0wVOKj7TXMDaiuomEEQmHgO5o,1590
 ai_edge_torch/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -109,8 +110,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=EIyKz-HY70DguWuSrJal8LpYXQ5ZSEUf3ZrVl7jikFM,4286
-ai_edge_torch_nightly-0.2.0.dev20240617.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240617.dist-info/METADATA,sha256=Z9rUO2CabVbBpydpRk8OxNlwK4yznGCb2QHGlJhqRsM,1748
-ai_edge_torch_nightly-0.2.0.dev20240617.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240617.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240617.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240618.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240618.dist-info/METADATA,sha256=aMhby_ftyg_8pWf8klYbHTCP7rMDcmuSTeryoRKt4U0,1748
+ai_edge_torch_nightly-0.2.0.dev20240618.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240618.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240618.dist-info/RECORD,,