ai-edge-torch-nightly 0.3.0.dev20241122__py3-none-any.whl → 0.3.0.dev20241123__py3-none-any.whl

ai_edge_torch/generative/layers/builder.py

@@ -74,6 +74,7 @@ def build_norm(dim: int, config: cfg.NormalizationConfig):
         dim,
         eps=config.epsilon,
         zero_centered_gamma=config.zero_centered,
+        enable_hlfb=config.enable_hlfb,
     )
   elif config.type == cfg.NormalizationType.LAYER_NORM:
     return normalization.LayerNorm(dim, config.epsilon, config.enable_hlfb)

ai_edge_torch/generative/layers/normalization.py

@@ -23,15 +23,24 @@ import torch.nn.functional as F
 # Implementation for RMSNorm from: https://arxiv.org/abs/1910.07467
 class RMSNorm(torch.nn.Module):
 
-  def __init__(self, dim: int, eps: float = 1e-6, zero_centered_gamma=False):
+  def __init__(
+      self,
+      dim: int,
+      eps: float = 1e-6,
+      zero_centered_gamma=False,
+      enable_hlfb: bool = False,
+  ):
     """Initialize the RMSNorm layer.
 
     Args:
       dim (int): dimension of the input tensor.
       eps (float): A small float value to ensure numerical stability (default:
         1e-6).
+      zero_centered_gamma (bool): Whether or not gamma has an offset.
+      enable_hlfb (bool): use HLFB in the op.
     """
     super().__init__()
+    self.enable_hlfb = enable_hlfb
     self.eps = eps
     self.weight = torch.nn.Parameter(torch.ones(dim))
     self.zero_centered_gamma = zero_centered_gamma
@@ -56,12 +65,20 @@ class RMSNorm(torch.nn.Module):
     Returns:
       torch.Tensor: output tensor after applying RMSNorm.
     """
-    output = self._norm(x.float()).type_as(x)
     if self.zero_centered_gamma:
-      return output * (1 + self.weight)
+      w = 1 + self.weight
     else:
-      return output * self.weight
+      w = self.weight
 
+    if self.enable_hlfb:
+      return rms_norm_with_hlfb(
+          x,
+          w,
+          self.eps,
+      )
+    else:
+      output = self._norm(x.float()).type_as(x)
+      return output * w
 
 class GroupNorm(torch.nn.Module):
 
@@ -194,6 +211,37 @@ def group_norm_with_hlfb(
   return y
 
 
+def rms_norm_with_hlfb(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    eps: float,
+):
+  """RMS Normalization with high-level function boundary enabled.
+
+  Args:
+    x (torch.Tensor): Input tensor for RMS Normalization, with BCHW shape.
+    w (torch.Tensor): The learned parameter tensor for normalization.
+    eps (float): A small float value to ensure numerical stability.
+
+  Returns:
+    The output tensor of RMS Normalization.
+  """
+  builder = StableHLOCompositeBuilder(
+      name="odml.rms_norm", attr={"epsilon": eps}
+  )
+
+  x, w = builder.mark_inputs(x, w)
+
+  def _norm(x):
+    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
+
+  output = _norm(x.float()).type_as(x)
+  out = output * w
+
+  out = builder.mark_outputs(out)
+  return out
+
+
 def layer_norm_with_hlfb(
     x: torch.Tensor,
     normalized_shape: list[int],

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241122.dist-info → ai_edge_torch_nightly-0.3.0.dev20241123.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20241122
+Version: 0.3.0.dev20241123
 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.dev20241122.dist-info → ai_edge_torch_nightly-0.3.0.dev20241123.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=B4r6opjqsPmDJdLbwvWto6dM-0KbsjszxSL6CXmi8K8,706
+ai_edge_torch/version.py,sha256=H-20IqCwUbnpEuOveG8xD4bixK8svU7k-n0hfdJ8AoY,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
@@ -115,11 +115,11 @@ ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=hhxSQvkD
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/attention.py,sha256=zN3BQjA25Ej_aRU0rFnyx--K74xf5ykc02zGvUpYHeE,13295
 ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
-ai_edge_torch/generative/layers/builder.py,sha256=Z5LyzCEThgnYZeyViakaE3yJVzTGHtw13acHsAQR15U,5050
+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=lbm-yJ1jGPtcgWS4C3FmSnB1IlxqDE7g0BLRh3PN4N4,6324
 ai_edge_torch/generative/layers/model_config.py,sha256=viX51T_naJ9sPpPxPoMnSueBPYE2zxWNOD0xn0f-_bM,7510
-ai_edge_torch/generative/layers/normalization.py,sha256=eKAGst9rPuyRFExMcQFJO7R3iHdCtlmjeF_lITjLhwE,6498
+ai_edge_torch/generative/layers/normalization.py,sha256=_2hps2m2MXEHQWbM-1B4he90hbq8wqOnIDIf-qXHhpc,7589
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=xxWtlVsGGJkEyXC6PwznubyhJnLPEfSpHOORE_hgxss,2670
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=gXxh3papKy4FBpGEX7VyZ7rZ1Js6aHK70Q6DKrVSckY,4154
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -196,8 +196,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.dev20241122.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241122.dist-info/METADATA,sha256=-YpC-ksRKR8hJ8pZET4Q2F5KbUiRmGOXPhBoEQgIuOA,1897
-ai_edge_torch_nightly-0.3.0.dev20241122.dist-info/WHEEL,sha256=bFJAMchF8aTQGUgMZzHJyDDMPTO3ToJ7x23SLJa1SVo,92
-ai_edge_torch_nightly-0.3.0.dev20241122.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241122.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241123.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241123.dist-info/METADATA,sha256=wsp1IYSM424I9ovzuvu-4b_T0VAcCcABGZuChuCdzWM,1897
+ai_edge_torch_nightly-0.3.0.dev20241123.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20241123.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241123.dist-info/RECORD,,

{ai_edge_torch_nightly-0.3.0.dev20241122.dist-info → ai_edge_torch_nightly-0.3.0.dev20241123.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.45.0)
+Generator: bdist_wheel (0.45.1)
 Root-Is-Purelib: true
 Tag: py3-none-any