ai-edge-torch-nightly 0.5.0.dev20250514__py3-none-any.whl → 0.5.0.dev20250516__py3-none-any.whl

ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py

@@ -15,8 +15,10 @@
 
 """Example of building a TinyLlama model."""
 
+from typing import Callable, Dict
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
+import torch
 from torch import nn
 
 TENSOR_NAMES = model_builder.TENSOR_NAMES_WITH_SEPARATE_LM_HEAD
@@ -81,10 +83,15 @@ def get_fake_model_config(**kwargs) -> cfg.ModelConfig:
   return config
 
 
-def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
+def build_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=TinyLlama,
+      custom_loader=custom_loader,
   )

ai_edge_torch/generative/layers/normalization.py

@@ -28,6 +28,8 @@ class RMSNorm(torch.nn.Module):
       dim: int,
       eps: float = 1e-6,
       zero_centered_gamma=False,
+      with_scale: bool = False,
+      scale_shift: float = 1.0,
       enable_hlfb: bool = False,
   ):
     """Initialize the RMSNorm layer.
@@ -37,13 +39,22 @@ class RMSNorm(torch.nn.Module):
       eps (float): A small float value to ensure numerical stability (default:
         1e-6).
       zero_centered_gamma (bool): Whether or not gamma has an offset.
+      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.
     """
     super().__init__()
+    self.dim = dim
     self.enable_hlfb = enable_hlfb
     self.eps = eps
-    self.weight = torch.nn.Parameter(torch.ones(dim))
+    self.weight = torch.nn.Parameter(torch.ones(dim), requires_grad=False)
     self.zero_centered_gamma = zero_centered_gamma
+    self.with_scale = with_scale
+    if with_scale:
+      self.scale = torch.nn.Parameter(
+          torch.zeros((dim,), dtype=torch.float32), requires_grad=False
+      )
+    self.scale_shift = scale_shift
 
   def _norm(self, x):
     """Apply RMSNorm normalization.
@@ -70,14 +81,20 @@ class RMSNorm(torch.nn.Module):
     else:
       w = self.weight
 
+    final_scale = (
+        self.scale + self.scale_shift
+        if self.with_scale
+        else torch.ones((self.dim,), dtype=torch.float32)
+    )
     if self.enable_hlfb:
       return rms_norm_with_hlfb(
           x,
           w,
           self.eps,
+          final_scale,
       )
     else:
-      output = self._norm(x.float()).type_as(x)
+      output = self._norm(x.float()).type_as(x) * final_scale
       return output * w
 
 
@@ -104,8 +121,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.empty(dim))
-    self.bias = torch.nn.Parameter(torch.empty(dim))
+    self.weight = torch.nn.Parameter(torch.empty(dim), requires_grad=False)
+    self.bias = torch.nn.Parameter(torch.empty(dim), requires_grad=False)
 
   def forward(self, x):
     """Running the forward pass of GroupNorm layer.
@@ -140,8 +157,8 @@ class LayerNorm(torch.nn.Module):
     self.enable_hlfb = enable_hlfb
     self.normalized_shape = (dim,)
     self.eps = eps
-    self.weight = torch.nn.Parameter(torch.empty(dim))
-    self.bias = torch.nn.Parameter(torch.empty(dim))
+    self.weight = torch.nn.Parameter(torch.empty(dim), requires_grad=False)
+    self.bias = torch.nn.Parameter(torch.empty(dim), requires_grad=False)
 
   def forward(self, x):
     """Running the forward pass of LayerNorm layer.
@@ -165,6 +182,7 @@ def rms_norm_with_hlfb(
     x: torch.Tensor,
     w: torch.Tensor,
     eps: float,
+    final_scale: torch.Tensor,
 ):
   """RMS Normalization with high-level function boundary enabled.
 
