ai-edge-torch-nightly 0.3.0.dev20241220__py3-none-any.whl → 0.3.0.dev20241224__py3-none-any.whl

ai_edge_torch/generative/examples/gemma/gemma1.py

@@ -72,12 +72,13 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       pre_attention_norm_config=norm_config,
       post_attention_norm_config=norm_config,
   )
+  embedding_dim = 2048
   config = cfg.ModelConfig(
       vocab_size=256000,
       num_layers=18,
       max_seq_len=8192,
-      embedding_dim=2048,
-      embedding_scale=2048**0.5,
+      embedding_dim=embedding_dim,
+      embedding_scale=embedding_dim**0.5,
       kv_cache_max_len=kv_cache_max_len,
       block_configs=block_config,
       final_norm_config=norm_config,

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -15,7 +15,7 @@
 
 """Example of building a Gemma2 model."""
 
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
 
 from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import builder
@@ -136,29 +136,45 @@ class Gemma2(nn.Module):
         f"Cannot forward sequence of length {seq_len}, max seq length is only"
         f" {self.config.max_seq_len}"
     )
-    assert len(self.transformer_blocks) == len(kv_cache.caches), (
-        "The number of transformer blocks and the number of KV cache entries"
-        " must be the same."
-    )
 
+    # token embeddings of shape (b, t, n_embd)
+    input_embeds = self.tok_embedding(tokens)
     # RoPE parameters are the same for all blocks. Use the first layer.
     attn_config = self.config.block_config(0).attn_config
     n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
     rope = rotary_pos_emb.build_rope(
         input_pos, n_elem, attn_config.head_dim, attn_config.rotary_base
     )
+    mask = [self.get_attention_mask(
+        self.config.block_config(i).attn_config.attn_type, input_pos
+    ) for i in range(self.config.num_layers)]
 
-    # token embeddings of shape (b, t, n_embd)
-    x = self.tok_embedding(tokens)
-    x = x * (self.config.embedding_dim**0.5)
+    return self._forward_with_embeds(
+        input_embeds, rope, mask, input_pos, kv_cache, export_config
+    )
+
+  def _forward_with_embeds(
+      self,
+      input_embeds: torch.Tensor,
+      rope: Tuple[torch.Tensor, torch.Tensor],
+      mask: List[torch.Tensor],
+      input_pos: torch.Tensor,
+      kv_cache: kv_utils.KVCache,
+      export_config: Optional[model_builder.ExportConfig] = None,
+  ) -> dict[torch.Tensor, kv_utils.KVCache]:
+    """Forwards the model with input embeddings."""
+    assert len(self.transformer_blocks) == len(kv_cache.caches), (
+        "The number of transformer blocks and the number of KV cache entries"
+        " must be the same."
+    )
 
+    if self.config.embedding_scale is not None:
+      input_embeds = input_embeds * self.config.embedding_scale
+    x = input_embeds
     updated_kv_entries = []
     for i, block in enumerate(self.transformer_blocks):
-      mask = self.get_attention_mask(
-          block.config.attn_config.attn_type, input_pos
-      )
       kv_entry = kv_cache.caches[i] if kv_cache else None
-      x, kv_entry = block(x, rope, mask, input_pos, kv_entry)
+      x, kv_entry = block(x, rope, mask[i], input_pos, kv_entry)
       if kv_entry:
         updated_kv_entries.append(kv_entry)
     updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entries))
@@ -227,11 +243,13 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
     )
 
   num_layers = 26
