ai-edge-torch-nightly 0.3.0.dev20240920__py3-none-any.whl → 0.3.0.dev20240921__py3-none-any.whl

ai_edge_torch/generative/examples/gemma/{convert_to_tflite.py → convert_gemma1_to_tflite.py}

@@ -13,14 +13,14 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Example of converting a Gemma model to multi-signature tflite model."""
+"""Example of converting a Gemma1 model to multi-signature tflite model."""
 
 import os
 import pathlib
 
 from absl import app
 from absl import flags
-from ai_edge_torch.generative.examples.gemma import gemma
+from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.utilities import converter
 
 _CHECKPOINT_PATH = flags.DEFINE_string(
@@ -51,7 +51,7 @@ _QUANTIZE = flags.DEFINE_bool(
 
 
 def main(_):
-  pytorch_model = gemma.build_2b_model(
+  pytorch_model = gemma1.build_2b_model(
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
   quant_suffix = 'q8' if _QUANTIZE.value else 'f32'

ai_edge_torch/generative/examples/gemma/{gemma.py → gemma1.py}

@@ -13,10 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Example of building a Gemma model."""
-
-import os
-import pathlib
+"""Example of building a Gemma1 model."""
 
 from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import builder
@@ -24,7 +21,6 @@ from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
-import numpy as np
 import torch
 from torch import nn
 
@@ -32,13 +28,11 @@ TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
     ff_up_proj="model.layers.{}.mlp.up_proj",
     ff_down_proj="model.layers.{}.mlp.down_proj",
     ff_gate_proj="model.layers.{}.mlp.gate_proj",
-    attn_query_proj="model.layers.{}.self_attn.q_proj",
-    attn_key_proj="model.layers.{}.self_attn.k_proj",
-    attn_value_proj="model.layers.{}.self_attn.v_proj",
+    attn_fused_qkv_proj="model.layers.{}.self_attn.qkv_proj",
     attn_output_proj="model.layers.{}.self_attn.o_proj",
     pre_attn_norm="model.layers.{}.input_layernorm",
     post_attn_norm="model.layers.{}.post_attention_layernorm",
-    embedding="model.embed_tokens",
+    embedding="embedder",
     final_norm="model.norm",
     lm_head=None,
 )
@@ -192,30 +186,3 @@ def build_2b_model(checkpoint_path: str, **kwargs) -> nn.Module:
   loader.load(model, strict=False)
   model.eval()
   return model
