ai-edge-torch-nightly 0.6.0.dev20250821__py3-none-any.whl → 0.6.0.dev20250823__py3-none-any.whl

ai_edge_torch/generative/layers/builder.py

@@ -54,7 +54,11 @@ class SwiGLU(nn.Module):
     return F.silu(x) * y
 
 
-def build_norm(dim: int, config: cfg.NormalizationConfig):
+def build_norm(
+    dim: int,
+    config: cfg.NormalizationConfig,
+    init_fn: Callable[..., torch.Tensor] = lambda *args, **kwargs: None,
+):
   """Builder function for normalizers.
 
   Args:
@@ -77,6 +81,7 @@ def build_norm(dim: int, config: cfg.NormalizationConfig):
         with_scale=config.with_scale,
         scale_shift=config.scale_shift,
         enable_hlfb=config.enable_hlfb,
+        init_fn=init_fn,
     )
   elif config.type == cfg.NormalizationType.LAYER_NORM:
     return normalization.LayerNorm(dim, config.epsilon, config.enable_hlfb)

ai_edge_torch/generative/layers/normalization.py

@@ -14,6 +14,8 @@
 # ==============================================================================
 # Common normalization layers.
 
+from typing import Callable
+
 from ai_edge_torch.hlfb import StableHLOCompositeBuilder
 import torch
 from torch import nn
@@ -31,6 +33,7 @@ class RMSNorm(torch.nn.Module):
       with_scale: bool = False,
       scale_shift: float = 1.0,
       enable_hlfb: bool = False,
+      init_fn: Callable[..., torch.Tensor] = lambda *args, **kwargs: None,
   ):
     """Initialize the RMSNorm layer.
 
@@ -42,12 +45,15 @@ class RMSNorm(torch.nn.Module):
       with_scale (bool): Whether or not to use a scale parameter.
       scale_shift (float): The shift to apply to the scale parameter.
       enable_hlfb (bool): use HLFB in the op.