+  embedding_dim = 2304
   config = cfg.ModelConfig(
       vocab_size=256000,
       num_layers=num_layers,
       max_seq_len=8192,
-      embedding_dim=2304,
+      embedding_dim=embedding_dim,
+      embedding_scale=embedding_dim**0.5,
       kv_cache_max_len=kv_cache_max_len,
       block_configs=[get_block_config(i) for i in range(num_layers)],
       final_norm_config=norm_config,
@@ -248,6 +266,7 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
   config.num_layers = 2
   config.max_seq_len = 2 * kv_cache_max_len
   config.embedding_dim = 128
+  config.embedding_scale = config.embedding_dim**0.5
   config.block_configs = config.block_configs[: config.num_layers]
   for block_config in config.block_configs:
     block_config.attn_config.num_heads = 4

ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py

@@ -29,9 +29,15 @@ from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities.model_builder import ExportConfig
 import torch
 
+_VERSION = flags.DEFINE_enum(
+    'version',
+    '2',
+    ['1', '2'],
+    'The version of PaliGemma model to verify.',
+)
 _CHECKPOINT_PATH = flags.DEFINE_string(
     'checkpoint_path',
-    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/paligemma-3b-224'),
+    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/paligemma2-3b-224'),
     'The path to the model checkpoint, or directory holding the checkpoint.',
 )
 _TFLITE_PATH = flags.DEFINE_string(
@@ -63,10 +69,12 @@ _QUANTIZE = flags.DEFINE_bool(
 
 def main(_):
   pytorch_model = paligemma.build_model(
-      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
+      _CHECKPOINT_PATH.value,
+      version=int(_VERSION.value),
+      kv_cache_max_len=_KV_CACHE_MAX_LEN.value,
   )
   quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
-  output_filename = f'paligemma_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
+  output_filename = f'paligemma{_VERSION.value}_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
       tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -19,6 +19,7 @@ from typing import Optional
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
@@ -54,6 +55,7 @@ class Decoder(model_builder.DecoderOnlyModel):
       kv_cache: kv_utils.KVCache,
       input_embeds: torch.Tensor = None,
       export_config: Optional[model_builder.ExportConfig] = None,
+      called_by_generate: bool = True,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if input_embeds is None:
       return super().forward(tokens, input_pos, kv_cache)
@@ -61,8 +63,12 @@ class Decoder(model_builder.DecoderOnlyModel):
     assert input_embeds is not None
 
     repo_pos = input_pos + 1  # PaliGemma position is 1-based.
-    cos, sin = self.rope_cache
-    rope = (cos.index_select(0, repo_pos), sin.index_select(0, repo_pos))
+    # ROPE parameters for all attn_configs are the same. Take the first one.
+    attn_config = self.config.block_config(0).attn_config
+    n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
+    rope = rotary_pos_emb.build_rope(
+        repo_pos, n_elem, attn_config.head_dim, attn_config.rotary_base
+    )
 
     # The first part of input_embeds are image embeddings. Diagonal causal mask
     # doesn't work here.
@@ -70,7 +76,7 @@ class Decoder(model_builder.DecoderOnlyModel):
     mask = torch.zeros(embeds_len, self.config.kv_cache_max)
     mask[:, embeds_len:] = float("-inf")
 
-    return self.forward_with_embeds(
+    return self._forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache
     )
 
@@ -108,12 +114,13 @@ def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       pre_attention_norm_config=norm_config,
       post_attention_norm_config=norm_config,
   )
+  embedding_dim = 2048
   config = cfg.ModelConfig(
       vocab_size=257216,
       num_layers=18,
       max_seq_len=8192,
-      embedding_dim=2048,
-      embedding_scale=2048**0.5,
+      embedding_dim=embedding_dim,
+      embedding_scale=embedding_dim**0.5,
       kv_cache_max_len=kv_cache_max_len,
       block_configs=block_config,
       final_norm_config=norm_config,
@@ -130,6 +137,8 @@ def get_fake_decoder_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
   config.vocab_size = 128
   config.num_layers = 2
   config.max_seq_len = 2 * kv_cache_max_len
+  config.embedding_dim = 128
+  config.embedding_scale = 128**0.5
   return config
 
 

ai_edge_torch/generative/examples/paligemma/decoder2.py

@@ -0,0 +1,174 @@
+# 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.
+# ==============================================================================
+
+"""Example of building a decoder of PaliGemma2 3B model which is Gemma2."""
+
+from typing import Optional
+
+from ai_edge_torch.generative.examples.gemma import gemma2
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
+from ai_edge_torch.generative.utilities import model_builder
+import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
+
+TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj="language_model.model.layers.{}.mlp.up_proj",
+    ff_down_proj="language_model.model.layers.{}.mlp.down_proj",
+    ff_gate_proj="language_model.model.layers.{}.mlp.gate_proj",
+    attn_query_proj="language_model.model.layers.{}.self_attn.q_proj",
+    attn_key_proj="language_model.model.layers.{}.self_attn.k_proj",
+    attn_value_proj="language_model.model.layers.{}.self_attn.v_proj",
+    attn_output_proj="language_model.model.layers.{}.self_attn.o_proj",
+    pre_attn_norm="language_model.model.layers.{}.input_layernorm",
+    post_attn_norm="language_model.model.layers.{}.post_attention_layernorm",
+    pre_ff_norm="language_model.model.layers.{}.pre_feedforward_layernorm",
+    post_ff_norm="language_model.model.layers.{}.post_feedforward_layernorm",
+    embedding="language_model.model.embed_tokens",
+    final_norm="language_model.model.norm",
+    lm_head=None,
+)
+
+
+class Decoder2(gemma2.Gemma2):
+  """A decoder of PaliGemma2 3B model which is Gemma2.
+
+  Besides a tensor of text token IDs, forward() can also take a tensor of
+  embeddings which may include text or image or both.
+  """
+
+  @torch.inference_mode
+  def forward(
+      self,
+      tokens: torch.Tensor,
+      input_pos: torch.Tensor,
+      kv_cache: kv_utils.KVCache,
+      input_embeds: torch.Tensor = None,
+      export_config: Optional[model_builder.ExportConfig] = None,
+      called_by_generate: bool = True,
+  ) -> dict[torch.Tensor, kv_utils.KVCache]:
+    if input_embeds is None:
+      return super().forward(tokens, input_pos, kv_cache)
+
+    assert input_embeds is not None
+
+    repo_pos = input_pos + 1  # PaliGemma2 position is 1-based.
+    # ROPE parameters for all attn_configs are the same. Take the first one.
+    attn_config = self.config.block_config(0).attn_config
+    n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
+    rope = rotary_pos_emb.build_rope(
+        repo_pos, n_elem, attn_config.head_dim, attn_config.rotary_base
+    )
+
+    if called_by_generate:
+      # PaliGemma2 generate() use a diagonal causal mask even with image embeds.
+      mask = [self.get_attention_mask(
+          self.config.block_config(i).attn_config.attn_type, input_pos
+      ) for i in range(self.config.num_layers)]
+    else:
+      # By default, don't mask image embeds with a diagonal causal mask.
+      embeds_len = input_embeds.shape[1]
+      mask = torch.zeros(embeds_len, self.config.kv_cache_max)
+      mask[:, embeds_len:] = float("-inf")
+      mask = [mask] * self.config.num_layers
+
+    return self._forward_with_embeds(
+        input_embeds, rope, mask, input_pos, kv_cache, export_config
+    )
+
+
+def get_decoder2_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
+  """Returns the model config for the decoder of a PaliGemma 3B model.
+
+  Args:
+    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
+      is 1024.
+
+  Returns:
+    The model config for the decoder of a PaliGemma 3B model.
+  """
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM,
+      epsilon=1e-6,
+      zero_centered=True,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_TANH),
+      intermediate_size=9216,
+      pre_ff_norm_config=norm_config,
+      post_ff_norm_config=norm_config,
+  )
+
+  def get_block_config(idx: int) -> cfg.TransformerBlockConfig:
+    attn_config = cfg.AttentionConfig(
+        num_heads=8,
+        head_dim=256,
+        num_query_groups=4,
+        rotary_base=10000,
+        rotary_percentage=1.0,
+        logit_softcap=50.0,
+        sliding_window_size=4096,
+        attn_type=(
+            cfg.AttentionType.GLOBAL
+            if idx % 2 == 0
+            else cfg.AttentionType.LOCAL_SLIDING
+        ),
+    )
+    return cfg.TransformerBlockConfig(
+        attn_config=attn_config,
+        ff_config=ff_config,
+        pre_attention_norm_config=norm_config,
+        post_attention_norm_config=norm_config,
+    )
+
+  num_layers = 26
+  embedding_dim = 2304
+  config = cfg.ModelConfig(
+      vocab_size=257216,
+      num_layers=num_layers,
+      max_seq_len=8192,
+      embedding_dim=embedding_dim,
+      embedding_scale=embedding_dim**0.5,
+      kv_cache_max_len=kv_cache_max_len,
+      block_configs=[get_block_config(i) for i in range(num_layers)],
+      final_norm_config=norm_config,
+      lm_head_use_bias=False,
+      enable_hlfb=True,
+      final_logit_softcap=30.0,
+  )
+  return config
+
+
+def get_fake_decoder2_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
+  config = get_decoder2_config(kv_cache_max_len)
+  # PaliGemma2 decoder has only one block config.
+  config.block_config(0).ff_config.intermediate_size = 128
+  config.vocab_size = 128
+  config.num_layers = 2
+  config.max_seq_len = 2 * kv_cache_max_len
+  config.embedding_dim = 128
+  config.embedding_scale = 128**0.5
+  return config
+
+
+def build_decoder2(checkpoint_path: str, **kwargs) -> torch.nn.Module:
+  return model_builder.build_decoder_only_model(
+      checkpoint_path=checkpoint_path,
+      config=get_decoder2_config(**kwargs),
+      tensor_names=TENSOR_NAMES,
+      model_class=Decoder2,
+  )

