ai-edge-torch-nightly 0.3.0.dev20241210__py3-none-any.whl → 0.3.0.dev20241211__py3-none-any.whl

ai_edge_torch/debug/test/test_culprit.py

@@ -15,14 +15,14 @@
 
 
 import ast
-import io
-import sys
 
-from ai_edge_torch.debug import find_culprits
+import ai_edge_torch.debug
 import torch
 
 from absl.testing import absltest as googletest
 
+find_culprits = ai_edge_torch.debug.find_culprits
+
 _test_culprit_lib = torch.library.Library("test_culprit", "DEF")
 
 _test_culprit_lib.define("non_lowerable_op(Tensor x) -> Tensor")
@@ -52,6 +52,11 @@ class BadModel(torch.nn.Module):
 
 class TestCulprit(googletest.TestCase):
 
+  def setUp(self):
+    super().setUp()
+    torch.manual_seed(0)
+    torch._dynamo.reset()
+
   def test_find_culprits(self):
     model = BadModel().eval()
     args = (torch.rand(10),)

ai_edge_torch/generative/test/test_model_conversion_large.py

@@ -151,7 +151,7 @@ class TestModelConversion(googletest.TestCase):
   )
   def test_openelm(self):
     config = openelm.get_fake_model_config()
-    pytorch_model = openelm.OpenElm(config).eval()
+    pytorch_model = openelm.OpenELM(config).eval()
     self._test_model(config, pytorch_model, "prefill", atol=1e-4, rtol=1e-5)
 
   @googletest.skipIf(

ai_edge_torch/odml_torch/lowerings/__init__.py

@@ -21,6 +21,6 @@ from . import _quantized_decomposed
 from . import context
 from . import registry
 from . import utils
-from .registry import decompositions
+from .decomp import decompositions
 from .registry import lookup
 from .registry import lower

ai_edge_torch/odml_torch/lowerings/decomp.py

@@ -0,0 +1,59 @@
+# 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.
+# ==============================================================================
+"""Torch export decompositions to run before lowering."""
+
+import functools
+
+import torch
+
+
+@functools.cache
+def decompositions():
+  # Base: Core ATen decompositions
+  decompositions = torch._decomp.core_aten_decompositions()
+
+  decompositions.update(
+      torch._decomp.get_decompositions([
+          torch.ops.aten.upsample_nearest2d,
+          torch.ops.aten._native_batch_norm_legit.no_stats,
+          torch.ops.aten._native_batch_norm_legit_functional,
+          torch.ops.aten._adaptive_avg_pool2d,
+          torch.ops.aten._adaptive_avg_pool3d,
+          torch.ops.aten.grid_sampler_2d,
+          torch.ops.aten.native_group_norm,
+          torch.ops.aten.native_dropout,
+          torch.ops.aten.reflection_pad1d,
+          torch.ops.aten.reflection_pad2d,
+          torch.ops.aten.reflection_pad3d,
+          torch.ops.aten.replication_pad1d,
+          torch.ops.aten.replication_pad2d,
+          torch.ops.aten.replication_pad3d,
+          torch.ops.aten.addmm,
+      ])
+  )
+
+  torch._decomp.remove_decompositions(
+      decompositions,
+      [torch.ops.aten.roll],
+  )
+
+  # Override _safe_softmax decompositions with regular softmax.
+  # _safe_softmax introduces additional check-select ops to guard extreme
+  # input values to softmax, which could make the converted model inefficient
+  # on-device.
+  if hasattr(torch.ops.aten, "_safe_softmax"):
+    decompositions[torch.ops.aten._safe_softmax.default] = torch.softmax
+
+  return decompositions

ai_edge_torch/odml_torch/lowerings/registry.py

@@ -26,7 +26,6 @@ class LoweringRegistry:
 
   def __init__(self):
     self.registered_ops = {}
-    self.decompositions = {}
 
   def lookup(self, op_or_name):
     candidate = self._get_lowering(op_or_name)
@@ -52,33 +51,6 @@ class LoweringRegistry:
 
 
 global_registry = LoweringRegistry()
-global_registry.decompositions.update(torch._decomp.core_aten_decompositions())
-global_registry.decompositions.update(
-    torch._decomp.get_decompositions([
-        torch.ops.aten.upsample_nearest2d,
-        torch.ops.aten._native_batch_norm_legit.no_stats,
-        torch.ops.aten._native_batch_norm_legit_functional,
-        torch.ops.aten._adaptive_avg_pool2d,
-        torch.ops.aten._adaptive_avg_pool3d,
-        torch.ops.aten.grid_sampler_2d,
-        torch.ops.aten.native_group_norm,
-        torch.ops.aten.native_dropout,
-        torch.ops.aten.reflection_pad1d,
-        torch.ops.aten.reflection_pad2d,
-        torch.ops.aten.reflection_pad3d,
-        torch.ops.aten.replication_pad1d,
-        torch.ops.aten.replication_pad2d,
-        torch.ops.aten.replication_pad3d,
-        torch.ops.aten.addmm,
-    ])
-)
-
-torch._decomp.remove_decompositions(
-    global_registry.decompositions,
-    [
-        torch.ops.aten.roll,
-    ],
-)
 
 
 def lookup(op):
@@ -91,7 +63,3 @@ def lower(op):
     return lowering
 
   return inner
-
-
-def decompositions():
-  return global_registry.decompositions

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241210.dist-info → ai_edge_torch_nightly-0.3.0.dev20241211.dist-info}/METADATA

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