+      init_fn: The initialization function to use for the parameters. This is
+        used to initialize the scale parameter.
     """
     super().__init__()
     self.dim = dim
     self.enable_hlfb = enable_hlfb
     self.eps = eps
     self.weight = torch.nn.Parameter(torch.ones(dim), requires_grad=False)
+    init_fn(self.weight)
     self.zero_centered_gamma = zero_centered_gamma
     self.with_scale = with_scale
     if with_scale:

ai_edge_torch/odml_torch/export.py

@@ -124,9 +124,11 @@ class LoweringInterpreter(torch.fx.Interpreter):
   def run_node(self, node: torch.fx.Node):
     loc = self._build_loc(node)
     with loc:
-      self.lctx = self.lctx.replace(ir_location=loc, node=node)
+      self.lctx.ir_location = loc
+      self.lctx.node = node
       res = super().run_node(node)
-      self.lctx = self.lctx.replace(ir_location=None, node=None)
+      self.lctx.ir_location = None
+      self.lctx.node = None
     return res
 
   def call_function(self, target, args, kwargs):

ai_edge_torch/odml_torch/export_utils.py

@@ -63,7 +63,10 @@ def inline(
   while True:
     is_changed = False
     for op in block.operations:
-      if op.OPERATION_NAME != func.CallOp.OPERATION_NAME:
+      if (
+          not hasattr(op, "OPERATION_NAME")
+          or op.OPERATION_NAME != func.CallOp.OPERATION_NAME
+      ):
         continue
 
       call_op = cast(func.CallOp, op)
@@ -96,7 +99,10 @@ def clone_func_body_ops(func_op: func.FuncOp, ir_inputs: Sequence[ir.Value]):
 
   for op in list(func_op.entry_block.operations):
     cloned_operands = [value_mapping[val] for val in op.operands]
-    if op.OPERATION_NAME == func.ReturnOp.OPERATION_NAME:
+    if (
+        hasattr(op, "OPERATION_NAME")
+        and op.OPERATION_NAME == func.ReturnOp.OPERATION_NAME
+    ):
       return cloned_operands
 
     cloned = cast(ir.Operation, op.operation.clone())

ai_edge_torch/odml_torch/lowerings/context.py

@@ -37,6 +37,3 @@ class LoweringContext:
   def loc(self):
     """Shortcut for ir_location."""
     return self.ir_location
-
-  def replace(self, **kwargs):
-    return dataclasses.replace(self, **kwargs)

ai_edge_torch/version.py

@@ -15,4 +15,4 @@
 
 # The next version of ai-edge-torch.
 # The minor version code should be bumped after every release.
-__version__ = "0.6.0.dev20250821"
+__version__ = "0.6.0.dev20250823"

{ai_edge_torch_nightly-0.6.0.dev20250821.dist-info → ai_edge_torch_nightly-0.6.0.dev20250823.dist-info}/METADATA

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

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=lemyLCNoGYRnJsmDuGZu7qOqLbLqG6CGDFtu3ue1syU,129
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=A7loFu8jE9CsXsfMmHYZ-KDFJiaD8Kkqwm_9d3IVzk0,5638
-ai_edge_torch/version.py,sha256=KCQW3wwoSg8zjssrnlyquD4aO-OeN7u1jWFHk_-trYY,806
+ai_edge_torch/version.py,sha256=JyfBN0yvWEvYs121XsXOTiW8p-4px3FWXcaDmfbzVlY,806
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=iQk3R-pLq4c1nfLqPB4xTRj78gghxPGzJCJtILLdg5o,6123
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -173,7 +173,7 @@ ai_edge_torch/generative/layers/attention.py,sha256=RaXENRRQo1MsLdt3U8h3kYTCmd6i
 ai_edge_torch/generative/layers/attention_test.py,sha256=9v8v96TLyFPdqxEylU1JOAeRFAp2s0YoDHZN83SFJJA,4764
 ai_edge_torch/generative/layers/attention_utils.py,sha256=2qfg7Tzk9ikKph5w3geOHC1I6EyOCdDsWXMr7F7IOZM,7630
 ai_edge_torch/generative/layers/attention_utils_test.py,sha256=22gQ1gcRPkwqFG3_p82GZfRKVE3udEssSy58wNOqv0w,2431
-ai_edge_torch/generative/layers/builder.py,sha256=2bUgkyowDkDznkF8XaHyZs4nowHr1QEHYLM7pMaFmIk,4921
+ai_edge_torch/generative/layers/builder.py,sha256=UiLOyvwd-bc0n5XcbVxi6JCn_qiKSC6zrDKSZT_TSDA,5030
 ai_edge_torch/generative/layers/einsum.py,sha256=LH4CNHr-pFfLUuCpwbYL3GpoAMgHJ4nLju3XCqA4VwM,1416
 ai_edge_torch/generative/layers/einsum_test.py,sha256=ltIE773bvvNLv_9aLQxFwe1MgQ762sez0c5E2tejxuA,1079
 ai_edge_torch/generative/layers/feed_forward.py,sha256=_GmtHxwL068l9gh_F_WFcFk7La-Tl5SfoQ9v2hMabZM,5541
@@ -181,7 +181,7 @@ ai_edge_torch/generative/layers/feed_forward_test.py,sha256=Y5l1eC9NgfYixHcfIfE1
 ai_edge_torch/generative/layers/kv_cache.py,sha256=A0IFXZ1HD2ZHOWRLfsDO4almgE0KQfjyBOdBFZIGnAs,10893
 ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
 ai_edge_torch/generative/layers/model_config.py,sha256=HP-vu1UmAiTmdLlTyZGDUF3le0gji8a61mLCy966NZw,10261
-ai_edge_torch/generative/layers/normalization.py,sha256=ijwCpi22NLX-Sygwy5sK9l9WjGvbPIhZvVwoBAonWAo,7014
+ai_edge_torch/generative/layers/normalization.py,sha256=WAhPcLbcC3SCEa6oIgdsojvN306_S8d90WyMQ7ZVP6I,7269
 ai_edge_torch/generative/layers/normalization_test.py,sha256=zwurZly-TgFxdgVVdpzu9vCpcLbd5RYt_gKg9Lfg1jI,2248
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIrILJ7blceAcxrNqX1ZCN0ECKG1gz-bV8,2655
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=2_AgwENsaOgaxgiSqgoj0V0JzQ09dFtP_nBhX-lJK2g,5648
@@ -235,8 +235,8 @@ ai_edge_torch/lowertools/translate_recipe.py,sha256=Oavs0ENKVnIeB-WidXvokTPqNlFf
 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_library.py,sha256=Lw1gqL2HWNRspdTwNhIkYAHDyafHedHtkXyKKxn-Wss,805
-ai_edge_torch/odml_torch/export.py,sha256=FDseiAOOcgN8HOuwv6HT9sALZgXGj7vZi6kwZPWPF14,14935
-ai_edge_torch/odml_torch/export_utils.py,sha256=QeA37Irlty6AiIBuqmHmJgn3lqahBQ5xsh6IKRoKm1g,4774
+ai_edge_torch/odml_torch/export.py,sha256=WPsATMFsc62aETExBsb1v6Ec1fL3r0CEvxwE7r2C_CA,14931
+ai_edge_torch/odml_torch/export_utils.py,sha256=Eax4QUefIzpmVuQxo1y9FqJ6g0qXjg4C0IVZ5uYPscs,4899
 ai_edge_torch/odml_torch/optimization_barrier.py,sha256=2lmSiu5iXWLFWpupZHvsVeNYNzG5AVGSK3K_CNhS5Sk,2290
 ai_edge_torch/odml_torch/tf_integration.py,sha256=NN29WeXmHZ0S1RPDFHUnBi2DEjMvAtwczStPYIsQ1w8,4849
 ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xKA9rknO93JS9kw9w_c,778
@@ -257,7 +257,7 @@ ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=xUuQjoR0NJhuwG36Guyc
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=khJIvDVk2s332Nd2Be-5dM6-wp5DGff61HCV5lskHmQ,3011
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQDJnaAAxhNFAd7NfQm5SOVEp_nno_krA,6178
 ai_edge_torch/odml_torch/lowerings/_rand.py,sha256=g6SuqDkuC6hD35lyP1-5H7ASDIzPSmKukeNT5naZSv8,4133
-ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
+ai_edge_torch/odml_torch/lowerings/context.py,sha256=dhGS9oFjKLKvc-aBFbNUhb4LqggP_AmPrgpT1EnOw0w,1216
 ai_edge_torch/odml_torch/lowerings/registry.py,sha256=Tp2h11l5uTifO0aIkuUOWAF_ibEjmd65Xx99w3EXuGE,1924
 ai_edge_torch/odml_torch/lowerings/utils.py,sha256=YagY2IDqtJTpPxjPz2dhb3eyCFTpTSu3ptEPSvEuDtk,10574
 ai_edge_torch/odml_torch/passes/__init__.py,sha256=AVwIwUTMx7rXacKjGy4kwrtMd3XB2v_ncdc40KOjUqQ,1245
@@ -269,8 +269,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 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.6.0.dev20250821.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.6.0.dev20250821.dist-info/METADATA,sha256=HNSeR4DnS8ax0bUi9CSIOzakw0lJk1ZA8JM889_0rSg,2074
-ai_edge_torch_nightly-0.6.0.dev20250821.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.6.0.dev20250821.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.6.0.dev20250821.dist-info/RECORD,,
+ai_edge_torch_nightly-0.6.0.dev20250823.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.6.0.dev20250823.dist-info/METADATA,sha256=ZrMG85KwcLQee8_Nj9jD25avEdJBJQ8zK4o3m7xFifk,2074
+ai_edge_torch_nightly-0.6.0.dev20250823.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.6.0.dev20250823.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.6.0.dev20250823.dist-info/RECORD,,