ai-edge-torch-nightly 0.3.0.dev20241213__py3-none-any.whl → 0.3.0.dev20241214__py3-none-any.whl

ai_edge_torch/__init__.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 # ==============================================================================
 
+from ai_edge_torch._config import config
 from ai_edge_torch._convert.converter import convert
 from ai_edge_torch._convert.converter import signature
 from ai_edge_torch._convert.to_channel_last_io import to_channel_last_io
 from ai_edge_torch.model import Model
 from ai_edge_torch.version import __version__
 
-
 def load(path: str) -> Model:
   """Imports an ai_edge_torch model from disk.
 

ai_edge_torch/_config.py

@@ -0,0 +1,52 @@
+# 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.
+# ==============================================================================
+
+"""Provides a configuration for the ai-edge-torch."""
+
+import functools
+import logging
+import os
+
+__all__ = ["config"]
+
+
+class _Config:
+  """ai-edge-torch global configs."""
+
+  @property
+  @functools.cache  # pylint: disable=method-cache-max-size-none
+  def use_torch_xla(self) -> bool:
+    """True if using torch_xla to lower torch ops to StableHLO.
+
+    To use torch_xla as the lowering backend, set environment variable
+    `USE_TORCH_XLA` to "true".
+    """
+    var = os.environ.get("USE_TORCH_XLA", "false")
+    var = var.lower().strip()
+    if var in ("y", "yes", "t", "true", "on", "1"):
+      return True
+    elif var in ("n", "no", "f", "false", "off", "0"):
+      return False
+    else:
+      logging.warning("Invalid USE_TORCH_XLA value is ignored: %s.", var)
+      return False
+
+  @property
+  def in_oss(self) -> bool:
+    """True if the code is not running in google internal environment."""
+    return True
+
+
+config = _Config()

ai_edge_torch/_convert/test/test_convert.py

@@ -19,7 +19,6 @@ import os
 from typing import Tuple
 
 import ai_edge_torch
-from ai_edge_torch import config
 from ai_edge_torch._convert import conversion_utils
 from ai_edge_torch.quantize import pt2e_quantizer
 from ai_edge_torch.testing import model_coverage
@@ -292,7 +291,7 @@ class TestConvert(googletest.TestCase):
     self.assertTrue(result)
 
   @googletest.skipIf(
-      not config.Config.use_torch_xla,
+      not ai_edge_torch.config.use_torch_xla,
       reason="Shape polymorphism is not yet support with odml_torch.",
   )
   def test_convert_model_with_dynamic_batch(self):

ai_edge_torch/generative/layers/kv_cache.py

@@ -20,6 +20,7 @@ from typing import List, Tuple
 
 from ai_edge_torch import hlfb
 from ai_edge_torch.generative.layers import model_config
+from ai_edge_torch.generative.utilities.dynamic_update_slice import dynamic_update_slice
 import torch
 import torch.utils._pytree as pytree
 
@@ -159,8 +160,6 @@ def update(
   Returns:
       KVCacheEntry: The updated KVCache entry based on the passed inputs.
   """
-  # Turn dynamic_update_slice updates off for now.
-  use_dus=False
   update_kv_cache = _update_kv_impl if use_dus else _update_kv_base_impl
   return update_kv_cache(cache, input_pos, k_slice, v_slice)
 

ai_edge_torch/generative/test/test_model_conversion.py

@@ -16,7 +16,6 @@
 """Testing model conversion for a few gen-ai models."""
 
 import ai_edge_torch
-from ai_edge_torch import config as ai_edge_config
 from ai_edge_torch.generative.examples.test_models import toy_model_with_kv_cache
 from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
 from ai_edge_torch.generative.layers import kv_cache
@@ -83,22 +82,22 @@ class TestModelConversion(googletest.TestCase):
     )
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_toy_model_with_kv_cache(self):
     self._test_model_with_kv_cache(enable_hlfb=False)
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_toy_model_with_kv_cache_with_hlfb(self):
     self._test_model_with_kv_cache(enable_hlfb=True)
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_toy_model_has_dus_op(self):
     """Tests that the model has the dynamic update slice op."""
@@ -179,8 +178,8 @@ class TestModelConversion(googletest.TestCase):
     )
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_tiny_llama_multisig(self):
     config = tiny_llama.get_fake_model_config()

ai_edge_torch/generative/test/test_model_conversion_large.py