-
-
-def define_and_run_2b(checkpoint_path: str) -> None:
-  """Instantiates and runs a Gemma 2B model."""
-
-  current_dir = pathlib.Path(__file__).parent.resolve()
-  gemma_goldens = torch.load(current_dir / "gemma_lm_logits.pt")
-
-  kv_cache_max_len = 1024
-  model = build_2b_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
-  idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
-  tokens[0, :4] = idx
-  input_pos = torch.arange(0, kv_cache_max_len, dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(model.config)
-  output = model.forward(tokens, input_pos, kv)
-  print("comparing with goldens..")
-  assert torch.allclose(
-      gemma_goldens, output["logits"][0, idx.shape[1] - 1, :], atol=1e-02
-  )
-
-
-if __name__ == "__main__":
-  input_checkpoint_path = os.path.join(
-      pathlib.Path.home(), "Downloads/llm_data/gemma-2b"
-  )
-  define_and_run_2b(input_checkpoint_path)

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -267,29 +267,3 @@ def build_2b_model(checkpoint_path: str, **kwargs) -> nn.Module:
   loader.load(model, strict=False)
   model.eval()
   return model
-
-
-def define_and_run_2b(checkpoint_path: str) -> None:
-  """Instantiates and runs a Gemma2 2B model."""
-
-  current_dir = pathlib.Path(__file__).parent.resolve()
-  gemma2_goldens = torch.load(current_dir / "gemma2it_2b_golden.pt")
-  print("Running GEMMA 2")
-  kv_cache_max_len = 1024
-  model = build_2b_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
-  toks = torch.from_numpy(
-      np.array([2, 651, 9456, 576, 573, 3520, 3858, 603, 235248])
-  )
-  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
-  tokens[0, :9] = toks
-  input_pos = torch.arange(0, kv_cache_max_len, dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(model.config)
-  out = model.forward(tokens, input_pos, kv)
-  out_final = out["logits"][0, 8, :]
-  assert torch.allclose(gemma2_goldens, out_final, atol=1e-04)
-
-
-if __name__ == "__main__":
-  torch.set_printoptions(sci_mode=True)
-  path = os.path.join(pathlib.Path.home(), "Downloads/llm_data/gemma2-2b")
-  define_and_run_2b(path)

ai_edge_torch/generative/examples/gemma/verify_gemma1.py

@@ -0,0 +1,55 @@
+# Copyright 2024 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.
+# ==============================================================================
+
+"""Verifies the reauthored Gemma1 model."""
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.gemma import gemma1
+from ai_edge_torch.generative.examples.gemma import verify_util
+from ai_edge_torch.generative.utilities import verifier
+import kagglehub
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "What is the meaning of life?",
+    "The input prompts to generate answers.",
+)
+_MAX_NEW_TOKENS = flags.DEFINE_integer(
+    "max_new_tokens",
+    30,
+    "The maximum size of the generated tokens.",
+)
+
+
+def main(_):
+  checkpoint = kagglehub.model_download("google/gemma/pyTorch/2b-it")
+
+  verifier.log_msg("Building the reauthored model from", checkpoint)
+  reauthored_model = gemma1.build_2b_model(checkpoint)
+
+  verify_util.verify_reauthored_gemma_model(
+      checkpoint=checkpoint,
+      variant="2b",
+      reauthored_model=reauthored_model,
+      weight_filename="gemma-2b-it.ckpt",
+      generate_prompts=_PROMPTS.value,
+      forward_input_ids=[[1, 2, 3, 4]],
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/gemma/verify_gemma2.py

@@ -0,0 +1,55 @@
+# Copyright 2024 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.
+# ==============================================================================
+
+"""Verifies the reauthored Gemma2 model."""
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.gemma import gemma2
+from ai_edge_torch.generative.examples.gemma import verify_util
+from ai_edge_torch.generative.utilities import verifier
+import kagglehub
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "What is the meaning of life?",
+    "The input prompts to generate answers.",
+)
+_MAX_NEW_TOKENS = flags.DEFINE_integer(
+    "max_new_tokens",
+    30,
+    "The maximum size of the generated tokens.",
+)
+
+
+def main(_):
+  checkpoint = kagglehub.model_download("google/gemma-2/pyTorch/gemma-2-2b-it")
+
+  verifier.log_msg("Building the reauthored model from", checkpoint)
+  reauthored_model = gemma2.build_2b_model(checkpoint)
+
+  verify_util.verify_reauthored_gemma_model(
+      checkpoint=checkpoint,
+      variant="2b-v2",
+      reauthored_model=reauthored_model,
+      generate_prompts=_PROMPTS.value,
+      forward_input_ids=[[2, 651, 9456, 576, 573, 3520, 3858, 603, 235248]],
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/gemma/verify_util.py

@@ -0,0 +1,142 @@
+# Copyright 2024 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.
+# ==============================================================================
+
+"""Utility functions to verify the reauthored Gemma model."""
+
+import dataclasses
+import os
+from typing import List, Tuple
+
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+from ai_edge_torch.generative.utilities import verifier
+from gemma import config as gemma_config
+from gemma import model as gemma_model
+import torch
+
+
+@dataclasses.dataclass
+class _Output:
+  logits: torch.Tensor
+
+
+class GemmaWrapper(verifier.ModelWrapper):
+  """Gemma model wrapper for verification.
+
+  Verifier calls model.forward() with maxium sequence length (1024) expecting
+  the output has 'logits' field while Gemma gets the input tokens with the
+  actual length and returns logits in a tuple.
+
+  Verifier runs tokenizer before model.generate() while Gemma runs the tokenizer
+  inside model.generate().
+  """
+
+  def __init__(self, model: torch.nn.Module, max_new_tokens: int):
+    super().__init__(model)
+    self.max_new_tokens = max_new_tokens
+
+  def _get_actual_input_len(self, tokens: torch.Tensor) -> int:
+    for i in range(tokens.shape[1]):
+      if tokens[0, i] == 0:
+        return i
+    return tokens.shape[1]
+
+  def _get_kv_caches(
+      self, max_seq_len: int
+  ) -> List[Tuple[torch.Tensor, torch.Tensor]]:
+    config = self.model.config
+    cache_size = (1, max_seq_len, config.num_key_value_heads, config.head_dim)
+    cache = torch.zeros(cache_size)
+    return [
+        (cache.clone(), cache.clone()) for _ in range(config.num_hidden_layers)
+    ]
+
+  def forward(self, tokens: torch.Tensor) -> _Output:
+    """Forwards the model after reducing input tokens to the actual length."""
+    actual_input_len = self._get_actual_input_len(tokens)
+    input_pos = torch.arange(0, actual_input_len, dtype=torch.long)
+    mask_cache = attn_utils.build_causal_mask_cache(tokens.shape[1])
+    _, logits = self.model.forward(
+        input_token_ids=tokens[0, :actual_input_len].unsqueeze(0),
+        input_positions=input_pos,
+        kv_write_indices=None,
+        kv_caches=self._get_kv_caches(tokens.shape[1]),
+        mask=mask_cache.index_select(2, input_pos),
+        output_positions=input_pos,
+        temperatures=None,
+        top_ps=torch.tensor([1.0], dtype=torch.float),
+        top_ks=torch.tensor([1], dtype=torch.long),
+    )
+    return _Output(logits.float())
+
+  def generate(self, tokens: torch.Tensor) -> torch.Tensor:
+    """Generates the response after decoding the tokens into a string."""
+    prompts = self.model.tokenizer.decode(tokens[0].tolist())
+    response = self.model.generate(
+        prompts, device="cpu", output_len=self.max_new_tokens, top_k=1
+    )
+    return torch.tensor([self.model.tokenizer.encode(prompts + response)])
+
+
+class TokenizerWrapper(torch.nn.Module):
+  """Tokenizer wrapper for verification.
+
+  Verifier expects the tokenizer to handle tokens in torch.Tensor while Gemma
+  tokenizer expects tokens in a list.
+  """
+
+  def __init__(self, tokenizer: torch.nn.Module):
+    super().__init__()
+    self.tokenizer = tokenizer
+
+  def encode(self, text: str, **_) -> torch.Tensor:
+    """Adds one more dimension to the output of the tokenizer."""
+    return torch.tensor([self.tokenizer.encode(text)])
+
+  def decode(self, tokens: torch.Tensor) -> str:
+    """Decodes the token sequence after converting to a list."""
+    return self.tokenizer.decode(tokens.tolist())
+
+
+def verify_reauthored_gemma_model(
+    checkpoint: str,
+    variant: str,
+    reauthored_model: torch.nn.Module,
+    generate_prompts: List[str],
+    forward_input_ids: List[List[int]],
+    weight_filename: str = "model.ckpt",
+    tokenizer_filename: str = "tokenizer.model",
+    max_new_tokens: int = 20,
+    rtol: float = 1e-05,
+    atol: float = 1e-05,
+):
+  """Verifies the reauthored Gemma model against the original model."""
+  config = gemma_config.get_model_config(variant)
+  config.tokenizer = os.path.join(checkpoint, tokenizer_filename)
+  # Use float32 to be compatible with the reauthored model.
+  config.dtype = torch.float32
+
+  verifier.log_msg("Loading the original model from", checkpoint)
+  original_model = gemma_model.GemmaForCausalLM(config).eval()
+  original_model.load_weights(os.path.join(checkpoint, weight_filename))
+
+  verifier.verify_reauthored_model(
+      original_model=GemmaWrapper(original_model, max_new_tokens),
+      reauthored_model=reauthored_model,
+      tokenizer=TokenizerWrapper(original_model.tokenizer),
+      generate_prompts=generate_prompts,
+      forward_input_ids=forward_input_ids,
+      rtol=rtol,
+      atol=atol,
+  )

ai_edge_torch/generative/examples/openelm/verify.py

@@ -55,7 +55,7 @@ def main(_):
       original_model=wrapper_model,
       reauthored_model=reauthored_model,
       tokenizer=tokenizer,
-      prompts=_PROMPTS.value,
+      generate_prompts=_PROMPTS.value,
   )
 
 

ai_edge_torch/generative/examples/phi/verify.py

@@ -54,7 +54,7 @@ def main(_):
       original_model=wrapper_model,
       reauthored_model=reauthored_model,
       tokenizer=tokenizer,
-      prompts=_PROMPTS.value,
+      generate_prompts=_PROMPTS.value,
       atol=1e-03,
   )
 

