ai-edge-torch-nightly 0.3.0.dev20240923__py3-none-any.whl → 0.3.0.dev20240924__py3-none-any.whl

ai_edge_torch/generative/examples/openelm/openelm.py

@@ -161,9 +161,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
         ),
         ff_config=cfg.FeedForwardConfig(
             type=cfg.FeedForwardType.SEQUENTIAL,
-            activation=cfg.ActivationConfig(
-                cfg.ActivationType.SILU_GLU, gate_is_front=True
-            ),
+            activation=cfg.ActivationConfig(cfg.ActivationType.SILU_GLU),
             intermediate_size=get_intermediate_size(idx),
             pre_ff_norm_config=norm_config,
         ),

ai_edge_torch/generative/examples/phi/verify.py

@@ -34,6 +34,7 @@ _MAX_NEW_TOKENS = flags.DEFINE_integer(
     "The maximum size of the generated tokens.",
 )
 
+
 def main(_):
   checkpoint = kagglehub.model_download("Microsoft/phi/transformers/2")
   verifier.log_msg("Loading the original model from", checkpoint)

ai_edge_torch/generative/layers/builder.py

@@ -23,34 +23,35 @@ from torch import nn
 import torch.nn.functional as F
 
 
-def build_glu(
-    act: Callable[[torch.Tensor], torch.Tensor], gate_is_front: bool = False
-) -> Callable[[torch.Tensor], torch.Tensor]:
-  """Builds an activation function with GLU (Gated Linear Unit).
+class GeGLU(nn.Module):
+  """GeGLU is an activation function which is a variant of GELU.
 
-  If gate_is_front is True,
-    f(x) = act(x) * y
-  otherwise,
-    f(x) = x * act(y),
-  where x is the first half of the input and y is the second half of the input.
+  GeGLU(x) = (xW+b) * GELU(xV+c)
+  See: https://arxiv.org/abs/2002.05202v1
+  """
 
-  Args:
-    act (Callable[[torch.Tensor], torch.Tensor]): activation function to apply
-      to the gate.
-    gate_is_front: whether the gate is in front half of the input. Other part is
-      the output in GLU.
+  def __init__(self, d_in: int, d_out: int):
+    super().__init__()
+    self.proj = nn.Linear(d_in, d_out * 2)
 
-  Returns:
-    A callable activation function with GLU.
+  def forward(self, x: torch.Tensor):
+    x, gate = self.proj(x).chunk(2, dim=-1)
+    return x * F.gelu(gate)
+
+
+class SwiGLU(nn.Module):
+  """SwiGLU is an activation function which is a variant of GLU.
+
+  SwiGLU is same as SiLU_GLU, because The SiLU function is also known as the
+  swish function.
+
+  SwiGLU(x) = Swish(xW+b) * (xV+c)
+  See: https://paperswithcode.com/method/swiglu
   """
 
-  def _glu(x):
+  def forward(self, x: torch.Tensor):
     x, y = x.chunk(2, dim=-1)
-    if gate_is_front:
-      return act(x) * y
-    return x * act(y)
-
-  return _glu
+    return F.silu(x) * y
 
 
 def build_norm(dim: int, config: cfg.NormalizationConfig):
@@ -151,10 +152,10 @@ def get_activation(config: cfg.ActivationConfig):
     # See: https://github.com/hendrycks/GELUs
     return lambda x: x * F.sigmoid(1.702 * x)
   elif config.type == cfg.ActivationType.GE_GLU:
-    return build_glu(F.gelu, config.gate_is_front)
+    return GeGLU(config.dim_in, config.dim_out)
   elif config.type == cfg.ActivationType.RELU:
     return F.relu
   elif config.type == cfg.ActivationType.SILU_GLU:
-    return build_glu(F.silu, config.gate_is_front)
+    return SwiGLU()
   else:
     raise ValueError("Unsupported activation type.")

ai_edge_torch/generative/layers/model_config.py

@@ -118,9 +118,9 @@ class AttentionConfig:
 @dataclass
 class ActivationConfig:
   type: ActivationType = ActivationType.LINEAR
-  # Whether to GLU gate is the front part instead of the back part of input
-  # when ActivationType is `GE_GLU` or `SILU_GLU`.
-  gate_is_front: bool = False
+  # Dimension of input and output, used in GeGLU.
+  dim_in: Optional[int] = None
+  dim_out: Optional[int] = None
 
 
 @dataclass

ai_edge_torch/generative/layers/normalization.py

@@ -183,8 +183,16 @@ def group_norm_with_hlfb(
   """
   x = torch.permute(x, (0, 2, 3, 1))
 
