ai-edge-torch-nightly 0.3.0.dev20241006__py3-none-any.whl → 0.3.0.dev20241007__py3-none-any.whl

ai_edge_torch/generative/examples/phi/phi2.py

@@ -60,7 +60,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   )
   norm_config = cfg.NormalizationConfig(
       type=cfg.NormalizationType.LAYER_NORM,
-      use_input_shape=False,  # Phi-2 does layer-norm with the weight shape.
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,

ai_edge_torch/generative/layers/builder.py

@@ -77,7 +77,7 @@ def build_norm(dim: int, config: cfg.NormalizationConfig):
     )
   elif config.type == cfg.NormalizationType.LAYER_NORM:
     return normalization.LayerNorm(
-        dim, config.epsilon, config.enable_hlfb, config.use_input_shape
+        dim, config.epsilon, config.enable_hlfb
     )
   elif config.type == cfg.NormalizationType.GROUP_NORM:
     return normalization.GroupNorm(

ai_edge_torch/generative/layers/model_config.py

@@ -69,9 +69,6 @@ class NormalizationConfig:
   zero_centered: bool = False
   # Number of groups used in group normalization.
   group_num: Optional[float] = None
-  # Whether to use the input shape to determine the dimension of normalization
-  # when type is LAYER_NORM.
-  use_input_shape: bool = True
 
 
 @dataclass

ai_edge_torch/generative/layers/normalization.py

@@ -86,8 +86,8 @@ class GroupNorm(torch.nn.Module):
     self.enable_hlfb = enable_hlfb
     self.group_num = group_num
     self.eps = eps
-    self.weight = torch.nn.Parameter(torch.ones(dim))
-    self.bias = torch.nn.Parameter(torch.ones(dim))
+    self.weight = torch.nn.Parameter(torch.empty(dim))
+    self.bias = torch.nn.Parameter(torch.empty(dim))
 
   def forward(self, x):
     """Running the forward pass of GroupNorm layer.
@@ -117,25 +117,22 @@ class LayerNorm(torch.nn.Module):
       dim: int,
       eps: float = 1e-5,
       enable_hlfb: bool = False,
-      use_input_shape: bool = True,
   ):
     """Initialize the LayerNorm layer.
 
     Args:
       dim (int): dimension of the input tensor.
       eps (float): A small float value to ensure numerical stability (default:
-        1e-6).
+        1e-5).
       enable_hlfb (bool): Whether to convert this normalization into a single
         op.
-      use_input_shape (bool): Whether to use the input shape to determine the
-        dimension of normalization (default: True).
     """
     super().__init__()
     self.enable_hlfb = enable_hlfb
-    self.use_input_shape = use_input_shape
+    self.normalized_shape = (dim,)
     self.eps = eps
-    self.weight = torch.nn.Parameter(torch.ones(dim))
-    self.bias = torch.nn.Parameter(torch.ones(dim))
+    self.weight = torch.nn.Parameter(torch.empty(dim))
+    self.bias = torch.nn.Parameter(torch.empty(dim))
 
   def forward(self, x):
     """Running the forward pass of LayerNorm layer.
@@ -148,18 +145,11 @@ class LayerNorm(torch.nn.Module):
     """
     if self.enable_hlfb:
       return layer_norm_with_hlfb(
-          x, self.weight, self.bias, self.eps, self.use_input_shape
+          x, self.normalized_shape, self.weight, self.bias, self.eps
       )
-
-    if self.use_input_shape:
-      normalized_shape = x.shape
-      weight = self.weight.broadcast_to(x.shape)
-      bias = self.bias.broadcast_to(x.shape)
-    else:
-      normalized_shape = self.weight.shape
-      weight = self.weight
-      bias = self.bias
-    return F.layer_norm(x, normalized_shape, weight, bias, self.eps)
+    return F.layer_norm(
+        x, self.normalized_shape, self.weight, self.bias, self.eps
+    )
 
 
 def group_norm_with_hlfb(
@@ -206,20 +196,20 @@ def group_norm_with_hlfb(
 
 def layer_norm_with_hlfb(
     x: torch.Tensor,
+    normalized_shape: list[int],
     w: torch.Tensor,
     b: torch.Tensor,
     eps: float,
-    use_input_shape: bool,
 ):
   """Layer Normalization with high-level function boundary enabled.
 
   Args:
     x (torch.Tensor): Input tensor for Layer Normalization, with BCHW shape.