ai_edge_torch/generative/examples/smollm/verify.py

@@ -51,7 +51,7 @@ def main(_):
       original_model=wrapper_model,
       reauthored_model=reauthored_model,
       tokenizer=tokenizer,
-      prompts=_PROMPTS.value,
+      generate_prompts=_PROMPTS.value,
       atol=1e-04,
   )
 

ai_edge_torch/generative/examples/tiny_llama/verify.py

@@ -53,7 +53,7 @@ def main(_):
       original_model=wrapper_model,
       reauthored_model=reauthored_model,
       tokenizer=tokenizer,
-      prompts=_PROMPTS.value,
+      generate_prompts=_PROMPTS.value,
       atol=1e-04,
   )
 

ai_edge_torch/generative/layers/feed_forward.py

@@ -18,7 +18,6 @@ from typing import Callable, Optional
 
 import torch
 from torch import nn
-import torch.nn.functional as F
 
 
 class SequentialFeedForward(nn.Module):

ai_edge_torch/generative/quantize/example.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 
 import ai_edge_torch
-from ai_edge_torch.generative.examples.gemma import gemma
+from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.quantize import quant_recipes
 import numpy as np
 import torch
@@ -22,8 +22,8 @@ import torch
 
 def main():
   # Build a PyTorch model as usual