@@ -16,7 +16,6 @@
 """Testing model conversion for a few gen-ai models."""
 
 import ai_edge_torch
-from ai_edge_torch import config as ai_edge_config
 from ai_edge_torch.generative.examples.amd_llama_135m import amd_llama_135m
 from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.examples.gemma import gemma2
@@ -91,8 +90,8 @@ class TestModelConversion(googletest.TestCase):
     )
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_gemma1(self):
     config = gemma1.get_fake_model_config()
@@ -100,8 +99,8 @@ class TestModelConversion(googletest.TestCase):
     self._test_model(config, pytorch_model, "prefill", atol=1e-3, rtol=1e-5)
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_gemma2(self):
     config = gemma2.get_fake_model_config()
@@ -109,8 +108,8 @@ class TestModelConversion(googletest.TestCase):
     self._test_model(config, pytorch_model, "prefill", atol=1e-4, rtol=1e-5)
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_llama(self):
     config = llama.get_fake_model_config()
@@ -118,8 +117,8 @@ class TestModelConversion(googletest.TestCase):
     self._test_model(config, pytorch_model, "prefill", atol=1e-3, rtol=1e-5)
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_phi2(self):
     config = phi2.get_fake_model_config()
@@ -128,8 +127,8 @@ class TestModelConversion(googletest.TestCase):
     self._test_model(config, pytorch_model, "prefill", atol=1e-3, rtol=1e-5)
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_phi3(self):
     config = phi3.get_fake_model_config()
@@ -137,8 +136,8 @@ class TestModelConversion(googletest.TestCase):
     self._test_model(config, pytorch_model, "prefill", atol=1e-5, rtol=1e-5)
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_smollm(self):
     config = smollm.get_fake_model_config()
@@ -146,8 +145,8 @@ class TestModelConversion(googletest.TestCase):
     self._test_model(config, pytorch_model, "prefill", atol=1e-4, rtol=1e-5)
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_openelm(self):
     config = openelm.get_fake_model_config()
@@ -155,8 +154,8 @@ class TestModelConversion(googletest.TestCase):
     self._test_model(config, pytorch_model, "prefill", atol=1e-4, rtol=1e-5)
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_qwen(self):
     config = qwen.get_fake_model_config()
@@ -164,8 +163,8 @@ class TestModelConversion(googletest.TestCase):
     self._test_model(config, pytorch_model, "prefill", atol=1e-3, rtol=1e-5)
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_amd_llama_135m(self):
     config = amd_llama_135m.get_fake_model_config()
@@ -173,8 +172,8 @@ class TestModelConversion(googletest.TestCase):
     self._test_model(config, pytorch_model, "prefill", atol=1e-5, rtol=1e-5)
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def disabled_test_paligemma(self):
     config = paligemma.get_fake_model_config()
@@ -222,8 +221,8 @@ class TestModelConversion(googletest.TestCase):
     )
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_stable_diffusion_clip(self):
     config = sd_clip.get_fake_model_config()
@@ -254,8 +253,8 @@ class TestModelConversion(googletest.TestCase):
     )
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_stable_diffusion_diffusion(self):
     config = sd_diffusion.get_fake_model_config(2)
@@ -296,8 +295,8 @@ class TestModelConversion(googletest.TestCase):
     )
 
   @googletest.skipIf(
-      ai_edge_config.Config.use_torch_xla,
-      reason="tests with custom ops are not supported on oss",
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
   )
   def test_stable_diffusion_decoder(self):
     config = sd_decoder.get_fake_model_config()

ai_edge_torch/lowertools/_shim.py

@@ -15,13 +15,15 @@
 
 from typing import Any, Optional
 
-from ai_edge_torch import config
+from ai_edge_torch import _config
 from ai_edge_torch._convert import signature
 from ai_edge_torch.quantize import quant_config as qcfg
 import torch
 
+config = _config.config
+
 # isort: off
-if config.Config.use_torch_xla:
+if config.use_torch_xla:
   from ai_edge_torch.lowertools import torch_xla_utils as utils
   from ai_edge_torch.lowertools.torch_xla_utils import exported_program_to_mlir_text
   from torch_xla.experimental.mark_pattern_utils import StableHLOCompositeBuilder

ai_edge_torch/lowertools/test_utils.py

@@ -15,9 +15,11 @@
 
 import re
 from typing import Optional