@@ -172,6 +190,7 @@ def rms_norm_with_hlfb(
     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.
+    final_scale (torch.Tensor): The final scale to apply to the normalization.
 
   Returns:
     The output tensor of RMS Normalization.
@@ -185,7 +204,7 @@ def rms_norm_with_hlfb(
   def _norm(x):
     return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
 
-  output = _norm(x.float()).type_as(x)
+  output = _norm(x.float()).type_as(x) * final_scale
   out = output * w
 
   out = builder.mark_outputs(out)

ai_edge_torch/generative/layers/normalization_test.py

@@ -0,0 +1,73 @@
+# Copyright 2025 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 normalization layers."""
+
+from ai_edge_torch.generative.layers import normalization
+import torch
+from absl.testing import absltest as googletest
+from absl.testing import parameterized
+
+
+class NormalizationTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="rms_norm_test_1",
+          model_dim=10,
+          with_scale=False,
+          scale_shift=1.0,
+          enable_hlfb=False,
+          expected_values=torch.ones((10,), dtype=torch.float32),
+      ),
+      dict(
+          testcase_name="rms_norm_test_2",
+          model_dim=10,
+          with_scale=True,
+          scale_shift=2.0,
+          enable_hlfb=False,
+          expected_values=torch.ones((10,), dtype=torch.float32) * 2.0,
+      ),
+      dict(
+          testcase_name="rms_norm_test_3",
+          model_dim=10,
+          with_scale=True,
+          scale_shift=2.0,
+          enable_hlfb=True,
+          expected_values=torch.ones((10,), dtype=torch.float32) * 2.0,
+      ),
+  )
+  def test_rms_norm(
+      self,
+      model_dim: int,
+      with_scale: bool,
+      scale_shift: float,
+      enable_hlfb: bool,
+      expected_values: torch.Tensor,
+  ):
+    rms_norm = normalization.RMSNorm(
+        dim=model_dim,
+        with_scale=with_scale,
+        scale_shift=scale_shift,
+        enable_hlfb=enable_hlfb,
+    )
+
+    x = torch.ones((model_dim,), dtype=torch.float32)
+    out = rms_norm(x)
+    self.assertEqual(out.shape, (model_dim,))
+    self.assertTrue(torch.allclose(out, expected_values))
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_torch/generative/utilities/converter.py

@@ -119,6 +119,12 @@ def define_conversion_flags(
       default_transpose_kv_cache,
       'If true, the model will be converted with transposed KV cache.',
   )
+  flags.DEFINE_bool(
+      'custom_checkpoint_loader',
+      False,
+      'If true, the conversion script will use a custom checkpoint loader which'
+      ' will read a checkpoint from a remote source.',
+  )
   return flags
 
 
@@ -397,13 +403,19 @@ def _export_helper(
       )
 
       if prefill_pixel_values is not None:
-        sample_kwargs['tokens'] = prefill_tokens_list_with_pixel[i]
-        sample_kwargs['input_pos'] = prefill_input_pos_list_with_pixel[i]
-        sample_kwargs['pixel_values'] = prefill_pixel_values
+        sample_pixel_kwargs = {
+            'tokens': prefill_tokens_list_with_pixel[i],
+            'input_pos': prefill_input_pos_list_with_pixel[i],
+            'kv_cache': prefill_kv,
+            'pixel_values': prefill_pixel_values,
+        }
+        # mask should be built internally when pixel values are passed.
+        if lora is not None:
+          sample_pixel_kwargs['lora'] = lora
         converter.add_signature(
             prefill_signature_name + '_pixel',
             mod,
-            sample_kwargs=sample_kwargs,
+            sample_kwargs=sample_pixel_kwargs,
         )
 
     sample_kwargs = {

ai_edge_torch/generative/utilities/loader.py

@@ -19,10 +19,55 @@ import os
 from typing import Callable, Dict, List, Tuple
 
 from ai_edge_torch.generative.layers import model_config
+import safetensors
 from safetensors import safe_open
 import torch
 
 
+def get_custom_loader(
+    checkpoint_path: str,
+) -> Callable[[str], Dict[str, torch.Tensor]]:
+  """Returns a custom loader for the given checkpoint path.
+
+  Those customer loaders can either support state dictionary or safetensors, and
+  the actual data might be fetched from a remote source.
+
+  Args:
+    checkpoint_path (string): The path to the checkpoint.
+
+  Returns:
+    Callable[[str], Dict[str, torch.Tensor]]: The custom loader.
+
+  Raises:
+    ValueError: If the checkpoint format is not supported.
+  """
+
+  if os.path.splitext(checkpoint_path)[1] in [".bin", ".pt", ".ckpt"]:
+    return lambda path: torch.load(path, weights_only=True)
+  if checkpoint_path.endswith(".safetensors"):
+    return safetensors.torch.load_file
+  raise ValueError(f"Unsupported checkpoint format: {checkpoint_path}")
+
+
+def maybe_get_custom_loader(
+    checkpoint_path: str,
+    use_custom_loader: bool = False,
+) -> Callable[[str], Dict[str, torch.Tensor]] | None:
+  """Returns a custom loader for the given checkpoint path.
+
+  If use_custom_loader is True, the function will return a custom loader.
+  Otherwise, it will return None.
+
+  Args:
+    checkpoint_path (string): The path to the checkpoint.
+    use_custom_loader (bool): Whether to use a custom loader.
+
+  Returns:
+    Callable[[str], Dict[str, torch.Tensor]] | None: The custom loader.
+  """
+  return get_custom_loader(checkpoint_path) if use_custom_loader else None
+
+
 def load_safetensors(full_path: str):
   """Loads safetensors into a single state dictionary.
 

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.5.0.dev20250514"
+__version__ = "0.5.0.dev20250516"

{ai_edge_torch_nightly-0.5.0.dev20250514.dist-info → ai_edge_torch_nightly-0.5.0.dev20250516.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250514
+Version: 0.5.0.dev20250516
 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.5.0.dev20250514.dist-info → ai_edge_torch_nightly-0.5.0.dev20250516.dist-info}/RECORD

@@ -1,12 +1,12 @@
-ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,1208
+ai_edge_torch/__init__.py,sha256=lemyLCNoGYRnJsmDuGZu7qOqLbLqG6CGDFtu3ue1syU,1290
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=wxjSFq_rBSxSqbUE8E8EJTCkgvgaRLjq_ZuAM-IZpCU,5606
-ai_edge_torch/version.py,sha256=ZvSDZpKkUslpMEN4pPp4xI6n8g3mHZMdfIcYeWth5Dg,706
+ai_edge_torch/version.py,sha256=qsmmOMVNJ3QYndWFHn1wZqGlFpjk3G1-KHlQvjpBSFg,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=QVugYVfbyaeBgSKKbhFzHG5oXA7t3M-40JcpcdSu6W8,5436
+ai_edge_torch/_convert/conversion.py,sha256=jidl5IOb3MhUPqhMLBNFRSzkqQyi3Y0R0ua-vOSahm0,6082
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
-ai_edge_torch/_convert/converter.py,sha256=075F8LRewk_033Ebsnft7FJr3KgtIbtZ_-8udIPy6ho,9980
+ai_edge_torch/_convert/converter.py,sha256=6MLKELzAwFoiXv-b7KRYi7gc7Z57XOeowcz9ArIl9TM,12100
 ai_edge_torch/_convert/signature.py,sha256=-YKJdLk-eNEHfhdPCtcQVtZf915SoVePEFxKXPPf16c,2572
 ai_edge_torch/_convert/to_channel_last_io.py,sha256=_31phf7TYgZY2ftpNbrdlB1RhDium1lz_BXEQ6IsMFc,2893
 ai_edge_torch/_convert/fx_passes/__init__.py,sha256=jbRCZmSduG_1qmngaEEtbofAyL1PKZ8P1uxzzsXQhsw,1253
@@ -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=6vQa0UJn2L3qxR967_-vkfLrO7JdrLLBk4BfguOtHRI,17874
+ai_edge_torch/_convert/test/test_convert.py,sha256=yQC0WZk_gzReguTOfgWWodK71jnfMjYoRF29_Kafnuw,18692
 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
@@ -52,83 +52,83 @@ ai_edge_torch/generative/custom_ops/bmm_4d.py,sha256=JmVbZCujG_wuBchma8QF3DSBfVc
 ai_edge_torch/generative/custom_ops/dynamic_update_slice.py,sha256=ZGAq2CfWZsfef5mHulsWmyUx0dDWJX6J6xPjhBrjQdM,2097
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/amd_llama_135m/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/amd_llama_135m/amd_llama_135m.py,sha256=XsDXx6k0kE_OYu_dr7GEC26jCepV1Kv39iH-kpuqA4M,2794
-ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py,sha256=hiuMFJ8QPymGMM6PiSQqQrfR4M1mblpPuDfjjabcr_w,1560
+ai_edge_torch/generative/examples/amd_llama_135m/amd_llama_135m.py,sha256=NyBlyUUk-3ksS5M2jFPeor6_1vSa8W_CofO8-lQ_4gE,2962
+ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py,sha256=s2f5TJos6rSgogqeFk0qsOpI30qsR04umk9hAAZ5918,1782
 ai_edge_torch/generative/examples/amd_llama_135m/verify.py,sha256=o13NkFlBgawBsjdJup05VMUjAPvDRAmig6VyEkX8q6U,2426
 ai_edge_torch/generative/examples/deepseek/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
-ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py,sha256=l0OrPGmX8WscuG9MIgtd0sqR4BeReNAu7fADzyPbnZw,1580
-ai_edge_torch/generative/examples/deepseek/deepseek.py,sha256=yhS_i2kR0GJWpWciCt4p9Z9nHYh6A5uJ8Ycy2ebFN9w,2909
+ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py,sha256=xTPfT3Mt_4bMfGkrqDKatLecZOuaE0WhxXs3uAsO_uU,1749
+ai_edge_torch/generative/examples/deepseek/deepseek.py,sha256=afKPeEjRUkLf5uhImvxtOdHrK2edfJ_R4lx92etEQpQ,3069
 ai_edge_torch/generative/examples/deepseek/verify.py,sha256=iYldze-pvZGvPkkqr6zA7EmitPnH9sXkzjNVx353IcE,2403
 ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=RRilUl2Ui08R9gy1Ua0jnaXNCrIJJb-oztgP62G3mX4,1526
-ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=9ozSw2-xuf5Wfh1HeLDTP3wJxxUZmrD3An1njJPMpdI,1594
-ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=6ImjTzJcq6JoKz2Z-z8pjv5BsRu5nUeEsTK3IPs3xgI,3521
-ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=JQLLiHNVBM9jOrZqUF0EmgAwtDD0yTRlmIbLaWM7qTg,11557
+ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=t2qZTjyM2imPenb14fzbQ-CHj5Cejw4M5xfEZpgX6Uc,1748
+ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=Yj-b4S9BNxArnGjruRIymCiWrlf7ZvwiG6keTVGldk4,1816
+ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=HqpNgJYL3X91Bpl9dAQsWEmaXJjDXGuGBVeyqK5hGTk,3682
+ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=zynxoe_9ESvTIsznpp44HUS3gVDaEltkapmjzoNOaqA,11691
 ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=ip-Gmk4CI5f0GWSdAIdrectxQWJ0t328KCsA4nfHuGg,1736
 ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=jhiyinOqPt5ZZjEadDRZt_wY5fiLSCpMo54PcxFaL_Q,1789
 ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=n7f2nF6Lin_tDvPs0JVldsuaBzo7pAwi5YAHAhlIxQg,6139
 ai_edge_torch/generative/examples/gemma3/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
-ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=JLXXn2mFEBs4DlHH_O6hpEG9KInJqsCdWy3DrgUjT1c,1827
+ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=wOrOV_jxCnjrhjC8X0-uIi0D-4aQjOfXw6XaxTSrM9k,2048
 ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=shdgLzKDUi0vyNOAsrIVAEFb3Adltsri6Rx1-wxzVf4,15089
 ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=ZorRtnbElWsctcA0nEbfwjx0C578voF7fjFEvWSR5Ck,6582
 ai_edge_torch/generative/examples/gemma3/image_encoder.py,sha256=uRoLoBWzFtQz5wFZfPCxbkvZsgPAqSkUUsV3977GbYc,5184
 ai_edge_torch/generative/examples/gemma3/verify_gemma3.py,sha256=v8oNXFICmVOtQxfO7IhZ8GnbvotEkDi9lzYHjoQyOso,2464
 ai_edge_torch/generative/examples/gemma3/verify_util.py,sha256=1vfAtayH_I_qTpqhzu6n9xnCuvhgTzhS8IzZviW2dJQ,9418
 ai_edge_torch/generative/examples/hammer/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/hammer/convert_to_tflite.py,sha256=9r8LXyaoBXYIIhhe1WQgEIjaxALQPE1dO2N6qopyWCk,1753
-ai_edge_torch/generative/examples/hammer/hammer.py,sha256=76INcjffvaNCQ02fzXcxJUW_6EKHs4sg3q1nDBbEpHE,3431
+ai_edge_torch/generative/examples/hammer/convert_to_tflite.py,sha256=XLmPuJCBJjKzMTG-mRmBX92juep2zl5yYeMrEhdqQQk,1975
+ai_edge_torch/generative/examples/hammer/hammer.py,sha256=s8arcxjETiyuERrFOvyQe_o8Lvr82gxmOIJO1hr2Dcs,3704
 ai_edge_torch/generative/examples/hammer/verify.py,sha256=MkzAGkbPy4LKRhyCDm1cw-9jUt4VUxLPdwK_25fCGSE,2705
 ai_edge_torch/generative/examples/llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/llama/convert_to_tflite.py,sha256=nz5h4m8bVnw8P7OEtqhA_fKfvaRzxhT2_75vkFCqHmU,1735
-ai_edge_torch/generative/examples/llama/llama.py,sha256=H7I5iNhIJ55gb0-9k7g-FPcG2IlthnA9XMR8qd__5bQ,6621
+ai_edge_torch/generative/examples/llama/convert_to_tflite.py,sha256=4qnMyvJHqhqf9k01wEsO23BKo6tSy2KD7sHdTGimKGg,1957
+ai_edge_torch/generative/examples/llama/llama.py,sha256=TJXU9yZwxPCnuT2uwlcXVLrs5pg1P-Csv4xY5WTcf8U,7005
 ai_edge_torch/generative/examples/llama/verify.py,sha256=X7oKQi85M789ugBrOlMvzk8eSRR3Kf1Mprfl-U-WIpo,2842
 ai_edge_torch/generative/examples/moonshine/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/moonshine/convert_moonshine_to_tflite.py,sha256=7m3rYRzThRDYb-7pGnpLr3ACi4PWX07Mg20Q98ArPc4,1714
+ai_edge_torch/generative/examples/moonshine/convert_moonshine_to_tflite.py,sha256=_GkaSkregS3NWN38UGXxj4pED5gtQGaaPZx5_CZ0TVM,1657
 ai_edge_torch/generative/examples/moonshine/moonshine.py,sha256=nZ2b8u4TmsB5sgdClgAuH8E78bcTv9RCnF9666HqP2M,3394
 ai_edge_torch/generative/examples/openelm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=wRdT7bWbCX8g4TbzKbjcLx6vmKtuT5-g-ipg19hJW-M,1525
-ai_edge_torch/generative/examples/openelm/openelm.py,sha256=hPcXYHj-nBP56TOeQQejB3HRzv6yHSftHOx0OEPP5M8,4574
+ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=S7OP8PJcOQbm8AHvi_Tc3qnQuVOtjMFNlwaZQ_oirUM,1747
+ai_edge_torch/generative/examples/openelm/openelm.py,sha256=2jkIbj_G0IuFi5nXz_yAIY4qRxgWGD5rKQDTSweRV9M,4734
 ai_edge_torch/generative/examples/openelm/verify.py,sha256=4W26ZtPF5Cb9mpHYuRM4b2QB_4W76zf4WV36KzexVjs,2446
 ai_edge_torch/generative/examples/paligemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=fkP-mWrih1s-vgJ41fLt8v5JE-UOs8Zrngh6ElQ6PMw,1997
-ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=-EYUZp55dfRY1E-N0Pr3b9i5c7Tt1XvYxvsRixguVS8,5527
-ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=WB8r-e_Crog1ItBq3Zse_nUG-foFyBcJsuEG26r_Ji8,6076
-ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=SvuR97sjkBtfkerH7Hu1UXB8kCFLpEATNbPfCbNAyfo,5614
-ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=CFIjOmrn4a4Udki7l3im0JR4zTC_NttnsIr9_qWjKTY,6110
-ai_edge_torch/generative/examples/paligemma/verify.py,sha256=zrCNz_QSQU6BbaFtx-J-MqxXWcNlsAlquaHpKodsyW4,5350
+ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=Fl4k-lcpiUaJS0A1E7HVVUW7iTcZAU4FbA4KcSkO5SQ,2212
+ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=ruUTonTErvuinWsJ3pnSbvKhCnDUlupT1MW4TUwcrMY,5551
+ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=C377j2ULpPvmY5SsNLUC8jskTNNHVDH8uYOLH5W7fOU,6100
+ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=IbneN2J-ASdUg7OHVRkrUBiZ0UXyCVRJXhnDAxjozl8,5644
+ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=nxvcurGkFJcCjjgVkK59SJgp8mZ71D56bEnrjvGgPs4,6264
+ai_edge_torch/generative/examples/paligemma/verify.py,sha256=myHdeIAtVTOqb915h661CnvjvFkwmihy3Vp4UrKHb5I,6195
 ai_edge_torch/generative/examples/paligemma/verify_decoder.py,sha256=al5wMPWri4IRVWrLmCplPi6uoCzwh0vBHMGnCt-XUqo,2690
 ai_edge_torch/generative/examples/paligemma/verify_decoder2.py,sha256=tm-UfLr0YeBRVcQsWLBOMWI9JUzHmtPEbYK2vpITpqY,2534
 ai_edge_torch/generative/examples/paligemma/verify_image_encoder.py,sha256=vNm-wTT8BD6zbX6GocfP1QrVoHl0zSvuVxoXN36eeiU,3540
 ai_edge_torch/generative/examples/phi/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py,sha256=k-0ZC-_zZZmkdcc6dr1QGXfX9lDZZXRQSuc6wT0n3Is,1514
-ai_edge_torch/generative/examples/phi/convert_phi4_to_tflite.py,sha256=5KSJRySjSc89FriCOnfBabD8zRLUcGAw3L0VInuJFUY,1512
-ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=wVIdGenHTi9xUffYddN_uXWMBO2tgo1e_hU4OG_NmHA,1513
-ai_edge_torch/generative/examples/phi/phi2.py,sha256=X9MfjK8rmyRSrfNzIaKQNSgqLM5_CBH-BrLFX_7BWL8,3494
-ai_edge_torch/generative/examples/phi/phi3.py,sha256=65Dbv8cA4WFdluflHQHzgDmDFjdmc6rxMO4hQukaxKU,6978
-ai_edge_torch/generative/examples/phi/phi4.py,sha256=y3CCZCW4MnvX74d4MNERRuQBE0p5dquC2M9vDXXqnZI,5760
+ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py,sha256=kYgZAIHXolUhOyDAYDuEK7RZ5ExL1YzpqtlcZjo622c,1736
+ai_edge_torch/generative/examples/phi/convert_phi4_to_tflite.py,sha256=3y3vYlJjLjRmz4Vsq-B8YKyp0LnC2fj1LAACW3pQivI,1734
+ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=tY5uwRu-4Jxro7Z9jsDqZR9SUDWB8PR6JKfswvsUSxM,1735
+ai_edge_torch/generative/examples/phi/phi2.py,sha256=nu18YKF95yg9Mo7TjpkgjA_br5fSYqaHmw0o86b5hDQ,3654
+ai_edge_torch/generative/examples/phi/phi3.py,sha256=c2h17Gmo9zLSEEdA7BzG_Jd8p4-3JmO6ZSEWLWXDGFU,7107
+ai_edge_torch/generative/examples/phi/phi4.py,sha256=TgoRbaW27X2tYAUi_z2GCb3j6uze5POhKGchRf-5eZw,5889
 ai_edge_torch/generative/examples/phi/verify.py,sha256=YPFCdbnfmvq38fbpBNr0kHPfSZo4p3_6WkLJAW3pLPo,2177
 ai_edge_torch/generative/examples/phi/verify_phi3.py,sha256=kVYaBVvddfQng0IyZGxyTJEzhiPO0G4VFJm2WOc2Q94,2360
 ai_edge_torch/generative/examples/phi/verify_phi4.py,sha256=BoCa5kUBRHtMQ-5ql6yD4pG4xHJMyUiQlpMOWVx-JgY,2356
 ai_edge_torch/generative/examples/qwen/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/qwen/convert_to_tflite.py,sha256=eOpv3scJr4mVsJ9Obl7PBhMgd3a0T1t8dqoPp_VzZaQ,1776
-ai_edge_torch/generative/examples/qwen/qwen.py,sha256=m8APYzo9N0SXsdvCxC8HtCcbN3W7gLKkRBL-Tg0BWXU,4223
+ai_edge_torch/generative/examples/qwen/convert_to_tflite.py,sha256=TnzyARHQgmWeOdYsV9WpRj5vhKGBH0kAbp3tMj8ZCYw,1998
+ai_edge_torch/generative/examples/qwen/qwen.py,sha256=XOLq1yTbW6nyAVrYYG3qu_8Cl0A74M2hkpjOT_UhyVs,4609
 ai_edge_torch/generative/examples/qwen/verify.py,sha256=9_AyEJTeUfvhhID64Rto2bflFPyXMFokdQLsseLUMiI,2775
 ai_edge_torch/generative/examples/qwen_vl/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
-ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py,sha256=4Gntv6LBIxd0CaKkb-koLzGTdBEOGgVf3ob99lAuvuY,2196
-ai_edge_torch/generative/examples/qwen_vl/decoder.py,sha256=7RFM25tDj_b0FkpSv8RUWir8K8v9p2jMtwZmP4VAUhw,4474
-ai_edge_torch/generative/examples/qwen_vl/image_encoder.py,sha256=nHzBe_YSPnUe1d5i09v4bePQomVifzJNeUjRfprmxC0,14878
-ai_edge_torch/generative/examples/qwen_vl/qwen_vl.py,sha256=mfLFrT8NPEPh9CqlJYHwh-I2y6ST7hH_vEmbZYartHQ,7764
-ai_edge_torch/generative/examples/qwen_vl/verify.py,sha256=JUwHoC_zvcC3RC3wZ3e3e6fGmrhbgdoztjK8HGSUG8I,5044
+ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py,sha256=BM-ed7KrmPwzI3MvDs2R7P-kJgE1SK_cNVqIfXhtJjs,2411
+ai_edge_torch/generative/examples/qwen_vl/decoder.py,sha256=yt3pO0x9t39dS2RWCM-0NRLl2ImcyWRIfL3E06bDg8k,4485
+ai_edge_torch/generative/examples/qwen_vl/image_encoder.py,sha256=vMZ6v6iVrps_NSFwycgG4OPG_RVQAxa80lKrbneMkaM,15023
+ai_edge_torch/generative/examples/qwen_vl/qwen_vl.py,sha256=1Ac28olo0OJExZRyxqm7vxcf7GtXdkUwEbHvhiCHi0o,7908
+ai_edge_torch/generative/examples/qwen_vl/verify.py,sha256=4WKgAFQNQzwmeJhC8ayI5vjGj9ko6VcU2HA3VAkhHug,5812
 ai_edge_torch/generative/examples/qwen_vl/verify_decoder.py,sha256=xPWoOBLh2eK12KEhELLYymfL7xvc0chmYC98c6x37oo,2602
 ai_edge_torch/generative/examples/qwen_vl/verify_image_encoder.py,sha256=PZ392nDoJG2OmHZ_7Jet3Zu1JkN6QErxKcDc7a-PPds,3126
 ai_edge_torch/generative/examples/smollm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/smollm/convert_to_tflite.py,sha256=jTM_tndbDqzq19uLz2n71S7M81L1Y6R7oVBPsMcYGzk,1785
-ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py,sha256=wU72MzpUIi2mQ8ZODW1x4L5KZPWvuXyB-_Eqo-RKqFw,1757
-ai_edge_torch/generative/examples/smollm/smollm.py,sha256=SFE8fIJx7Y_oan0vXSmhEmI0Ib2HD3k9cyKLU_4MxfI,3807
-ai_edge_torch/generative/examples/smollm/verify.py,sha256=KpYxVz_lv61YWy6HLfwT68n0owZMvty5Rr3W7ZNWWSw,2702
+ai_edge_torch/generative/examples/smollm/convert_to_tflite.py,sha256=QVRX_ovqBQi8fKAG6PezaO1qoRvMGpVxNH-_sds0pf8,1997
+ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py,sha256=rOVYSaS68_otJcGewQSconBCPD4GhDEIIyquD4dSUWc,1979
+ai_edge_torch/generative/examples/smollm/smollm.py,sha256=OXSN0Vu1MXnWb_H-aW9acgjpeLIhPIXGq2fx7RaojcM,4080
+ai_edge_torch/generative/examples/smollm/verify.py,sha256=sH3rn1TbaCusPiUD5XlECiHY0rvoHIXALbk7ECOiinI,2720
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=kDWG6MlIGa89zC5KSRcJlw2c4ITuw8KcchtfmF55f4g,3545
-ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=5M4auM33SgCTODt0VT8TO-EVILruqGDRiNILBPeB83Y,6072
+ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=fPSg217F9xBvqMZwujCAQvYq5MRZzXTYOxjiPLqD7ZU,6102
 ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=_yk6wVoZm1_FRMFJF5URaPZNNdmMR89fwmKz81BEyao,5601
 ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=afyHXc86h-ij5zTULmZnM1h313N9VWCyIVriH6pqeSo,16368
 ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=ylqXOZhYc6XFCaNBKQw0jAnYrCtRFFQKzQzEsFIntvo,34890
@@ -150,8 +150,8 @@ ai_edge_torch/generative/examples/test_models/convert_toy_model.py,sha256=6-WaNH
 ai_edge_torch/generative/examples/test_models/toy_model.py,sha256=Crpj-vOwSViHpblXOrRJmsIn4DrHyuB3XZ8kHifb7LA,5203
 ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=-z5tkQzGHbo37eAl9sDAJuT1Egxm8xI9CZmYLcmqIfU,4761
 ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=XM-dCBW2HG6FlwwPjlJi0I_TEaVqdv7aWpFEv-XUdLc,1539
-ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=6Qhml-XB8_RjQdYN948OaSsPJNrfi-Mr7PFB73C79Ug,2828
+ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=urWkWjOaGzV2gwMXoGEs1mfHNEXfEKgwuXmQ0lrWcbM,1761
+ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=HRyq5nzoljWEWGYw0kCHAZH-GNiNHxh7E2qNoupjA-4,2988
 ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=LRu6PSw7Lqu6HGbv1tO2i0nUCqe-VkRgboA10VZ7KNg,2431
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=PFSMsA1vfBfrV9ssBCkYJNl8Hx_bLdWjN01iyjPM5jE,1094
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=myGjal5A8yIBoqgArd2k40rZmCgD1Ya369KR7182bhI,2129
@@ -166,7 +166,8 @@ ai_edge_torch/generative/layers/feed_forward_test.py,sha256=8ZGy79BBpsyS6yKKDEKr
 ai_edge_torch/generative/layers/kv_cache.py,sha256=b-7shzDaKexmvQF7P3SiAmIz4ZofjYWv3m5u71GojsA,10460
 ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
 ai_edge_torch/generative/layers/model_config.py,sha256=X_gjN5524DCDBNXsX5GrOBlkKM4UHzj_RfdCD0-VOxQ,8572
-ai_edge_torch/generative/layers/normalization.py,sha256=MbwH-n80Fob5YvjBzdqDjBizMHLzSJGYRDdbD-rL5C0,6174
+ai_edge_torch/generative/layers/normalization.py,sha256=ijwCpi22NLX-Sygwy5sK9l9WjGvbPIhZvVwoBAonWAo,7014
+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
 ai_edge_torch/generative/layers/scaled_dot_product_attention_test.py,sha256=c6JBMQsq9XeMmR1XvGEIidNsoh-YIvichXo2LwVHgr4,3301
@@ -192,9 +193,9 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=vQWmpzMkJ2hP
 ai_edge_torch/generative/test/test_quantize.py,sha256=kKJ01wscTC2t_Ylr7huO5gNKES01gm3dT1gx52z15PA,7356
 ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv9og0DDLHU,2656
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/converter.py,sha256=4zcDlhgCQQyLylH8NLgVjnelou2pW6HWJHBFYsFyHuw,15020
+ai_edge_torch/generative/utilities/converter.py,sha256=mM8Vgd6zWkOrGt4-waa8cNjJwfhhTp-VNJ306NhXrV8,15425
 ai_edge_torch/generative/utilities/export_config.py,sha256=5IvR3grlMd4mWO5c_Y4x9Fk1b1xa57MzlYNE8XUaN28,2049
-ai_edge_torch/generative/utilities/loader.py,sha256=nw2REQ9sGWDwphShfRqNFICYmwIjqLp6bDcwVmsNTtg,14067
+ai_edge_torch/generative/utilities/loader.py,sha256=y1uSkUBiR0b9U4aoCQQk9qk7ctya_vEeY28Wc0A5e2s,15504
 ai_edge_torch/generative/utilities/model_builder.py,sha256=tBfOcsI_NcneggHqkCSydYN3ZgmkzPc6nW0AJrA81wI,6461
 ai_edge_torch/generative/utilities/moonshine_loader.py,sha256=_RpFabSqtGH5PHiP3_1f6QfO14qMADUxr_HGRlVDFB0,4891
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=dqPD9qRXEWtU3ombslOC-BE2l_dMwHoCNu7NsIJhsso,36158
@@ -251,8 +252,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.5.0.dev20250514.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/METADATA,sha256=4_d1LvNhvXOHKlqYZDcBYSLdYDmoGvWMgCK5PJasNiU,2074
-ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250516.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250516.dist-info/METADATA,sha256=669y6k49WKfsyVCxQ-N-xiyLc5U2lR90qfnNDoPpedA,2074
+ai_edge_torch_nightly-0.5.0.dev20250516.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250516.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250516.dist-info/RECORD,,