+    normalized_shape (list[int]): Input shape from an expected input of size,
+      same as https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html.
     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.
-    use_input_shape (bool): Whether to use the input shape to determine the
-      dimension of normalization.
 
   Returns:
     The output tensor of Layer Normalization.
@@ -229,12 +219,6 @@ def layer_norm_with_hlfb(
       attr={"num_groups": 1, "epsilon": eps, "channel_axis": 1},
   )
   x, w, b = builder.mark_inputs(x, w, b)
-  if use_input_shape:
-    normalized_shape = x.shape
-    w = w.broadcast_to(x.shape)
-    b = b.broadcast_to(x.shape)
-  else:
-    normalized_shape = w.shape
   y = F.layer_norm(x, normalized_shape, w, b, eps=eps)
   y = builder.mark_outputs(y)
   return y

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241006.dist-info → ai_edge_torch_nightly-0.3.0.dev20241007.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20241006
+Version: 0.3.0.dev20241007
 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.dev20241006.dist-info → ai_edge_torch_nightly-0.3.0.dev20241007.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=VXOiPAspd6Q0ZIwnjLplRQtIy97iE7lFvqOIWCFNigI,706
+ai_edge_torch/version.py,sha256=G4658C8Iyg3OGjdMypYdKW18bjZ3ysuHhek8VfS6mCc,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
@@ -57,7 +57,7 @@ ai_edge_torch/generative/examples/openelm/verify.py,sha256=VkigoqhAr8ew95neb3Tif
 ai_edge_torch/generative/examples/phi/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py,sha256=rkbTtMaqSVG48cm-NTxR_LDgZmXAEBqayTm9O49oMXc,2171
 ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=3go690yX6PFeXMdpY7y4JZorAwxX0HT_b_pKZieauvk,2169
-ai_edge_torch/generative/examples/phi/phi2.py,sha256=CQ55KfOdoOM43CxF7yNQsgq8b-j0S50bXpxYzgq-keM,3418
+ai_edge_torch/generative/examples/phi/phi2.py,sha256=nbivDwZREd-sypy_ittO59-yaAdPvHv1YEV6Fo5buCo,3341
 ai_edge_torch/generative/examples/phi/phi3.py,sha256=GkHOaYfsFEbHvfZCaLlb3Us_h19ezqPDUakoz_DiG9A,7123
 ai_edge_torch/generative/examples/phi/verify.py,sha256=YPFCdbnfmvq38fbpBNr0kHPfSZo4p3_6WkLJAW3pLPo,2177
 ai_edge_torch/generative/examples/phi/verify_phi3.py,sha256=kVYaBVvddfQng0IyZGxyTJEzhiPO0G4VFJm2WOc2Q94,2360
@@ -101,11 +101,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=oE8DdqLA-oWkBC2zySSCh8JNAJg_hk8-W_UoMSrgDVk,5088
+ai_edge_torch/generative/layers/builder.py,sha256=XyZS1RrnMbvypeLMfwU7h1Y4x5r4WGgOx2YGJF0OUNQ,5064
 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=xZt4xaNZJPvtdy4hfbnRencEENr689zO0WnZbhpNTIs,7137
-ai_edge_torch/generative/layers/normalization.py,sha256=cpo88JUXbF9j3sJTU4JuwOap9ryGV05C1QkPij-YQwU,6999
+ai_edge_torch/generative/layers/model_config.py,sha256=DdsdhTP5tZAtyWim-qW2m8HDBsYbs7boqSDb83vwgmE,6998
+ai_edge_torch/generative/layers/normalization.py,sha256=eKAGst9rPuyRFExMcQFJO7R3iHdCtlmjeF_lITjLhwE,6498
 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
@@ -180,8 +180,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.dev20241006.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241006.dist-info/METADATA,sha256=WCkiyqsYONZTjrZ5OS0O3V4jK8kHc9MHDI6iJiywO9k,1897
-ai_edge_torch_nightly-0.3.0.dev20241006.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20241006.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241006.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241007.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241007.dist-info/METADATA,sha256=GgrT0zBBoJAVXVdeYnu9-ZaI2uFn9zpd1BS4OEVPNSY,1897
+ai_edge_torch_nightly-0.3.0.dev20241007.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20241007.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241007.dist-info/RECORD,,