-from ai_edge_torch import config
+from ai_edge_torch import _config
 from absl.testing import absltest as googletest
 
+config = _config.config
+
 
 def _extract_backend_configs(mlir):
   mlir = mlir.replace("\\22", '"')
@@ -38,7 +40,7 @@ def assert_string_count(
   if odml_torch_attr_counter is None:
     odml_torch_attr_counter = {}
 
-  if config.Config.use_torch_xla:
+  if config.use_torch_xla:
     for key in torch_xla_pattern_counter:
       test_case.assertEqual(
           mlir.count(key),

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -276,11 +276,13 @@ def _aten_slice_scatter(lctx, self, src, dim=0, start=None, end=None, step=1):
           interior_padding if i == dim else 0 for i in range(rank)
       ],
   )
-  pred = np.ones(self.type.shape, dtype=np.bool_)
-  pred[*[
+
+  slices = [
       slice(start, end, step) if i == dim else slice(None, None, None)
       for i in range(rank)
-  ]] = False
+  ]
+  pred = np.ones(self.type.shape, dtype=np.bool_)
+  pred[np.index_exp[tuple(slices)]] = False
   pred = stablehlo.constant(
       ir.DenseElementsAttr.get(
           np.packbits(pred, bitorder="little"),

ai_edge_torch/odml_torch/lowerings/_convolution.py

@@ -232,7 +232,9 @@ def _aten_convolution(
 
   if bias is not None:
     # broadcast [C] to [NCHW]
-    broadcasted_bias = stablehlo.broadcast_in_dim(output_type, bias, [1])
+    broadcasted_bias = stablehlo.broadcast_in_dim(
+        output_type, bias, ir.DenseI64ArrayAttr.get([1])
+    )
     res = stablehlo.add(
         lhs=res,
         rhs=broadcasted_bias,

ai_edge_torch/odml_torch/lowerings/_layer_norm.py

@@ -20,6 +20,7 @@ from ai_edge_torch.odml_torch.lowerings import registry
 from ai_edge_torch.odml_torch.lowerings import utils
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import hlo as stablehlo
+import numpy as np
 import torch
 
 
@@ -66,12 +67,20 @@ def _aten_native_layer_norm(
   normalized_rank = len(normalized_shape)
   if weight is not None:
     weight = stablehlo.broadcast_in_dim(
-        data_type, weight, list(range(data_rank - normalized_rank, data_rank))
+        data_type,
+        weight,
+        ir.DenseI64ArrayAttr.get(
+            list(range(data_rank - normalized_rank, data_rank))
+        ),
     )
     output = stablehlo.multiply(weight, output)
   if bias is not None:
     bias = stablehlo.broadcast_in_dim(
-        data_type, bias, list(range(data_rank - normalized_rank, data_rank))
+        data_type,
+        bias,
+        ir.DenseI64ArrayAttr.get(
+            list(range(data_rank - normalized_rank, data_rank))
+        ),
     )
     output = stablehlo.add(bias, output)
 

ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 """Lowerings for PT2E torch.ops.quantized_decomposed ops."""
-from typing import Union, cast
+from typing import Optional, Union, cast
 
 from ai_edge_torch.odml_torch.lowerings import context
 from ai_edge_torch.odml_torch.lowerings import utils
@@ -30,15 +30,15 @@ LoweringContext = context.LoweringContext
 
 
 def _uniform_quantized_type(
-    stored_type: str | ir.Type,
-    expressed_type: str | ir.Type,
+    stored_type: Union[str, ir.Type],
+    expressed_type: Union[str, ir.Type],
     *,
-    scale=float | list[float] | tuple[float],
-    zero_point=float | list[float] | tuple[float],
-    storage_type_min: int | None = None,
-    storage_type_max: int | None = None,
-    channel_axis: int | None = None,
-    channel_axis_size: int | None = None,
+    scale=Union[float, list[float], tuple[float]],
+    zero_point=Union[float, list[float], tuple[float]],
+    storage_type_min: Optional[int] = None,
+    storage_type_max: Optional[int] = None,
+    channel_axis: Optional[int] = None,
+    channel_axis_size: Optional[int] = None,
 ):
   """Polyfill for quant.UniformQuantizedType."""
   if storage_type_min and storage_type_max:

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241213.dist-info → ai_edge_torch_nightly-0.3.0.dev20241214.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20241213
+Version: 0.3.0.dev20241214
 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
@@ -11,7 +11,6 @@ Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Topic :: Scientific/Engineering
@@ -20,7 +19,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development
 Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy
@@ -28,10 +27,13 @@ Requires-Dist: scipy
 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.dev20241201
 Requires-Dist: ai-edge-litert-nightly
 Requires-Dist: ai-edge-quantizer-nightly
+Requires-Dist: jax
+Requires-Dist: torch-xla2[odml]>=0.0.1.dev20241201
+Provides-Extra: torch-xla
+Requires-Dist: torch-xla>=2.4.0; extra == "torch-xla"
 
 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.3.0.dev20241213.dist-info → ai_edge_torch_nightly-0.3.0.dev20241214.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/__init__.py,sha256=rq9ZtMJLG8yYNC4tNE4rpl94UAUClZW7f4GAr6HBVDQ,1208
+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=ZKmpBJjKnl93nlPyg2RYz15citIW0ntqZ-0diRjwTt8,706
+ai_edge_torch/version.py,sha256=iCH8lnlOrtbGwvxnT3knpY_keeu2UnrJ_ZXNK2LSvf4,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
@@ -26,7 +26,7 @@ ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitio
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py,sha256=L_x8BrF7UDah-SYl-pG11I6CIckdU9kBTUHcmwW4cts,2420
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=mzfL9cf0qBnpmxM_OlMQFvQsEZV2B_Mia9yEJV4J7rI,7135
 ai_edge_torch/_convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/test/test_convert.py,sha256=v6AhfWqRBuHT7uBDueTbntaQtDSMMrvQOqlIDXNUaMA,17250
+ai_edge_torch/_convert/test/test_convert.py,sha256=gK9QJuLbpjXt0l6tVnzl9Miq6GLkJR-hB67i3VE13Og,17224
 ai_edge_torch/_convert/test/test_convert_composites.py,sha256=BCIODgxMI_3MxMLfNWYMGjcz-al-J3z5eDHCiZJXNwY,7992
 ai_edge_torch/_convert/test/test_convert_multisig.py,sha256=6_C2R9--KyNR7_oezZIAfyTSR97tOeEWy4XGcbSxBDE,5778
 ai_edge_torch/_convert/test/test_to_channel_last_io.py,sha256=1o-gUiwzIuO67FNAJ8DeyKv8fVUeZVNNNwofNVDjYeU,3024
@@ -119,7 +119,7 @@ ai_edge_torch/generative/layers/attention.py,sha256=aOoVM1hY7qjvzVQI1-m26p_f9qoT
 ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
 ai_edge_torch/generative/layers/builder.py,sha256=LXGuSHIx6QZAzLFm7aJvlzoMPgQwbXLFchGEKYwOOUA,5090
 ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
-ai_edge_torch/generative/layers/kv_cache.py,sha256=dOhk3ec21189uPyCDYyxuznYQL6s4od-ln-FoDQ2cE0,6269
+ai_edge_torch/generative/layers/kv_cache.py,sha256=DhHIggaOQ2IAY4aRuMAuCLWZv1dBz5PYtmOEjkx9EQY,6291
 ai_edge_torch/generative/layers/model_config.py,sha256=viX51T_naJ9sPpPxPoMnSueBPYE2zxWNOD0xn0f-_bM,7510
 ai_edge_torch/generative/layers/normalization.py,sha256=h2btgRHMMjOcyLm8adEmcT0pG6imq4QcWblKJK5MYXA,7479
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=xxWtlVsGGJkEyXC6PwznubyhJnLPEfSpHOORE_hgxss,2670
@@ -139,8 +139,8 @@ ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudj
 ai_edge_torch/generative/test/test_custom_dus.py,sha256=gxG78CcTpXF3iLzDR15Rlz1ey1tNTlSdkp6TeYEijp0,3301
 ai_edge_torch/generative/test/test_kv_cache.py,sha256=2AulHBS3hC4b_68PNNBkRVOrypy4IM5YjC4p-6dgCMM,3793
 ai_edge_torch/generative/test/test_loader.py,sha256=9mQUeeZKOVApOWSWl2cN9c10axZjMKM1-0Zd823CCS4,3449
-ai_edge_torch/generative/test/test_model_conversion.py,sha256=4d_UF19KYf5xFa3yhQGe1nu3TKzmXrbr9PFiEZpPlyk,6274
-ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=4lfZXjXqfjiyxc2s8vMuYbdOZzD8VuPelr2AQo9PFNI,11656
+ai_edge_torch/generative/test/test_model_conversion.py,sha256=jfqkECCX7XKHeBAuDXrkwQJf0vM72eG3LMc5rluha84,6191
+ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=mVuax3MPRmuNjnDRKXqtc9YmswCy7MnhD1CHADK-3nk,11501
 ai_edge_torch/generative/test/test_quantize.py,sha256=bEJMhpQ9bIDUZVBXTW888728FcH-i3SyE4JSZZUgU0A,6071
 ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv9og0DDLHU,2656
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
@@ -159,12 +159,11 @@ ai_edge_torch/hlfb/mark_pattern/passes.py,sha256=pjkKcI1nHECPluAt87cFBrt1DP0f3ge
 ai_edge_torch/hlfb/mark_pattern/pattern.py,sha256=NP2mYhe5D2GjtqQfqqldp-ko3xtNghuFKKJOQskUJFI,10041
 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/hlfb/test/test_stablehlo_composite_builder.py,sha256=j8WpeS-mz3Zr4I7p7NwanQzkQNeH0asZ7lz5y7twgQ4,8447
 ai_edge_torch/lowertools/__init__.py,sha256=A8WBXvWtuFYYWtNTqPD7waVntLaSVAnSMwx5ugjZBIw,761
-ai_edge_torch/lowertools/_shim.py,sha256=ilL7x1ebUBj1clg7bagrX4y_nVSHiGrvDrOVfuTeenE,3039
+ai_edge_torch/lowertools/_shim.py,sha256=xJIHDSWNoF4PkkT0JkjeJxgguQ9JGEwooJf9xZNkVRU,3058
 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/test_utils.py,sha256=bPgc2iXX16KYtMNvmsRdKfrCY6UJmcfitfCOvHoD7Oc,1930
+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/translate_recipe.py,sha256=ymkBpFqAUiupRWqrPOWiVphKcXR1K5vHK0RjgBFtxlE,5652
 ai_edge_torch/odml_torch/__init__.py,sha256=S8jOzE9nLof-6es3XDiGJRN-9H_XTxsVm9dE7lD3RWo,812
@@ -183,12 +182,12 @@ ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=Jco5zvejxuyl9xHQxZICAKbkg
 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=0GytV1dGnqe1mKityqQDNFNS8T4QBg3UZuRJcGHwGyA,993
-ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=ufvnaAh6rM_yfoc8ybI3VErHEVBv5W_p4iOe9slfwKM,9948
+ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=8mZTp_ybcMO3tDRQdlDP68BVeTw560XsTR4XH-ldTdc,9987
 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
+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=1ePJs7oIdUkVdMddFsXMc53qTkEKqGz0ZhQQoNzBa10,2862
-ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=rFmzqcdjYrwhcxH8j9zCFStPy21HFF7hkUV_GQ8FPAk,6056
+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/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
@@ -201,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.dev20241213.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241213.dist-info/METADATA,sha256=Yzw2YkrbFAe1EbxfoKBDME5NQe0GIzzseUOVylqFpnM,1897
-ai_edge_torch_nightly-0.3.0.dev20241213.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20241213.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241213.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241214.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241214.dist-info/METADATA,sha256=fUbq26zB0WUU1l6eUud8vq3Nm3KSIhox74pzFSFTmoM,1966
+ai_edge_torch_nightly-0.3.0.dev20241214.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20241214.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241214.dist-info/RECORD,,

ai_edge_torch/config.py

@@ -1,27 +0,0 @@
-# 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.
-# ==============================================================================
-
-"""Provides a configuration for the AI Edge Torch library."""
-
-import dataclasses
-import os
-
-
-@dataclasses.dataclass
-class Config:
-  use_torch_xla: bool = os.environ.get("USE_TORCH_XLA", "true").lower() in (
-      "1",
-      "true",
-  )

ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py

@@ -1,283 +0,0 @@
-# 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.
-# ==============================================================================
-"""Tests for StableHLOCompositeBuilder."""
-
-import math
-
-from ai_edge_torch import config
-from ai_edge_torch import lowertools
-from ai_edge_torch.hlfb import StableHLOCompositeBuilder
-import torch
-import torch.nn.functional as F
-
-from absl.testing import absltest as googletest
-
-
-def _export_stablehlo_mlir(model, args):
-  ep = torch.export.export(model, args)
-  return lowertools.exported_program_to_mlir_text(ep)
-
-
-@googletest.skipIf(
-    not config.Config.use_torch_xla,
-    reason="The odml_torch counter part is in odml_torch.",
-)
-class TestStableHLOCompositeBuilder(googletest.TestCase):
-
-  def test_build_composite(self):
-    class SampleModel(torch.nn.Module):
-
-      def forward(self, x):
-        builder = StableHLOCompositeBuilder(name="test.plus_two")
-        y = x + 1
-        y = builder.mark_inputs(y)
-        z = y + 2
-        z = builder.mark_outputs(z)
-        return z
-
-    mlir = _export_stablehlo_mlir(SampleModel().eval(), (torch.rand((2, 2)),))
-    self.assertEqual(mlir.count('stablehlo.composite "test.plus_two"'), 1)
-
-  def test_build_multiple_composites(self):
-    class SampleModel(torch.nn.Module):
-
-      def plus_one(self, x: torch.Tensor):
-        builder = StableHLOCompositeBuilder("test.plus_one")
-        x = builder.mark_inputs(x)
-        y = x + 1
-        y = builder.mark_outputs(y)
-        return y
-
-      def plus_two(self, x: torch.Tensor):
-        builder = StableHLOCompositeBuilder("test.plus_two")
-        x = builder.mark_inputs(x)
-        y = x + 2
-        y = builder.mark_outputs(y)
-        return y
-
-      def forward(self, x):
-        x = self.plus_two(x)
-        x = x + 3
-        x = self.plus_one(x)
-        x = x + 4
-        x = self.plus_two(x)
-        return x
-
-    mlir = _export_stablehlo_mlir(SampleModel().eval(), (torch.rand((2, 2)),))
-    self.assertEqual(mlir.count('stablehlo.composite "test.plus_one"'), 1)
-    self.assertEqual(mlir.count('stablehlo.composite "test.plus_two"'), 2)
-
-  def test_build_composite_with_attr(self):
-    class SampleModel(torch.nn.Module):
-
-      def __init__(self):
-        super().__init__()
-
-      def log_softmax(self, x: torch.Tensor, dim: int):
-        builder = StableHLOCompositeBuilder(
-            name="test.log_softmax", attr={"dim": dim}
-        )
-        x = builder.mark_inputs(x)
-        y = torch.nn.functional.log_softmax(x, dim=dim)
-        y = builder.mark_outputs(y)
-        return y
-
-      def forward(self, x):
-        x = x + 1
-        x = self.log_softmax(x, 0)
-        x = self.log_softmax(x, 1)
-        return x
-
-    mlir = _export_stablehlo_mlir(SampleModel().eval(), (torch.rand((2, 2)),))
-    self.assertEqual(mlir.count('stablehlo.composite "test.log_softmax"'), 2)
-    self.assertEqual(mlir.count("composite_attributes = {dim = 0 : i64}"), 1)
-    self.assertEqual(mlir.count("composite_attributes = {dim = 1 : i64}"), 1)
-
-  def test_build_composite_with_mix_type_attrs(self):
-    class SampleModel(torch.nn.Module):
-
-      def __init__(self):
-        super().__init__()
-
-      def log_softmax(self, x: torch.Tensor, dim: int):
-        builder = StableHLOCompositeBuilder(
-            name="test.log_softmax",
-            attr={
-                "dim": dim,
-                "source": "torch.nn",
-                "version": 1.0,
-            },
-        )
-        x = builder.mark_inputs(x)
-        y = torch.nn.functional.log_softmax(x, dim=dim)
-        y = builder.mark_outputs(y)
-        return y
-
-      def forward(self, x):
-        x = x + 1
-        x = self.log_softmax(x, 0)
-        return x
-
-    mlir = _export_stablehlo_mlir(SampleModel().eval(), (torch.rand((2, 2)),))
-    self.assertEqual(mlir.count('stablehlo.composite "test.log_softmax"'), 1)
-    self.assertEqual(
-        mlir.count(
-            'composite_attributes = {dim = 0 : i64, source = "torch.nn",'
-            " version = 1.000000e+00 : f32}"
-        ),
-        1,
-    )
-
-  def test_sdpa_composite(self):
-    class SDPAModel(torch.nn.Module):
-
-      def scaled_dot_product_attention(
-          self,
-          q: torch.Tensor,
-          k: torch.Tensor,
-          v: torch.Tensor,
-          head_size: int,
-          mask: torch.Tensor,
-      ):
-        builder = StableHLOCompositeBuilder("test.scaled_dot_product_attention")
-        q, k, v, mask = builder.mark_inputs(q, k, v, mask)
-
-        scale = 1.0 / math.sqrt(head_size)
-
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
-        y = F.scaled_dot_product_attention(
-            q,
-            k,
-            v,
-            attn_mask=mask,
-            dropout_p=0.0,
-            is_causal=mask is None,
-            scale=scale,
-        )
-        result = y.transpose(1, 2)
-        result = builder.mark_outputs(result)
-        return result
-
-      def forward(self, q, k, v, mask):
-        x = self.scaled_dot_product_attention(
-            q,
-            k,
-            v,
-            8,
-            mask,
-        )
-        return x
-
-    query = torch.rand(1, 1, 32, 4)
-    key = torch.rand(1, 500, 1, 4)
-    value = torch.rand(1, 500, 1, 4)
-    mask = torch.rand(1, 1, 1, 500)
-
-    mlir = _export_stablehlo_mlir(
-        SDPAModel().eval(),
-        (query, key, value, mask),
-    )
-    self.assertEqual(
-        mlir.count('stablehlo.composite "test.scaled_dot_product_attention"'), 1
-    )
-
-  def test_sdpa_composite_with_attr(self):
-    class SDPAModel(torch.nn.Module):
-
-      def scaled_dot_product_attention(
-          self,
-          q: torch.Tensor,
-          k: torch.Tensor,
-          v: torch.Tensor,
-          head_size: int,
-          include_captanh: bool,
-      ):
-        builder = StableHLOCompositeBuilder(
-            name="test.scaled_dot_product_attention",
-            attr={"include_captanh": include_captanh},
-        )
-        q, k, v = builder.mark_inputs(q, k, v)
-
-        scale = 1.0 / math.sqrt(head_size)
-
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
-        y = F.scaled_dot_product_attention(
-            q,
-            k,
-            v,
-            attn_mask=None,
-            dropout_p=0.0,
-            is_causal=True,
-            scale=scale,
-        )
-        result = y.transpose(1, 2)
-        result = builder.mark_outputs(result)
-        return result
-
-      def forward(self, q, k, v):
-        x = self.scaled_dot_product_attention(q, k, v, 8, True)
-        y = self.scaled_dot_product_attention(q, k, v, 8, False)
-        return x + y
-
-    query = torch.rand(1, 1, 32, 4)
-    key = torch.rand(1, 500, 1, 4)
-    value = torch.rand(1, 500, 1, 4)
-    mlir = _export_stablehlo_mlir(
-        SDPAModel().eval(),
-        (query, key, value),
-    )
-    self.assertEqual(
-        mlir.count('stablehlo.composite "test.scaled_dot_product_attention"'), 2
-    )
-    self.assertEqual(
-        mlir.count("composite_attributes = {include_captanh = true}"), 1
-    )
-    self.assertEqual(
-        mlir.count("composite_attributes = {include_captanh = false}"), 1
-    )
-
-  def test_build_composite_with_multiple_inputs_outputs(self):
-    class SampleModel(torch.nn.Module):
-
-      def mimo_sample(self, a, b, c):
-        builder = StableHLOCompositeBuilder(name="test.mimo_sample")
-
-        a, b, c = builder.mark_inputs(a, b, c)
-        x = a + b + c
-        y = (a - b) * x
-        z = (c + 1.0) * a
-        x, y, z = builder.mark_outputs(x, y, z)
-
-        result = x + y * z
-        return result
-
-      def forward(self, a, b, c):
-        x = self.mimo_sample(a, b, c)
-        x = self.mimo_sample(a, b, x)
-        x = self.mimo_sample(x, x, c)
-        return x
-
-    mlir = _export_stablehlo_mlir(
-        SampleModel().eval(), (torch.rand(2), torch.rand(2), torch.rand(2))
-    )
-    self.assertEqual(mlir.count('stablehlo.composite "test.mimo_sample"'), 3)
-
-
-if __name__ == "__main__":
-  googletest.main()