ai_edge_torch/generative/examples/paligemma/paligemma.py

@@ -19,6 +19,7 @@ from dataclasses import dataclass
 from typing import Optional
 
 from ai_edge_torch.generative.examples.paligemma import decoder
+from ai_edge_torch.generative.examples.paligemma import decoder2
 from ai_edge_torch.generative.examples.paligemma import image_encoder
 import ai_edge_torch.generative.layers.kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
@@ -38,13 +39,14 @@ class PaliGemmaConfig:
   decoder_config: cfg.ModelConfig
 
   image_token_id: int
+  image_projection_scale: float
   image_projection_use_bias: bool = False
 
 
 class PaliGemma(nn.Module):
   """PaliGemma model from the Edge Generative API."""
 
-  def __init__(self, config: PaliGemmaConfig):
+  def __init__(self, config: PaliGemmaConfig, decoder_class: nn.Module):
     super().__init__()
 
     self.image_encoder = image_encoder.SiglipVisionEncoder(
@@ -55,7 +57,7 @@ class PaliGemma(nn.Module):
         config.decoder_config.embedding_dim,
         bias=config.image_projection_use_bias,
     )
-    self.decoder = decoder.Decoder(config.decoder_config)
+    self.decoder = decoder_class(config.decoder_config)
     image_embedding_config = config.image_encoder_config.image_embedding
     self.num_patches = (
         image_embedding_config.image_size // image_embedding_config.patch_size
@@ -70,6 +72,7 @@ class PaliGemma(nn.Module):
       kv_cache: kv_utils.KVCache,
       pixel_values: torch.Tensor = None,
       export_config: Optional[model_builder.ExportConfig] = None,
+      called_by_generate: bool = True,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if pixel_values is None:
       return self.decoder(
@@ -77,15 +80,15 @@ class PaliGemma(nn.Module):
           input_pos=input_pos,
           kv_cache=kv_cache,
           input_embeds=None,
-          export_config=export_config
+          export_config=export_config,
+          called_by_generate=called_by_generate,
       )
 
     input_embeds = self.decoder.tok_embedding(tokens)
 
     image_encoded = self.image_encoder(pixel_values=pixel_values)
     image_embeds = self.image_projection(image_encoded)
-    if self.config.decoder_config.embedding_scale is not None:
-      image_embeds = image_embeds / self.config.decoder_config.embedding_scale
+    image_embeds = image_embeds / self.config.image_projection_scale
 
     # Merging image_embeds into text_embeds as PaliGemmaForConditionalGeneration
     # can be done like:
@@ -110,10 +113,11 @@ class PaliGemma(nn.Module):
         kv_cache=kv_cache,
         input_embeds=input_embeds,
         export_config=export_config,
+        called_by_generate=called_by_generate,
     )
 
 
-def get_model_config(**kwargs) -> PaliGemmaConfig:
+def get_model_config(get_decoder_config, **kwargs) -> PaliGemmaConfig:
   """Returns the model config for a PaliGemma 3B-224 model.
 
   Returns:
@@ -121,31 +125,42 @@ def get_model_config(**kwargs) -> PaliGemmaConfig:
   """
   return PaliGemmaConfig(
       image_encoder_config=image_encoder.get_image_encoder_config(),
-      decoder_config=decoder.get_decoder_config(**kwargs),
-      image_projection_use_bias=True,
+      decoder_config=get_decoder_config(**kwargs),
       image_token_id=257152,
+      image_projection_scale=2048**0.5,
+      image_projection_use_bias=True,
   )
 
 
-def get_fake_model_config() -> PaliGemmaConfig:
+def get_fake_model_config(get_decoder_config, **kwargs) -> PaliGemmaConfig:
   return PaliGemmaConfig(
       image_encoder_config=image_encoder.get_fake_image_encoder_config(),
-      decoder_config=decoder.get_fake_decoder_config(),
+      decoder_config=get_decoder_config(**kwargs),
+      image_token_id=127,
+      image_projection_scale=128**0.5,
       image_projection_use_bias=True,
-      image_token_id=257152,
   )
 
 
-def build_model(checkpoint_path: str, **kwargs) -> PaliGemma:
-  config = get_model_config(**kwargs)
-  model = PaliGemma(config)
+def build_model(checkpoint_path: str, version: int = 2, **kwargs) -> PaliGemma:
+  if version == 1:
+    decoder_class = decoder.Decoder
+    decoder_tensor_names = decoder.TENSOR_NAMES
+    get_decoder_config = decoder.get_decoder_config
+  else:
+    decoder_class = decoder2.Decoder2
+    decoder_tensor_names = decoder2.TENSOR_NAMES
+    get_decoder_config = decoder2.get_decoder2_config
+
+  config = get_model_config(get_decoder_config, **kwargs)
+  model = PaliGemma(config, decoder_class)
   # Load the parameters of image encoder.
   loader = loading_utils.ModelLoader(
       checkpoint_path, image_encoder.TENSOR_NAMES
   )
   loader.load(model.image_encoder, strict=False)
   # Load the parameters of decoder.
-  loader = loading_utils.ModelLoader(checkpoint_path, decoder.TENSOR_NAMES)
+  loader = loading_utils.ModelLoader(checkpoint_path, decoder_tensor_names)
   loader.load(model.decoder, strict=False)
 
   # Load the parameters of image projection.

ai_edge_torch/generative/examples/paligemma/verify.py

@@ -22,11 +22,18 @@ from absl import flags
 from ai_edge_torch.generative.examples.paligemma import paligemma
 from ai_edge_torch.generative.layers import kv_cache
 from ai_edge_torch.generative.utilities import verifier