-  config = gemma.get_fake_model_config()
-  model = gemma.Gemma(config)
+  config = gemma1.get_fake_model_config()
+  model = gemma1.Gemma(config)
   idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
   tokens = torch.full((1, 10), 0, dtype=torch.int, device="cpu")
   tokens[0, :4] = idx

ai_edge_torch/generative/test/test_model_conversion_large.py

@@ -17,7 +17,7 @@
 
 import ai_edge_torch
 from ai_edge_torch import config as ai_edge_config
-from ai_edge_torch.generative.examples.gemma import gemma
+from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.examples.gemma import gemma2
 from ai_edge_torch.generative.examples.openelm import openelm
 from ai_edge_torch.generative.examples.phi import phi2
@@ -82,9 +82,9 @@ class TestModelConversion(googletest.TestCase):
       ai_edge_config.Config.use_torch_xla,
       reason="tests with custom ops are not supported on oss",
   )
-  def test_gemma(self):
-    config = gemma.get_fake_model_config()
-    pytorch_model = gemma.Gemma(config).eval()
+  def test_gemma1(self):
+    config = gemma1.get_fake_model_config()
+    pytorch_model = gemma1.Gemma(config).eval()
     self._test_model(
         config, pytorch_model, "serving_default", atol=1e-2, rtol=1e-5
     )

ai_edge_torch/generative/utilities/verifier.py

@@ -19,7 +19,6 @@ import datetime
 from typing import List, Optional, Union
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
-import numpy as np
 import torch
 import transformers
 