+  # TODO: b/366544750 - Change "reduction_axes" field as an array, rather than
+  # int32 when the bug is fixed.
   builder = StableHLOCompositeBuilder(
-      name="odml.group_norm", attr={"num_groups": num_groups, "eps": eps}
+      name="odml.group_norm",
+      attr={
+          "num_groups": num_groups,
+          "eps": eps,
+          "reduction_axes": 3,
+          "channel_axis": 3,
+      },
   )
   x, w, b = builder.mark_inputs(x, w, b)
   x = torch.permute(x, (0, 3, 1, 2))
@@ -206,7 +214,7 @@ def layer_norm_with_hlfb(
   """Layer Normalization with high-level function boundary enabled.
 
   Args:
-    x (torch.Tensor): Input tensor for Layer Normalization.
+    x (torch.Tensor): Input tensor for Layer Normalization, with BCHW shape.
     w (torch.Tensor): The weight tensor for the normalization.
     b (torch.Tensor): The bias tensor for the normalization.
     eps (float): A small float value to ensure numerical stability.
@@ -216,7 +224,10 @@ def layer_norm_with_hlfb(
   Returns:
     The output tensor of Layer Normalization.
   """
-  builder = StableHLOCompositeBuilder(name="odml.layer_norm", attr={"eps": eps})
+  builder = StableHLOCompositeBuilder(
+      name="odml.group_norm",
+      attr={"num_groups": 1, "eps": eps, "channel_axis": 1},
+  )
   x, w, b = builder.mark_inputs(x, w, b)
   if use_input_shape:
     normalized_shape = x.shape

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20240923.dist-info → ai_edge_torch_nightly-0.3.0.dev20240924.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20240923
+Version: 0.3.0.dev20240924
 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
@@ -30,6 +30,7 @@ Requires-Dist: tabulate
 Requires-Dist: torch>=2.4.0
 Requires-Dist: torch-xla>=2.4.0
 Requires-Dist: tf-nightly>=2.18.0.dev20240722
+Requires-Dist: ai-edge-litert-nightly
 Requires-Dist: ai-edge-quantizer-nightly
 
 Library that supports converting PyTorch models into a .tflite format, which can

{ai_edge_torch_nightly-0.3.0.dev20240923.dist-info → ai_edge_torch_nightly-0.3.0.dev20240924.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=D86Gw3pIRcpnTebUPKlnPbPGJae1S6Fw4DZZ3ZkD0zw,3730
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=oxtOOEY9LJkV5vRrgr1EoSjAjuetYVNq7WQqMuauRkc,706
+ai_edge_torch/version.py,sha256=sQUcRP5rShDk3vfblz87j26JciN6PV8S8DJkiiZP5o8,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=5uPwHhmc6kwiIz-CqaiHDejf2SOWMHrb-rYEHm69wKc,3801
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -48,12 +48,12 @@ ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=kSzn1ITJXqrtNQax
 ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=HBK2d8FcWFoxVDF5zk9sLSbKZEtwZQhX-K_zm4AvQtQ,5160
 ai_edge_torch/generative/examples/openelm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=85FVEt6cKFP2UzCLC78tAkbwGlGhAArtG7Wa75NxJik,2185
-ai_edge_torch/generative/examples/openelm/openelm.py,sha256=gGkHELNrt4xqnu11fCh3sJbZ7OsPyvoiF1J1aKCs5r8,7532
+ai_edge_torch/generative/examples/openelm/openelm.py,sha256=VcU8A0B9nQR-FTPHXqNHSHZzeIZZ_As4yvKZMnoU2P4,7482
 ai_edge_torch/generative/examples/openelm/verify.py,sha256=QdFKymQSCYFJcYVvA63u5uIsn1YxJ0JZD5UqN6gxraI,2112
 ai_edge_torch/generative/examples/phi/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=3go690yX6PFeXMdpY7y4JZorAwxX0HT_b_pKZieauvk,2169
 ai_edge_torch/generative/examples/phi/phi2.py,sha256=YwAszA53aOjvaMJ5wua2-5rP79N21Un_Y5yBCfFSYNU,6189
-ai_edge_torch/generative/examples/phi/verify.py,sha256=QPYX6weEZGMEXt_Vb2hNARPAECQBKzx-KCivd4dzOrw,2145
+ai_edge_torch/generative/examples/phi/verify.py,sha256=SwPyRjiupD4AsmWW_7FDcMSWaNRmDBu6uVFcBQRoM40,2146
 ai_edge_torch/generative/examples/smollm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/smollm/convert_to_tflite.py,sha256=zPrDTDeRVWFi9DS32uNi-RLpzOStFOk5MhNla4ixeew,2179
 ai_edge_torch/generative/examples/smollm/smollm.py,sha256=hyhMk-b5762Q2xmjdD47g85dcbBSNJXNPIsifm1DRto,3239
@@ -89,11 +89,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=Z0Y_G8IG0LmvLX2u9D8__Fkr22szB-az6wMNnZpzhkA,13233
 ai_edge_torch/generative/layers/attention_utils.py,sha256=68GXGR2HSWBFViTxX7cHifzVG-kcLS2IL2tQJPIpupg,7344
-ai_edge_torch/generative/layers/builder.py,sha256=toT9Tl1x9o5KbG-eGOEViUr4fd_4f-XLZdMQT0Ae5_8,5130
+ai_edge_torch/generative/layers/builder.py,sha256=oE8DdqLA-oWkBC2zySSCh8JNAJg_hk8-W_UoMSrgDVk,5088
 ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
 ai_edge_torch/generative/layers/kv_cache.py,sha256=2El7kZYnQRCRcVc63xgiAdBh9oVOksDu35p9XggvaGE,6148
-ai_edge_torch/generative/layers/model_config.py,sha256=d0Y-EFb4Rr7iLZ4Bsdf1i92KuhY1BXRqyeUN2kuu510,6923
-ai_edge_torch/generative/layers/normalization.py,sha256=l_36uFdruJwqqyubnBTM0M-iGiJfeFafyXKPPK8KHVo,6713
+ai_edge_torch/generative/layers/model_config.py,sha256=l5Rb3h3GK2pux-Lg3BONTD6b7klxXqUbDDtYs_bGKLk,6879
+ai_edge_torch/generative/layers/normalization.py,sha256=LDczSHujMgo1WV8IhYVQe-egPkaBEmWFt8wZQ_tgshg,6991
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=CZqOoibLcHvUgrgaIIWAlmk3XgE2inzx340MN-npLoU,1347
 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
@@ -166,8 +166,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.dev20240923.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20240923.dist-info/METADATA,sha256=BgwLxDJ3AOPVn0fkngAQpf3YdmShufhMt3bANFevtiQ,1859
-ai_edge_torch_nightly-0.3.0.dev20240923.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20240923.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20240923.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20240924.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20240924.dist-info/METADATA,sha256=BotYlw1pMxClnHOi8rSb5v6jX0zE7EqUo8b11xvqEII,1897
+ai_edge_torch_nightly-0.3.0.dev20240924.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20240924.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20240924.dist-info/RECORD,,