@@ -3,7 +3,7 @@ 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=518ziQ0TUxqum2qZXqlD8qr65pHPh8ZNLnwFC6zvK3k,4253
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=AYxcupivW-iYIlCjWXl-QtEvpRsQqFcNK9I6uyGDqaU,706
+ai_edge_torch/version.py,sha256=_uS2Df0H-aUbz-7M-gLxfjDVOJxr03EeNDfbVC_cBrE,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
@@ -34,7 +34,7 @@ ai_edge_torch/debug/__init__.py,sha256=N05Mmvi41KgSuK0JhuMejERESgP8QekiGdp9_PEyu
 ai_edge_torch/debug/culprit.py,sha256=7UYVpVWpiCXbMAyThVtHt_kc_poT7sCTh5UUPvcycgk,14832
 ai_edge_torch/debug/utils.py,sha256=vOAL4t6Lj47uhKapfEsc_WHmvwew3eKO9hSJyzvPXnU,1625
 ai_edge_torch/debug/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/debug/test/test_culprit.py,sha256=SLX4rC-5Dlna8MWHhGRNe72K71AHTFufDrWLlFQn50c,3773
+ai_edge_torch/debug/test/test_culprit.py,sha256=fRN-8jJicawJ2mhPRQNAQUZ8AdGg-s0tYMXyhnLAlWw,3875
 ai_edge_torch/debug/test/test_search_model.py,sha256=-RuU0QsjqkfzZF2IbeA55MoeVOawhbgiSEu96PmioPE,1668
 ai_edge_torch/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -140,7 +140,7 @@ ai_edge_torch/generative/test/test_custom_dus.py,sha256=gxG78CcTpXF3iLzDR15Rlz1e
 ai_edge_torch/generative/test/test_kv_cache.py,sha256=W6Bh0gYDzmwb0j9HdD5_D7Z7FPToP2HSyFrmwIXuFqo,3793
 ai_edge_torch/generative/test/test_loader.py,sha256=9mQUeeZKOVApOWSWl2cN9c10axZjMKM1-0Zd823CCS4,3449
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=i3tQ6mEAo9lCctNoqFAnULk94hgKncC4ywn8IvgbUOo,6341
-ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=IBuvXvORtHu3khr3mLJzYXyCd-zQLUdURTfH28Oo9e0,11079
+ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=MNxJHxDOx8SXi5dJ--JH7KnpJUgtUjwfvCm7X_rnNzA,11079
 ai_edge_torch/generative/test/test_quantize.py,sha256=bEJMhpQ9bIDUZVBXTW888728FcH-i3SyE4JSZZUgU0A,6071
 ai_edge_torch/generative/test/utils.py,sha256=eQ-hjd1eXuHJF3SJK6_CrjgOZVzmG_4VEdH7Z1gH_lA,1897
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
@@ -182,7 +182,7 @@ ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py,sha256=IvOBQyROI9WHS3umHRxsDW
 ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=Jco5zvejxuyl9xHQxZICAKbkgH7x38qPlwUUpD7S15Q,730
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=LqwZ1vCJTSOzgzvH8LUAN-sAkF-l_pGj1AMEIzAqHCA,6638
 ai_edge_torch/odml_torch/jax_bridge/utils.py,sha256=T8isGc896VrHZ6c_L5pYmLpolQ7ibcOlgWfPuVFPzIg,2264
-ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=1lMKPoStK3SUA8yYTPZBRhESN33BghGXnfqOOg4oeVk,995
+ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=0GytV1dGnqe1mKityqQDNFNS8T4QBg3UZuRJcGHwGyA,993
 ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=ufvnaAh6rM_yfoc8ybI3VErHEVBv5W_p4iOe9slfwKM,9948
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=v1VdKmL8YLJv3PR9VgyNghO83A25PpTzY2ZUAJqlq3Q,6847
@@ -190,7 +190,8 @@ ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=4UyNyaR2W-vCOvj-P5ly
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=1ePJs7oIdUkVdMddFsXMc53qTkEKqGz0ZhQQoNzBa10,2862
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=rFmzqcdjYrwhcxH8j9zCFStPy21HFF7hkUV_GQ8FPAk,6056
 ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
-ai_edge_torch/odml_torch/lowerings/registry.py,sha256=itTt8MLbq2LoHTzRidCF2TTbh0TP7L836u99qCjP3FA,2953
+ai_edge_torch/odml_torch/lowerings/decomp.py,sha256=aR6JPFP2Iq-aR0qPxJEHehmAVTjiGhgQEoycZV_1vPY,2130
+ai_edge_torch/odml_torch/lowerings/registry.py,sha256=Tp2h11l5uTifO0aIkuUOWAF_ibEjmd65Xx99w3EXuGE,1924
 ai_edge_torch/odml_torch/lowerings/utils.py,sha256=pqM6mumpviFDHRaabp93CUAngzEZmWcAHl0nTDgyI2g,6167
 ai_edge_torch/odml_torch/passes/__init__.py,sha256=AVwIwUTMx7rXacKjGy4kwrtMd3XB2v_ncdc40KOjUqQ,1245
 ai_edge_torch/quantize/__init__.py,sha256=aB5dXot04bqyUhpsDFvxt9CIi15QAC4euvqOndJ0XLU,714
@@ -200,8 +201,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.dev20241210.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241210.dist-info/METADATA,sha256=SM6aXiKe6YYFKtS0NbSZwwYIdZES74y0X7wautX45S4,1897
-ai_edge_torch_nightly-0.3.0.dev20241210.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20241210.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241210.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241211.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241211.dist-info/METADATA,sha256=Lyub5vadYf6Yu6mGY7l1PFk8Jg2rB36ojIBHm9CxhBM,1897
+ai_edge_torch_nightly-0.3.0.dev20241211.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20241211.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241211.dist-info/RECORD,,