@@ -126,7 +125,7 @@ def generate(
 def verify_with_input_ids(
     original_model: ModelWrapper,
     reauthored_model: torch.nn.Module,
-    input_ids: torch.Tensor = torch.from_numpy(np.array([[1, 2, 3, 4]])).int(),
+    input_ids: List[int],
     kv_cache_max_len: int = 1024,
     rtol: float = 1e-05,
     atol: float = 1e-05,
@@ -139,7 +138,7 @@ def verify_with_input_ids(
     original_model (ModelWrapper): The original model.
     reauthored_model (torch.nn.Module): The model reauthored with ai_edge_torch
       Generative API.
-    input_ids (torch.Tensor): The input token IDs to forward.
+    input_ids (List[int]): The input token IDs to forward with.
     kv_cache_max_len (int): The maximum sequence length of the KV cache.
     rtol (float): The relative tolerance for the comparison.
     atol (float): The absolute tolerance for the comparison.
@@ -148,18 +147,17 @@ def verify_with_input_ids(
     True if the model reauthored generates the same output of the original.
   """
   tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
-  input_ids_len = input_ids.shape[1]
-  tokens[0, :input_ids_len] = input_ids
+  tokens[0, : len(input_ids)] = torch.tensor([input_ids]).int()
 
   log_msg("Forwarding the original model...")
   outputs_original = original_model.forward(tokens)
-  logits_original = outputs_original.logits[0, input_ids_len - 1, :]
+  logits_original = outputs_original.logits[0, len(input_ids) - 1, :]
   log_msg("logits_original: ", logits_original)
 
   log_msg("Forwarding the reauthored model...")
   kv_cache = kv_utils.KVCache.from_model_config(reauthored_model.config)
   outputs_reauthored = forward(reauthored_model, tokens, kv_cache)
-  logits_reauthored = outputs_reauthored[0][0, input_ids_len - 1, :]
+  logits_reauthored = outputs_reauthored[0][0, len(input_ids) - 1, :]
   log_msg("logits_reauthored:", logits_reauthored)
 
   return torch.allclose(
@@ -208,7 +206,8 @@ def verify_reauthored_model(
     original_model: ModelWrapper,
     reauthored_model: torch.nn.Module,
     tokenizer: torch.nn.Module,
-    prompts: List[str],
+    generate_prompts: List[str],
+    forward_input_ids: List[List[int]] = [[1, 2, 3, 4]],
     rtol: float = 1e-05,
     atol: float = 1e-05,
 ):
@@ -227,22 +226,25 @@ def verify_reauthored_model(
     reauthored_model (torch.nn.Module): The model reauthored with ai_edge_torch
       Generative API.
     tokenizer (torch.nn.Module): The tokenizer.
-    prompts (List[str]): List of the input prompts to generate answers.
+    generate_prompts (List[str]): List of the input prompts to generate answers.
+    forward_input_ids (List[torch.Tensor]): List if ihe input token IDs to
+      forward with.
     rtol (float): The relative tolerance for the comparison.
     atol (float): The absolute tolerance for the comparison.
   """
-  log_msg("Verifying the reauthored model with an arbitrary input...")
-  if verify_with_input_ids(
-      original_model, reauthored_model, rtol=rtol, atol=atol
-  ):
-    log_msg("PASS")
-  else:
-    log_msg("FAILED")
+  for input_ids in forward_input_ids:
+    log_msg("Verifying the reauthored model with input IDs:", input_ids)
+    if verify_with_input_ids(
+        original_model, reauthored_model, input_ids, rtol=rtol, atol=atol
+    ):
+      log_msg("PASS")
+    else:
+      log_msg("FAILED")
 
-  for p in prompts:
-    log_msg("Verifying the reauthored model with prompts:", p)
+  for prompts in generate_prompts:
+    log_msg("Verifying the reauthored model with prompts:", prompts)
     if verify_model_with_prompts(
-        original_model, reauthored_model, tokenizer, p
+        original_model, reauthored_model, tokenizer, prompts
     ):
       log_msg("PASS")
     else:

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20240920.dist-info → ai_edge_torch_nightly-0.3.0.dev20240921.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20240920
+Version: 0.3.0.dev20240921
 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.dev20240920.dist-info → ai_edge_torch_nightly-0.3.0.dev20240921.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=-oH0R07HZpydzqltOWclHB1dbcc4VycTlZcnDYtS89g,706
+ai_edge_torch/version.py,sha256=t9zajdsiowClI2fG0RkKVonPF-SUx9UBuUDOEZFU9y4,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
@@ -39,22 +39,25 @@ ai_edge_torch/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrK
 ai_edge_torch/generative/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
+ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=evmUj_4yygQthSRU-ke-Xn1qFNDCZKbegqINWfruKwU,2184
 ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=RZDs6oY-NLYrPNtfuJDweIHzGUL2kzpIc3AW_1p8gGg,2186
-ai_edge_torch/generative/examples/gemma/convert_to_tflite.py,sha256=t8Qg10obnEzeoMeyHnZhyNBN7G85SGy-au8Y8nehq8E,2181
-ai_edge_torch/generative/examples/gemma/gemma.py,sha256=hjpSPzEjPHuxwRJ-vHHtCCf2PSTnm30Mp0ajYYtDivo,7489
-ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=gCLOti-4xHunjphNBbx9St6faRteSakm8Oex6R1Xek0,10272
+ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=cahMzvJNJfShIw4uqoBRX5iBZrI3rvsha6wpNHzeYJ0,6369
+ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=KsM6VlzluTqbodG24IFr3biPxBrLay0z0gmnG0bcU2U,9277
+ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=B14IR4mIw6qBVUbiIRdfdUzHMCIJCJ0RFPsYOxA46qc,1776
+ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=kSzn1ITJXqrtNQax1m4WTAnC3gMxBpcvGA7-xTO1Iuw,1802
+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/verify.py,sha256=BvK4c8jodQBy2l3NnvCjlBB0qaA7EYwPNKklvFR4k_o,2103
+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=5bKONolW8JIsQAzMHIvh_OSytoJVVJqDZEcxjhciFnI,2136
+ai_edge_torch/generative/examples/phi/verify.py,sha256=QPYX6weEZGMEXt_Vb2hNARPAECQBKzx-KCivd4dzOrw,2145
 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
-ai_edge_torch/generative/examples/smollm/verify.py,sha256=wsoy3CaHZhrdJjkJJYir7xxxwgCvLprMnh8QxT0hEkc,2013
+ai_edge_torch/generative/examples/smollm/verify.py,sha256=G2dAcl-VhAbx1E1PEqM6hpzPF24HqFZaz7UBEpJSQ3w,2022
 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=tL6w2dr6VP66IXjSKo9StDNP-wl0RO3fh6dIliiYlFA,4656
@@ -80,14 +83,14 @@ ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=
 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=ekxd8efjMgEvauUu3PidWOC-DszPHn5sqU753F7sJIM,2201
 ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=tlWpa7Aun3u3w5b-9EBtW7olhmSf8W-tn5bKUIwC-ys,6044
-ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=27oBf706_AKX7amfp2THF9J0G3AUEEecGaXv025idKA,2086
+ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=LUChL5tA7FHL_DlTg5QKvGInmH9AwVVw9a-omcndiz8,2095
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=jrzCB3ZyY_t5jJM1e2Czdt3DjAIL43R0_a-T-I7wOzw,1155
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=hhxSQvkDMv0isZJhmuLiod66ZODaJ8uSPSVTJVHBabQ,1931
 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/feed_forward.py,sha256=dfS1psdmomgs4EbwzkYyV_xx1xl3P1lU-3GoS8m0Avw,4221
+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
@@ -98,7 +101,7 @@ ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=c8rtlfDaeKmUfiiTKPmQhNW
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
 ai_edge_torch/generative/layers/unet/model_config.py,sha256=8ze9kVWMuyZVQcgK7hWYw9TM1W9lXD-2j0iMHlxoGX4,9267
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/quantize/example.py,sha256=n_YFFP3dpKjeNKYZicDGL5LqtjqwhYEIaDrC6-Ci2vE,1539
+ai_edge_torch/generative/quantize/example.py,sha256=tlACaRsz6lqOxakzpXVFJZYfFKOiFqetcYVJqWVRdPE,1542
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=n1Fm8BFC8gJa_oiwwAOOghJyHtOXYZ4q-5ZRy4pHrIw,1957
 ai_edge_torch/generative/quantize/quant_recipe.py,sha256=tKnuJq6hPD23JPCB9nPAlE1UHAwdbChkgPShiVaz4CE,5156
 ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=4fgmP_GgeiFUOkIaC9ZZXC12eO3DQZdrWDXRz5YXiwU,2270
@@ -108,7 +111,7 @@ ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudj
 ai_edge_torch/generative/test/test_kv_cache.py,sha256=W6Bh0gYDzmwb0j9HdD5_D7Z7FPToP2HSyFrmwIXuFqo,3793
 ai_edge_torch/generative/test/test_loader.py,sha256=8y74ChO3CZCfEi1eCf3-w47kRgAI4qPYCXpi8rTQXMA,3378
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=s-EVLOQGjIeVtgNI8Ggs37pkRdErAliT6NhrrFigPOE,5459
-ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=PtePuBqVMLjxq2cDIIXXqaz7zsn3R19oilFyIVJRFi8,4490
+ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=mAK8Pm4mgGyilDSBtFazCRDetoqYKKB0sGC83MPKE0M,4494
 ai_edge_torch/generative/test/test_quantize.py,sha256=8geJhKwYBU20m0mdGPD1BUFwQ0lZKNtCB04SOLO18y4,5980
 ai_edge_torch/generative/test/utils.py,sha256=YvEhO2HIj1LkBs5du1UxY-cGRW9HMyAYsOUhgsTrTpA,1796
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
@@ -116,7 +119,7 @@ ai_edge_torch/generative/utilities/converter.py,sha256=MQUg2ZLmfk_2csWmQWKD_II0b
 ai_edge_torch/generative/utilities/loader.py,sha256=b9iotIhVDX-Zc9XjIDUaLxnV395AyBnkQe3dV5YA7Co,13297
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=dqPD9qRXEWtU3ombslOC-BE2l_dMwHoCNu7NsIJhsso,36158
 ai_edge_torch/generative/utilities/t5_loader.py,sha256=tEsfy8-ymzbbjOIc-oesXF3yGyyWtJgFXn2s7VOavt8,16961
-ai_edge_torch/generative/utilities/verifier.py,sha256=7DoYtkilz4wjWnXfdydIGNgTG1udZIydFxdbpIcKbMQ,8625
+ai_edge_torch/generative/utilities/verifier.py,sha256=vU9KgmFS7I9jNS_3H2SWROx-rbNqtMKgQC2MRhdqQ4g,8803
 ai_edge_torch/hlfb/__init__.py,sha256=sH4um75na-O8tzxN6chFyp6Y4xnexsE7kUQpZySv6dE,735
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=cjTprggj_cuktSCm7-A25e7Shop3k63ylp7sdZmtZ8o,4790
 ai_edge_torch/hlfb/mark_pattern/passes.py,sha256=pjkKcI1nHECPluAt87cFBrt1DP0f3ge7rHq1NhCkBIE,1936
@@ -163,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.dev20240920.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20240920.dist-info/METADATA,sha256=m60oD-H8W2EMVolDGw02tMYcKDrotaTaLtsZwzr_Kyk,1859
-ai_edge_torch_nightly-0.3.0.dev20240920.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20240920.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20240920.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20240921.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20240921.dist-info/METADATA,sha256=SWy7BhOQDe0_SBF17deNndzt1bEYy7iXUxy0KznIPYM,1859
+ai_edge_torch_nightly-0.3.0.dev20240921.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20240921.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20240921.dist-info/RECORD,,