+import kagglehub
 from PIL import Image
 import requests
 import torch
 import transformers
 
+_VERSION = flags.DEFINE_enum(
+    "version",
+    "2",
+    ["1", "2"],
+    "The version of PaliGemma model to verify.",
+)
 _IMAGE_URL = flags.DEFINE_string(
     "image_url",
     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true",
@@ -34,7 +41,7 @@ _IMAGE_URL = flags.DEFINE_string(
 )
 _PROMPTS = flags.DEFINE_string(
     "prompts",
-    "Caption en",
+    "describe en",
     "The input prompts to generate answers.",
 )
 _MAX_NEW_TOKENS = flags.DEFINE_integer(
@@ -43,28 +50,47 @@ _MAX_NEW_TOKENS = flags.DEFINE_integer(
     "The maximum size of the generated tokens.",
 )
 
+_CHECKPOINT = {
+    "1": "google/paligemma-3b-mix-224",
+    "2": "google/paligemma-2/transformers/paligemma2-3b-pt-224",
+}
+
 
 class ReauthoredPaliGemmaWrapper(verifier.ReauthoredModelWrapper):
   """Reauthored PaliGemma model wrapper."""
 
+  def __init__(self, model: torch.nn.Module):
+    super().__init__(model)
+    self.forward_called_by_generate = False
+
   def _init_kv_cache(self):
     return kv_cache.KVCache.from_model_config(self.model.config.decoder_config)
 
+  def _get_extra_args_for_forward(self):
+    return {"called_by_generate": self.forward_called_by_generate}
+
 
 def main(_):
-  checkpoint = "google/paligemma-3b-mix-224"
+  if _VERSION.value == "1":
+    checkpoint = _CHECKPOINT[_VERSION.value]
+    # Locate the cached dir.
+    cached_config_file = transformers.utils.cached_file(
+        checkpoint, transformers.utils.CONFIG_NAME
+    )
+    reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  else:
+    checkpoint = kagglehub.model_download(_CHECKPOINT[_VERSION.value])
+    reauthored_checkpoint = checkpoint
+
   logging.info("Loading the original model from: %s", checkpoint)
   original_model = (
       transformers.PaliGemmaForConditionalGeneration.from_pretrained(checkpoint)
   )
 
-  # Locate the cached dir.
-  cached_config_file = transformers.utils.cached_file(
-      checkpoint, transformers.utils.CONFIG_NAME
-  )
-  reauthored_checkpoint = pathlib.Path(cached_config_file).parent
   logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
-  reauthored_model = paligemma.build_model(reauthored_checkpoint)
+  reauthored_model = paligemma.build_model(
+      reauthored_checkpoint, version=int(_VERSION.value)
+  )
 
   logging.info("Loading the processor from: %s", checkpoint)
   # It works only when GemmaTokenizerFast is available. In some environments,
@@ -93,7 +119,7 @@ def main(_):
   logging.info("outputs_reauthored: %s", outputs_reauthored)
 
   try:
-    assert torch.allclose(outputs_original, outputs_reauthored, atol=1e-03)
+    assert torch.allclose(outputs_original, outputs_reauthored, atol=1e-02)
   except AssertionError as e:
     logging.error("*** FAILED *** verify with forward()")
     raise e
@@ -111,6 +137,7 @@ def main(_):
   logging.info("outputs_from_original_model: [[%s]]", response_original)
 
   logging.info("Generating answer with the reauthored model...")
+  wrapped_reauthored_model.forward_called_by_generate = True
   outputs_reauthored = wrapped_reauthored_model.generate(
       prompts=inputs["input_ids"],
       pixel_values=inputs["pixel_values"],

ai_edge_torch/generative/examples/paligemma/verify_decoder2.py

@@ -0,0 +1,72 @@
+# 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 decoder of PaliGemma2 3B model."""
+
+import logging
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.paligemma import decoder2
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import kagglehub
+import transformers
+
+_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/paligemma-2/transformers/paligemma2-3b-pt-224"
+  )
+  logging.info("Loading the original model from: %s", checkpoint)
+  original_full_model = (
+      transformers.PaliGemmaForConditionalGeneration.from_pretrained(checkpoint)
+  )
+  original_language_model = original_full_model.eval().language_model
+
+  logging.info("Building the reauthored model from: %s", checkpoint)
+  reauthored_model = decoder2.build_decoder2(checkpoint)
+
+  logging.info("Loading the tokenizer from: %s", checkpoint)
+  # It works only when GemmaTokenizerFast is available. In some environments,
+  # use_fast=False doeesn't work either if the tokenizer cannot load the
+  # sentencepiece model file properly.
+  processor = transformers.AutoProcessor.from_pretrained(checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_language_model
+      ),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=verifier.TokenizerWrapper(processor.tokenizer),
+      generate_prompts=_PROMPTS.value,
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/paligemma/verify_image_encoder.py

@@ -20,31 +20,48 @@ import pathlib
 from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.paligemma import image_encoder
+import kagglehub
 from PIL import Image
 import requests
 import torch
 import transformers
 
+_VERSION = flags.DEFINE_enum(
+    "version",
+    "2",
+    ["1", "2"],
+    "The version of PaliGemma vision model to verify.",
+)
 _IMAGE_URL = flags.DEFINE_string(
     "image_url",
     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true",
     "The image URI to encode.",
 )
 
+_CHECKPOINT = {
+    "1": "google/paligemma-3b-mix-224",
+    "2": "google/paligemma-2/transformers/paligemma2-3b-pt-224",
+}
+
 
 def main(_):
-  checkpoint = "google/paligemma-3b-mix-224"
+  if _VERSION.value == "1":
+    checkpoint = _CHECKPOINT[_VERSION.value]
+    # Locate the cached dir.
+    cached_config_file = transformers.utils.cached_file(
+        checkpoint, transformers.utils.CONFIG_NAME
+    )
+    reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  else:
+    checkpoint = kagglehub.model_download(_CHECKPOINT[_VERSION.value])
+    reauthored_checkpoint = checkpoint
+
   logging.info("Loading the original model from: %s", checkpoint)
   original_full_model = (
       transformers.PaliGemmaForConditionalGeneration.from_pretrained(checkpoint)
   )
   original_vision_model = original_full_model.eval().vision_tower
 
-  # Locate the cached dir.
-  cached_config_file = transformers.utils.cached_file(
-      checkpoint, transformers.utils.CONFIG_NAME
-  )
-  reauthored_checkpoint = pathlib.Path(cached_config_file).parent
   logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
   reauthored_model = image_encoder.build_image_encoder(reauthored_checkpoint)
 
@@ -69,7 +86,7 @@ def main(_):
 
   try:
     assert torch.allclose(
-        outputs_original, outputs_reauthored, atol=1e-04, rtol=1e-04
+        outputs_original, outputs_reauthored, atol=1e-03, rtol=1e-04
     )
   except AssertionError as e:
     logging.error("*** FAILED *** verify with an image")

ai_edge_torch/generative/test/test_model_conversion_large.py

@@ -21,6 +21,8 @@ from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.examples.gemma import gemma2
 from ai_edge_torch.generative.examples.llama import llama
 from ai_edge_torch.generative.examples.openelm import openelm
+from ai_edge_torch.generative.examples.paligemma import decoder
+from ai_edge_torch.generative.examples.paligemma import decoder2
 from ai_edge_torch.generative.examples.paligemma import paligemma
 from ai_edge_torch.generative.examples.phi import phi2
 from ai_edge_torch.generative.examples.phi import phi3
@@ -171,13 +173,9 @@ class TestModelConversion(googletest.TestCase):
     pytorch_model = amd_llama_135m.AmdLlama(config).eval()
     self._test_model(config, pytorch_model, "prefill", atol=1e-5, rtol=1e-5)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
-  def disabled_test_paligemma(self):
-    config = paligemma.get_fake_model_config()
-    pytorch_model = paligemma.PaliGemma(config).eval()
+  def _test_paligemma_model(self, decoder_class, decoder_config, atol, rtol):
+    config = paligemma.get_fake_model_config(decoder_config)
+    pytorch_model = paligemma.PaliGemma(config, decoder_class).eval()
 
     image_embedding_config = config.image_encoder_config.image_embedding
     num_patches = (
@@ -215,11 +213,32 @@ class TestModelConversion(googletest.TestCase):
             kv,
             pixel_values=pixel_values,
             signature_name="prefill_pixel",
-            atol=1e-3,
-            rtol=1e-5,
+            atol=atol,
+            rtol=rtol,
         )
     )
 
+  @googletest.skipIf(
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
+  )
+  def disabled_test_paligemma1(self):
+    self._test_paligemma_model(
+        decoder.Decoder, decoder.get_fake_decoder_config, atol=1e-3, rtol=1e-5
+    )
+
+  @googletest.skipIf(
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
+  )
+  def disabled_test_paligemma2(self):
+    self._test_paligemma_model(
+        decoder2.Decoder2,
+        decoder2.get_fake_decoder2_config,
+        atol=1e-3,
+        rtol=1e-5,
+    )
+
   @googletest.skipIf(
       ai_edge_torch.config.in_oss,
       reason="tests with custom ops are not supported in oss",

ai_edge_torch/generative/utilities/model_builder.py

@@ -107,8 +107,6 @@ class DecoderOnlyModel(nn.Module):
 
     # token embeddings of shape (b, t, n_embd)
     input_embeds = self.tok_embedding(tokens)
-    mask = self.mask_cache.index_select(2, input_pos)
-    mask = mask[:, :, :, : self.config.kv_cache_max]
 
     # ROPE parameters for all attn_configs are the same. Take the first one.
     attn_config = self.config.block_config(0).attn_config
@@ -117,11 +115,14 @@ class DecoderOnlyModel(nn.Module):
         input_pos, n_elem, attn_config.head_dim, attn_config.rotary_base
     )
 
-    return self.forward_with_embeds(
+    mask = self.mask_cache.index_select(2, input_pos)
+    mask = mask[:, :, :, : self.config.kv_cache_max]
+
+    return self._forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache, export_config
     )
 
-  def forward_with_embeds(
+  def _forward_with_embeds(
       self,
       input_embeds: torch.Tensor,
       rope: Tuple[torch.Tensor, torch.Tensor],

ai_edge_torch/generative/utilities/verifier.py

@@ -16,7 +16,7 @@
 """Common utility functions to verify the reauthored models."""
 
 import logging
-from typing import List
+from typing import Any,List
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.generative.utilities.model_builder import ExportConfig
@@ -87,6 +87,10 @@ class ReauthoredModelWrapper(ModelWrapper):
     """Returns an initialized KV cache."""
     return kv_utils.KVCache.from_model_config(self.model.config)
 
+  def _get_extra_args_for_forward(self) -> dict[str, Any]:
+    """Returns extra arguments for the forward() method."""
+    return {}
+
   def _forward_with_kv_cache(
       self,
       tokens: torch.Tensor,
@@ -105,26 +109,15 @@ class ReauthoredModelWrapper(ModelWrapper):
     Returns:
       The output logits and the updated KV cache.
     """
-    # Verification requires logit outputs on prefill for comparison.
-    if (
-        self.export_config is not None
-        and not self.export_config.output_logits_on_prefill
-    ):
-      raise ValueError("Verifier requires logit output on prefill.")
-    # Since the reauthored model doesn't include keyword arguments, pass
-    # pixel_values only when it is not None. Otherwise, it may raise an error.
-    if pixel_values is None:
-      output = self.model.forward(
-          tokens, input_pos, kv_cache, export_config=self.export_config
-      )
-    else:
-      output = self.model.forward(
-          tokens,
-          input_pos,
-          kv_cache,
-          pixel_values=pixel_values,
-          export_config=self.export_config,
-      )
+    extra_args = self._get_extra_args_for_forward()
+    if self.export_config is not None:
+      # Verification requires logit outputs on prefill for comparison.
+      if not self.export_config.output_logits_on_prefill:
+        raise ValueError("Verifier requires logit output on prefill.")
+      extra_args["export_config"] = self.export_config
+    if pixel_values is not None:
+      extra_args["pixel_values"] = pixel_values
+    output = self.model.forward(tokens, input_pos, kv_cache, **extra_args)
     return output["logits"], output["kv_cache"]
 
   def forward(
@@ -141,6 +134,7 @@ class ReauthoredModelWrapper(ModelWrapper):
       prompts: torch.Tensor,
       max_new_tokens: int,
       pixel_values: torch.Tensor = None,
+      eos_token_id: int = 1,
   ) -> torch.IntTensor:
     input_ids = prompts[0].int().tolist()
     tokens = torch.tensor([input_ids])
@@ -152,6 +146,8 @@ class ReauthoredModelWrapper(ModelWrapper):
       )
       generated_token = logits[0][-1].argmax().item()
       input_ids.append(generated_token)
+      if generated_token == eos_token_id:
+        break
       tokens = torch.tensor([[generated_token]])
       input_pos = torch.tensor([len(input_ids) - 1])
       pixel_values = None  # Pass only for the first time.
@@ -254,7 +250,11 @@ def verify_model_with_prompts(
   logging.info("outputs_from_original_model: [[%s]]", response_original)
 
   logging.info("Generating answer with the reauthored model...")
-  outputs_reauthored = reauthored_model.generate(prompt_tokens, max_new_tokens)
+  outputs_reauthored = reauthored_model.generate(
+      prompt_tokens,
+      max_new_tokens,
+      eos_token_id=tokenizer.tokenizer.eos_token_id,
+  )
   response_reauthored = tokenizer.decode(outputs_reauthored[0])
   logging.info("outputs from reauthored model: [[%s]]", response_reauthored)
 

ai_edge_torch/odml_torch/export.py

@@ -198,7 +198,12 @@ class MlirLowered:
     # build, which may not have the same StableHLO version as what used in
     # TFLite converter. Therefore we always serialize MLIR module in VHLO.
     # TODO(b/362798610) Build MLIR pybinding in ai-edge-torch release.
-    target_version = stablehlo.get_minimum_version()
+    if stablehlo.get_api_version() < 9:
+      target_version = stablehlo.get_minimum_version()
+    else:
+      target_version = stablehlo.get_version_from_compatibility_requirement(
+          stablehlo.StablehloCompatibilityRequirement.WEEK_4
+      )
     module_bytecode = xla_extension.mlir.serialize_portable_artifact(
         self.module_bytecode, target_version
     )

ai_edge_torch/odml_torch/jax_bridge/__init__.py

@@ -12,4 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from ai_edge_torch.odml_torch.jax_bridge._wrap import wrap
+from ai_edge_torch.odml_torch.jax_bridge import _wrap
+from ai_edge_torch.odml_torch.jax_bridge import utils
+
+wrap = _wrap.wrap

ai_edge_torch/odml_torch/lowerings/__init__.py

@@ -18,6 +18,7 @@ from . import _convolution
 from . import _jax_lowerings
 from . import _layer_norm
 from . import _quantized_decomposed
+from . import _rand
 from . import context
 from . import registry
 from . import utils

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -26,6 +26,7 @@ import torch_xla2.ops.ops_registry  # Import to load torch_xla2 ops
 
 LoweringContext = context.LoweringContext
 
+
 @functools.cache
 def _log_usage(op):
   logging.warning("Use jax lowering: %s", str(op))
@@ -184,8 +185,6 @@ lower_by_torch_xla2(torch.ops.aten.permute_copy)
 lower_by_torch_xla2(torch.ops.aten.pixel_shuffle)
 lower_by_torch_xla2(torch.ops.aten.pow)
 lower_by_torch_xla2(torch.ops.aten.prod)
-lower_by_torch_xla2(torch.ops.aten.rand)
-lower_by_torch_xla2(torch.ops.aten.randn)
 lower_by_torch_xla2(torch.ops.aten.reciprocal)
 lower_by_torch_xla2(torch.ops.aten.reflection_pad1d)
 lower_by_torch_xla2(torch.ops.aten.relu)

ai_edge_torch/odml_torch/lowerings/_rand.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.
+# ==============================================================================
+import uuid
+
+from ai_edge_torch.odml_torch import export_utils
+from ai_edge_torch.odml_torch.lowerings import context
+from ai_edge_torch.odml_torch.lowerings import registry
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import func
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import numpy as np
+import torch
+import torch.utils._pytree as pytree
+
+LoweringContext = context.LoweringContext
+lower = registry.lower
+
+
+def _random_lowering(
+    lctx: LoweringContext,
+    size: list[int],
+    generator,
+    dtype: torch.dtype,
+    rand_tensor,
+    composite_name: str,
+):
+  if dtype is None:
+    dtype = torch.float32
+
+  rand_tensor = rand_tensor.type(dtype)
+  data = rand_tensor.detach().numpy()
+
+  shape, _ = pytree.tree_flatten(size)
+  elty = export_utils.torch_dtype_to_ir_element_type(dtype)
+
+  decomp_name = f"{composite_name}.impl_{uuid.uuid4().hex[:8]}"
+
+  with ir.InsertionPoint(lctx.ir_module.body):
+
+    @func.FuncOp.from_py_func(
+        ir.RankedTensorType.get(
+            [len(shape)],
+            ir.IntegerType.get_signless(32),
+        ),
+        name=decomp_name,
+    )
+    def _rand_impl(_):
+      return [stablehlo.constant(ir.DenseElementsAttr.get(data))]
+
+  seed, seed2 = (
+      torch.randint(
+          torch.iinfo(torch.int64).min,
+          torch.iinfo(torch.int64).max,
+          (2,),
+          dtype=torch.int64,
+          generator=generator,
+      )
+      .detach()
+      .numpy()
+  )
+
+  shape_ = stablehlo.constant(
+      ir.DenseElementsAttr.get(np.array(shape, dtype=np.int32))
+  )
+  return stablehlo.CompositeOp(
+      result=[ir.RankedTensorType.get(shape, elty)],
+      inputs=[shape_],
+      name=composite_name,
+      composite_attributes=ir.DictAttr.get({
+          "seed": ir.IntegerAttr.get(ir.IntegerType.get_signless(64), seed),
+          "seed2": ir.IntegerAttr.get(ir.IntegerType.get_signless(64), seed2),
+      }),
+      decomposition=decomp_name,
+  ).results[0]
+
+
+# Schema:
+# - aten::rand(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None,
+#     Device? device=None, bool? pin_memory=None) -> Tensor
+# - aten::rand.generator(SymInt[] size, *, Generator? generator,
+#     ScalarType? dtype=None, Layout? layout=None, Device? device=None,
+#     bool? pin_memory=None) -> Tensor
+@registry.lower(torch.ops.aten.rand)
+def _aten_rand(
+    lctx: LoweringContext,
+    size,
+    generator=None,
+    dtype=None,
+    layout=torch.strided,
+    device=None,
+    pin_memory=False,
+):
+  return _random_lowering(
+      lctx,
+      size,
+      generator,
+      dtype,
+      rand_tensor=torch.ops.aten.rand.generator(
+          size, generator=generator, dtype=dtype
+      ),
+      composite_name="odml.random_uniform",
+  )
+
+
+# Schema:
+# - aten::randn(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None,
+#     Device? device=None, bool? pin_memory=None) -> Tensor
+# - aten::randn.generator(SymInt[] size, *, Generator? generator,
+#     ScalarType? dtype=None, Layout? layout=None, Device? device=None,
+#     bool? pin_memory=None) -> Tensor
+@registry.lower(torch.ops.aten.randn)
+def _aten_randn(
+    lctx: LoweringContext,
+    size,
+    generator=None,
+    dtype=None,
+    layout=torch.strided,
+    device=None,
+    pin_memory=False,
+):
+  return _random_lowering(
+      lctx,
+      size,
+      generator,
+      dtype,
+      rand_tensor=torch.ops.aten.randn.generator(
+          size, generator=generator, dtype=dtype
+      ),
+      composite_name="odml.random_standard_normal",
+  )

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241220.dist-info → ai_edge_torch_nightly-0.3.0.dev20241224.dist-info}/METADATA

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

@@ -3,7 +3,7 @@ ai_edge_torch/_config.py,sha256=QIrerb6uHMahRvMilmhodJ_6jfiRps3qgLOBeidPnS4,1614
 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=xD-MWAEa1ROHhyF3rY7MaL28xsuON0aJwaiXbJ04qfc,706
+ai_edge_torch/version.py,sha256=TkfJYt2lJC8A_AcieO1xVmMQ2xdnoTOwF8CZ5dZeaqc,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=SzbR16V2JEfkCjjPwRVAFUbFnzu-_1iHPKgGT9Yz7gQ,5678
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -46,8 +46,8 @@ ai_edge_torch/generative/examples/amd_llama_135m/verify.py,sha256=-9Nb9D818YSJR3
 ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=GhwtQZ1xuMyKJl8qdxU6uKavQnlm5US9xhKJvdmgACc,2309
 ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=hsy4Gd7Inchi0p_Cc5yecH6vr9A7X4MvmQNfTt8N2sQ,2311
-ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=N0jKVZA3qWKOaHVbIM3WmQh3u0Sq7MTw_oO3Zo16wCw,3456
-ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=roEwWVXASbk5BFj7jojjEJpHui6gCelT51l-TtN_ZaQ,9367
+ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=w8oWYibZzvEvCDyp39EYyAWmjgJljhzdYPyFCfAWxZA,3497
+ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=VTM2nO3TqK2d1DyEb2MiHc-Tyw2lMcUXyOhvg0H5ENY,10147
 ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=ip-Gmk4CI5f0GWSdAIdrectxQWJ0t328KCsA4nfHuGg,1736
 ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=IoBhEMwH07-tFm5-U6F2hpCsI8xynglhq1x9tIOdaPQ,1322
 ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=tR8RflXocDZqvuStyw9aFlzuiTllEC8rNnjrxms6_Is,5727
@@ -63,13 +63,15 @@ ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=nji1oDgf6x
 ai_edge_torch/generative/examples/openelm/openelm.py,sha256=sIJ8Ie1oxFrJM-1jvv2ukiJbQOTIUGuMEZvmwZbt3n0,4556
 ai_edge_torch/generative/examples/openelm/verify.py,sha256=VkigoqhAr8ew95neb3TifYv-SLOSheaWKv2AH0iKDrc,2441
 ai_edge_torch/generative/examples/paligemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=rPFqcsv8RHvjmgfBW9OL6EKxMtVX-ySjBsMP4N8FErk,2816
-ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=eICKQkJsJuEUkuvn5ymUsI9CGB-oNbgV7VH7BlmklfQ,4961
+ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=LFCcnkmOksySDa_5bLBzoGMijYdFVjXIMidUlyzAbNk,2996
+ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=amN96oBMTPolOFvGa47vG92AZ-BNLm8j0bBYd-IrMvI,5407
+ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=0V_CX0Pn5Fj_-koOGjc_Av2KMSAaVjAlD-G8P6FBGyY,6385
 ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=yKPWG8aBp-GuzeyQntlzwTTcGBBjvUywVGRjnlNprmo,5574
-ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=nDyI-wUFJSawu57uLbFENei5l4cciqZ8lM5S5beN0FU,5604
-ai_edge_torch/generative/examples/paligemma/verify.py,sha256=Bkbgy-GFjnMNYjduWUM7YLWarPTwmj1v38eHY-PdBlM,4874
+ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=FwGlFHl9zktGDxnoOpEtbS6NYN5RyzcOXH7lvNUCwEU,6257
+ai_edge_torch/generative/examples/paligemma/verify.py,sha256=KT3Ruy40tSESxQuy-Sw01NAI3zId1BZr6Bp7FZj1wZk,5622
 ai_edge_torch/generative/examples/paligemma/verify_decoder.py,sha256=al5wMPWri4IRVWrLmCplPi6uoCzwh0vBHMGnCt-XUqo,2690
-ai_edge_torch/generative/examples/paligemma/verify_image_encoder.py,sha256=pSekf1BybhieQz3cQx_llbRQHxczXbTqool8fOyGj_0,3114
+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=cD8rtwgYeGrXB9sYVV_D1AB8Up1AWNS-1XtrRlyzE5o,2296
 ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=G1i_ybDCTBaOD1OOCTk6jqOf__xYYZvhXcxY8MXhPHw,2294
@@ -140,19 +142,19 @@ ai_edge_torch/generative/test/test_custom_dus.py,sha256=gxG78CcTpXF3iLzDR15Rlz1e
 ai_edge_torch/generative/test/test_kv_cache.py,sha256=2AulHBS3hC4b_68PNNBkRVOrypy4IM5YjC4p-6dgCMM,3793
 ai_edge_torch/generative/test/test_loader.py,sha256=9mQUeeZKOVApOWSWl2cN9c10axZjMKM1-0Zd823CCS4,3449
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=jfqkECCX7XKHeBAuDXrkwQJf0vM72eG3LMc5rluha84,6191
-ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=mVuax3MPRmuNjnDRKXqtc9YmswCy7MnhD1CHADK-3nk,11501
+ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=NctnggTSFh0XEQbTu55diZ35rFD2QIARO-8PzLktRWg,12165
 ai_edge_torch/generative/test/test_quantize.py,sha256=bEJMhpQ9bIDUZVBXTW888728FcH-i3SyE4JSZZUgU0A,6071
 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=hIwWUWjgPvWLATtsYYG6RWbFQWhOr2RpPlMrd-4Am9U,5959
 ai_edge_torch/generative/utilities/dynamic_update_slice.py,sha256=e2mhx-Vp8sUK4EXoPtpZLSx3TViqLAKs67EhKcXBjAQ,2121
 ai_edge_torch/generative/utilities/loader.py,sha256=A3SOjPXp--AsvoP1hqj5QKWE4sgxoFc3H5EBUz_Eogc,13531
-ai_edge_torch/generative/utilities/model_builder.py,sha256=q82-1E2zYlzpbFW6Vw-MWrJivRXHKpRh8jUxpR-w0sY,6349
+ai_edge_torch/generative/utilities/model_builder.py,sha256=S08WNqVKCmxd2QjtMlwETd7J97UnlME_bTKdz5LMkGU,6352
 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
 ai_edge_torch/generative/utilities/t5_loader.py,sha256=tEsfy8-ymzbbjOIc-oesXF3yGyyWtJgFXn2s7VOavt8,16961
 ai_edge_torch/generative/utilities/transformers_verifier.py,sha256=8sp9m_FMcXn7nqOrochtu2jIANkJKhnhIBUmH0ZTDR4,1549
-ai_edge_torch/generative/utilities/verifier.py,sha256=ESSA8W1EYNsd4ntwmXbr-dn-BcIS27hf53XL5RTwjEU,11941
+ai_edge_torch/generative/utilities/verifier.py,sha256=awO-sQrEpsFxIkZw72ysWZenYEmkLOLOuj62o2c7XeQ,11994
 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
@@ -169,7 +171,7 @@ ai_edge_torch/lowertools/translate_recipe.py,sha256=ymkBpFqAUiupRWqrPOWiVphKcXR1
 ai_edge_torch/odml_torch/__init__.py,sha256=S8jOzE9nLof-6es3XDiGJRN-9H_XTxsVm9dE7lD3RWo,812
 ai_edge_torch/odml_torch/_torch_future.py,sha256=AJ0klpsbu2ZBTfiZlqSOoaYzBVITt40a1fYN8xKkEPw,3044
 ai_edge_torch/odml_torch/_torch_library.py,sha256=Lw1gqL2HWNRspdTwNhIkYAHDyafHedHtkXyKKxn-Wss,805
-ai_edge_torch/odml_torch/export.py,sha256=Wc_JM7U2IjZeBmXA6t1AZxREGOWjZ6EB-PIhEevWWeU,13207
+ai_edge_torch/odml_torch/export.py,sha256=QzOPmcNPB7R-KhhPEP0oGVbDRgGPptIxRSoz3S8py9I,13405
 ai_edge_torch/odml_torch/export_utils.py,sha256=QeA37Irlty6AiIBuqmHmJgn3lqahBQ5xsh6IKRoKm1g,4774
 ai_edge_torch/odml_torch/tf_integration.py,sha256=lTFJPPEijLPFmn6qq2jbpVTQOo0YaOTK36kK6rCiyIE,5956
 ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xKA9rknO93JS9kw9w_c,778
@@ -178,16 +180,17 @@ ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py,sha256=2Y52E_g
 ai_edge_torch/odml_torch/debuginfo/__init__.py,sha256=9ag6-WWRG50rPCtIV7OpIokEKu2YRyGlMZZqVPWUH6g,762
 ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=1xCXOs3-9UcsOyLFH0uyQwLu7c06iYFTo0NQ7Ckbl2I,1465
 ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py,sha256=IvOBQyROI9WHS3umHRxsDW-1YElU9BPWzKtJA2eKWOI,1739
-ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=Jco5zvejxuyl9xHQxZICAKbkgH7x38qPlwUUpD7S15Q,730
+ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNitEeg-IoBUGNfUxsDSA,798
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=LqwZ1vCJTSOzgzvH8LUAN-sAkF-l_pGj1AMEIzAqHCA,6638
 ai_edge_torch/odml_torch/jax_bridge/utils.py,sha256=T8isGc896VrHZ6c_L5pYmLpolQ7ibcOlgWfPuVFPzIg,2264
-ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=0GytV1dGnqe1mKityqQDNFNS8T4QBg3UZuRJcGHwGyA,993
+ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=GWFl7WWgExLXu6FEYxnig5_g6hd_Sfnl8690uFg2-CU,1013
 ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=8mZTp_ybcMO3tDRQdlDP68BVeTw560XsTR4XH-ldTdc,9987
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=Q0aDzyUcZMoSzSbOU-r3LJMgPe6fble0QwdYVIOHHHk,6887
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=OVmlPGwyhDXKhmG4SAeEsa6iLpJHEHV_jKqwfjYvetA,11643
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=CJHWkmY4aAVQ5dmFsVc3Ox9TPkoLSNOfa96psD4CLRo,11561
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=khJIvDVk2s332Nd2Be-5dM6-wp5DGff61HCV5lskHmQ,3011
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQDJnaAAxhNFAd7NfQm5SOVEp_nno_krA,6178
+ai_edge_torch/odml_torch/lowerings/_rand.py,sha256=g6SuqDkuC6hD35lyP1-5H7ASDIzPSmKukeNT5naZSv8,4133
 ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
 ai_edge_torch/odml_torch/lowerings/decomp.py,sha256=UoJeZVcr4zAN_11i-HzfOhxGCxUm-7b1JXPVBxR2hSs,2414
 ai_edge_torch/odml_torch/lowerings/registry.py,sha256=Tp2h11l5uTifO0aIkuUOWAF_ibEjmd65Xx99w3EXuGE,1924
@@ -200,8 +203,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.dev20241220.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241220.dist-info/METADATA,sha256=PfyYhqbf7VEibw2TEDRb8tBOIPG9dfXhT9tNNou_iZg,1966
-ai_edge_torch_nightly-0.3.0.dev20241220.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20241220.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241220.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241224.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241224.dist-info/METADATA,sha256=J8nAtbMNmSIcHuThVv0omkhpldggz91pIIYy-6ATJgM,1966
+ai_edge_torch_nightly-0.3.0.dev20241224.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20241224.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241224.dist-info/RECORD,,