ai-edge-torch-nightly 0.3.0.dev20241110__py3-none-any.whl → 0.3.0.dev20241115__py3-none-any.whl

ai_edge_torch/_convert/fx_passes/build_interpolate_composite_pass.py

@@ -51,6 +51,7 @@ def _get_upsample_bilinear2d_pattern():
     return {
         "output": (int(output_h), int(output_w)),
         "align_corners": False,
+        "is_nchw_op": True,
     }
 
   return pattern
@@ -74,6 +75,7 @@ def _get_upsample_bilinear2d_align_corners_pattern():
     return {
         "output": (int(output_h), int(output_w)),
         "align_corners": True,
+        "is_nchw_op": True,
     }
 
   return pattern

ai_edge_torch/generative/examples/gemma/verify_gemma2.py

@@ -15,10 +15,8 @@
 
 """Verifies the reauthored Gemma2 model."""
 
-import logging
 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
 import kagglehub
 
@@ -38,18 +36,7 @@ _MAX_NEW_TOKENS = flags.DEFINE_integer(
 def main(_):
   checkpoint = kagglehub.model_download("google/gemma-2/pyTorch/gemma-2-2b-it")
 
-  logging.info("Building the reauthored model from: %s", 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,
-  )
+  verify_util.verify_gemma2(checkpoint, _PROMPTS.value, _MAX_NEW_TOKENS.value)
 
 
 if __name__ == "__main__":

ai_edge_torch/generative/examples/gemma/verify_util.py

@@ -19,6 +19,7 @@ import logging
 import os
 from typing import List, Tuple
 
+from ai_edge_torch.generative.examples.gemma import gemma2
 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
@@ -109,8 +110,11 @@ def verify_reauthored_gemma_model(
     max_new_tokens: int = 20,
     rtol: float = 1e-05,
     atol: float = 1e-05,
-):
-  """Verifies the reauthored Gemma model against the original model."""
+) -> bool:
+  """Verifies the reauthored Gemma model against the original model.
+
+  Returns True if the verification passes, False otherwise.
+  """
   config = gemma_config.get_model_config(variant)
   config.tokenizer = os.path.join(checkpoint, tokenizer_filename)
   # Use float32 to be compatible with the reauthored model.
@@ -120,7 +124,7 @@ def verify_reauthored_gemma_model(
   original_model = gemma_model.GemmaForCausalLM(config).eval()
   original_model.load_weights(os.path.join(checkpoint, weight_filename))
 
-  verifier.verify_reauthored_model(
+  return verifier.verify_reauthored_model(
       original_model=GemmaWrapper(original_model),
       reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
       tokenizer=GemmaTokenizerWrapper(original_model.tokenizer),
@@ -130,3 +134,24 @@ def verify_reauthored_gemma_model(
       rtol=rtol,
       atol=atol,
   )
+
+
+def verify_gemma2(
+    gemma2_model_path: str, prompts: List[str], max_new_tokens: int
+) -> bool:
+  """Verifies the reauthored Gemma2 model.
+
+  Return True if the verification passes, False otherwise.
+  """
+  logging.info("Building the reauthored model from: %s", gemma2_model_path)
+  reauthored_model = gemma2.build_2b_model(gemma2_model_path)
+
+  return verify_reauthored_gemma_model(
+      checkpoint=gemma2_model_path,
+      variant="2b-v2",
+      reauthored_model=reauthored_model,
+      generate_prompts=prompts,
+      forward_input_ids=[[2, 651, 9456, 576, 573, 3520, 3858, 603, 235248]],
+      max_new_tokens=max_new_tokens,
+      atol=1e-04,
+  )

ai_edge_torch/generative/examples/paligemma/image_encoder.py

@@ -0,0 +1,158 @@
+# 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 an image encoder of PaliGemma model which is Siglip."""
+
+from ai_edge_torch.generative.layers import attention
+from ai_edge_torch.generative.layers import builder
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
+from torch import nn
+
+TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj="vision_tower.vision_model.encoder.layers.{}.mlp.fc1",
+    ff_down_proj="vision_tower.vision_model.encoder.layers.{}.mlp.fc2",
+    attn_query_proj=(
+        "vision_tower.vision_model.encoder.layers.{}.self_attn.q_proj"
+    ),
+    attn_key_proj=(
+        "vision_tower.vision_model.encoder.layers.{}.self_attn.k_proj"
+    ),
+    attn_value_proj=(
+        "vision_tower.vision_model.encoder.layers.{}.self_attn.v_proj"
+    ),
+    attn_output_proj=(
+        "vision_tower.vision_model.encoder.layers.{}.self_attn.out_proj"
+    ),
+    pre_attn_norm="vision_tower.vision_model.encoder.layers.{}.layer_norm1",
+    post_attn_norm="vision_tower.vision_model.encoder.layers.{}.layer_norm2",
+    embedding="vision_tower.vision_model.embeddings.patch_embedding",
+    embedding_position=(
+        "vision_tower.vision_model.embeddings.position_embedding.weight"
+    ),
+    final_norm="vision_tower.vision_model.post_layernorm",
+)
+
+
+class SiglipVisionEncoder(nn.Module):
+  """Signlip vision encoder from the Edge Generative API."""
+
+  def __init__(self, config: cfg.ModelConfig):
+    super().__init__()
+
+    # Construct model layers.
+    self.tok_embedding = nn.Conv2d(
+        in_channels=config.image_embedding.channels,
+        out_channels=config.embedding_dim,
+        kernel_size=config.image_embedding.patch_size,
+        stride=config.image_embedding.patch_size,
+        padding="valid",
+    )
+    num_patches = (
+        config.image_embedding.image_size // config.image_embedding.patch_size
+    ) ** 2
+    self.tok_embedding_position = nn.Parameter(
+        torch.zeros((num_patches, config.embedding_dim))
+    )
+
+    self.transformer_blocks = nn.ModuleList(
+        attention.TransformerBlock(config.block_config(idx), config)
+        for idx in range(config.num_layers)
+    )
+    self.final_norm = builder.build_norm(
+        config.embedding_dim,
+        config.final_norm_config,
+    )
+    self.config = config
+
+  @torch.inference_mode
+  def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+    # Embed the image according to SiplipVisionEmbeddings.
+    x = self.tok_embedding(pixel_values)
+    x = x.flatten(2).transpose(1, 2) + self.tok_embedding_position
+
+    # Pass a dummy mask because SDPA attention impl expects non-None mask.
+    mask = torch.zeros(x.shape[:2])
+    for _, block in enumerate(self.transformer_blocks):
+      x = block(x, mask=mask)
+    return self.final_norm(x)
+
+
+def get_image_encoder_config() -> cfg.ModelConfig:
+  """Returns the model config for the image encoder of a PaliGemma 3B-224 model.
+
+  Returns:
+    The model config for the image encoder of a PaliGemma 3B model.
+  """
+  image_embedding_config = cfg.ImageEmbeddingConfig(
+      channels=3,
+      image_size=224,
+      patch_size=14,
+  )
+  attn_config = cfg.AttentionConfig(
+      num_heads=16,
+      head_dim=72,
+      num_query_groups=16,
+      qkv_use_bias=True,
+      output_proj_use_bias=True,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.SEQUENTIAL,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_TANH),
+      intermediate_size=4304,
+      use_bias=True,
+  )
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.LAYER_NORM,
+      epsilon=1e-6,
+      enable_hlfb=True,
+  )
+  block_config = cfg.TransformerBlockConfig(
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+  )
+  config = cfg.ModelConfig(
+      vocab_size=0,  # Not used in image encoder.
+      num_layers=27,
+      max_seq_len=0,  # Not used in image encoder.
+      embedding_dim=1152,
+      embedding_use_bias=True,
+      image_embedding=image_embedding_config,
+      block_configs=block_config,
+      final_norm_config=norm_config,
+      enable_hlfb=True,
+  )
+  return config
+
+
+def get_fake_image_encoder_config() -> cfg.ModelConfig:
+  config = get_image_encoder_config()
+  # PaliGemma image encoder has only one block config.
+  config.block_config(0).ff_config.intermediate_size = 128
+  config.num_layers = 2
+  return config
+
+
+def build_image_encoder(checkpoint_path: str) -> SiglipVisionEncoder:
+  config = get_image_encoder_config()
+  encoder = SiglipVisionEncoder(config)
+  loader = loading_utils.ModelLoader(checkpoint_path, TENSOR_NAMES)
+  # Loose the strictness because only image encoder is being loaded.
+  loader.load(encoder, strict=False)
+  encoder.eval()
+  return encoder

ai_edge_torch/generative/examples/paligemma/verify_image_encoder.py

@@ -0,0 +1,82 @@
+# 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 image encoder of PaliGemma 3B model."""
+
+import logging
+import pathlib
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.paligemma import image_encoder
+from PIL import Image
+import requests
+import torch
+import transformers
+
+_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.",
+)
+
+
+def main(_):
+  checkpoint = "google/paligemma-3b-mix-224"
+  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)
+
+  logging.info("Loading the processor 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)
+
+  logging.info("Loading the image from: %s", _IMAGE_URL.value)
+  image = Image.open(requests.get(_IMAGE_URL.value, stream=True).raw)
+  pixel_values = processor(images=image, return_tensors="pt")["pixel_values"]
+
+  logging.info("Forwarding the original model...")
+  outputs_original = original_vision_model.forward(pixel_values=pixel_values)
+  outputs_original = outputs_original.last_hidden_state
+  logging.info("outputs_original: %s", outputs_original)
+
+  logging.info("Forwarding the reauthored model...")
+  outputs_reauthored = reauthored_model.forward(pixel_values=pixel_values)
+  logging.info("outputs_reauthored: %s", outputs_reauthored)
+
+  try:
+    assert torch.allclose(
+        outputs_original, outputs_reauthored, atol=1e-04, rtol=1e-04
+    )
+  except AssertionError as e:
+    logging.error("*** FAILED *** verify with an image")
+    raise e
+  else:
+    logging.info("*** PASSED *** verify with an image")
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/layers/attention.py

@@ -235,9 +235,10 @@ class CausalSelfAttention(nn.Module):
     k = k.reshape(B, T, -1, self.config.head_dim)
     v = v.reshape(B, T, -1, self.config.head_dim)
 
-    # Compute rotary positional embedding for query and key.
-    n_elem = int(self.config.rotary_percentage * self.config.head_dim)
-    q, k = _embed_rope(q, k, n_elem, rope)
+    if rope is not None:
+      # Compute rotary positional embedding for query and key.
+      n_elem = int(self.config.rotary_percentage * self.config.head_dim)
+      q, k = _embed_rope(q, k, n_elem, rope)
 
     if kv_cache is not None:
       kv_cache = kv_utils.update(
@@ -372,9 +373,10 @@ class CrossAttention(nn.Module):
     k = k.view(interim_shape)
     v = v.view(interim_shape)
 
-    # Compute rotary positional embedding for query and key.
-    n_elem = int(self.config.rotary_percentage * self.config.head_dim)
-    q, k = _embed_rope(q, k, n_elem, rope)
+    if rope is not None:
+      # Compute rotary positional embedding for query and key.
+      n_elem = int(self.config.rotary_percentage * self.config.head_dim)
+      q, k = _embed_rope(q, k, n_elem, rope)
 
     if kv_cache is not None:
       kv_cache = kv_utils.update(

ai_edge_torch/generative/layers/model_config.py

@@ -163,6 +163,16 @@ class TransformerBlockConfig:
   relative_attention: bool = False
 
 
+@dataclass
+class ImageEmbeddingConfig:
+  """Image embedding parameters."""
+
+  channels: int
+  # All images should be normalized to the size of [image_size * image_size].
+  image_size: int
+  patch_size: int
+
+
 @dataclass
 class ModelConfig:
   """Base configurations for building a transformer architecture."""
@@ -183,6 +193,10 @@ class ModelConfig:
 
   # Scale factor of the embedding.
   embedding_scale: Optional[float] = None
+  # Use bias term within embedding.
+  embedding_use_bias: bool = False
+  # Image embedding parameters.
+  image_embedding: Optional[ImageEmbeddingConfig] = None
 
   # Use bias term within LLM's HEAD.
   lm_head_use_bias: bool = False

ai_edge_torch/generative/layers/unet/blocks_2d.py

@@ -115,12 +115,15 @@ class AttentionBlock2D(nn.Module):
     """
     super().__init__()
     self.config = config
+    hidden_dim = config.hidden_dim
+    if not hidden_dim:
+      hidden_dim = config.dim
     self.norm = layers_builder.build_norm(
-        config.dim, config.normalization_config
+        hidden_dim, config.normalization_config
     )
     self.attention = SelfAttention(
         config.attention_batch_size,
-        config.dim,
+        hidden_dim,
         config.attention_config,
         enable_hlfb=config.enable_hlfb,
     )
@@ -172,7 +175,7 @@ class CrossAttentionBlock2D(nn.Module):
     super().__init__()
     self.config = config
     self.norm = layers_builder.build_norm(
-        config.query_dim, config.normalization_config
+        config.output_dim, config.normalization_config
     )
     self.attention = CrossAttention(
         config.attention_batch_size,
@@ -294,21 +297,30 @@ class TransformerBlock2D(nn.Module):
       hidden_states
   """
 
-  def __init__(self, config: unet_cfg.TransformerBlock2DConfig):
+  def __init__(
+      self, config: unet_cfg.TransformerBlock2DConfig, dim_override=None
+  ):
     """Initialize an instance of the TransformerBlock2D.
 
     Args:
       config (unet_cfg.TransformerBlock2Dconfig): the configuration of this
         block.
+      dim_override: in case specified, overrides config.attention_block_config.hidden_dim. Set to None by default.
     """
     super().__init__()
     self.config = config
+    attention_block_config_dim = config.attention_block_config.dim
+    attention_block_config_hidden_dim = config.attention_block_config.hidden_dim
+    if dim_override:
+      attention_block_config_dim = dim_override
+    if not attention_block_config_hidden_dim:
+      attention_block_config_hidden_dim = attention_block_config_dim
     self.pre_conv_norm = layers_builder.build_norm(
-        config.attention_block_config.dim, config.pre_conv_normalization_config
+        attention_block_config_dim, config.pre_conv_normalization_config
     )
     self.conv_in = nn.Conv2d(
-        config.attention_block_config.dim,
-        config.attention_block_config.dim,
+        attention_block_config_dim,
+        attention_block_config_hidden_dim,
         kernel_size=1,
         padding=0,
     )
@@ -318,8 +330,8 @@ class TransformerBlock2D(nn.Module):
     )
     self.feed_forward = FeedForwardBlock2D(config.feed_forward_block_config)
     self.conv_out = nn.Conv2d(
-        config.attention_block_config.dim,
-        config.attention_block_config.dim,
+        attention_block_config_hidden_dim,
+        attention_block_config_dim,
         kernel_size=1,
         padding=0,
     )
@@ -385,14 +397,18 @@ class DownEncoderBlock2D(nn.Module):
     self.config = config
     resnets = []
     transformers = []
+    hidden_channels = config.hidden_channels
+    if not hidden_channels:
+      hidden_channels = config.out_channels
     for i in range(config.num_layers):
       input_channels = config.in_channels if i == 0 else config.out_channels
       resnets.append(
           ResidualBlock2D(
               unet_cfg.ResidualBlock2DConfig(
                   in_channels=input_channels,
-                  hidden_channels=config.out_channels,
+                  hidden_channels=hidden_channels,
                   out_channels=config.out_channels,
+                  residual_out_channels=config.out_channels,
                   time_embedding_channels=config.time_embedding_channels,
                   normalization_config=config.normalization_config,
                   activation_config=config.activation_config,
@@ -589,23 +605,37 @@ class SkipUpDecoderBlock2D(nn.Module):
     """
     super().__init__()
     self.config = config
+    hidden_channels = config.hidden_channels
+    if not hidden_channels:
+      hidden_channels = config.out_channels
+    sub_block_channels = config.sub_block_channels
+    if sub_block_channels:
+      assert len(sub_block_channels) == config.num_layers, (
+          "Assertion failed: The length of 'sub_block_channels'"
+          f" ({len(sub_block_channels)}) does not match 'config.num_layers'"
+          f" ({config.num_layers})."
+      )
+    else:
+      sub_block_channels = [config.out_channels] * config.num_layers
     resnets = []
     transformers = []
     for i in range(config.num_layers):
+      resnet_in_channels = (
+          config.prev_out_channels if i == 0 else sub_block_channels[i - 1]
+      )
       res_skip_channels = (
           config.in_channels
           if (i == config.num_layers - 1)
           else config.out_channels
       )
-      resnet_in_channels = (
-          config.prev_out_channels if i == 0 else config.out_channels
-      )
+      residual_out_channel = sub_block_channels[i]
       resnets.append(
           ResidualBlock2D(
               unet_cfg.ResidualBlock2DConfig(
                   in_channels=resnet_in_channels + res_skip_channels,
-                  hidden_channels=config.out_channels,
-                  out_channels=config.out_channels,
+                  hidden_channels=hidden_channels,
+                  out_channels=sub_block_channels[i],
+                  residual_out_channels=residual_out_channel,
                   time_embedding_channels=config.time_embedding_channels,
                   normalization_config=config.normalization_config,
                   activation_config=config.activation_config,
@@ -613,7 +643,12 @@ class SkipUpDecoderBlock2D(nn.Module):
           )
       )
       if config.transformer_block_config:
-        transformers.append(TransformerBlock2D(config.transformer_block_config))
+        transformers.append(
+            TransformerBlock2D(
+                config.transformer_block_config,
+                dim_override=sub_block_channels[i],
+            )
+        )
     self.resnets = nn.ModuleList(resnets)
     self.transformers = (
         nn.ModuleList(transformers) if len(transformers) > 0 else None
@@ -623,7 +658,7 @@ class SkipUpDecoderBlock2D(nn.Module):
       if config.upsample_conv:
         self.upsample_conv = nn.Conv2d(
             config.out_channels,
-            config.out_channels,
+            sub_block_channels[0],
             kernel_size=3,
             stride=1,
             padding=1,
@@ -711,6 +746,7 @@ class MidBlock2D(nn.Module):
                 in_channels=config.in_channels,
                 hidden_channels=config.in_channels,
                 out_channels=config.in_channels,
+                residual_out_channels=config.in_channels,
                 time_embedding_channels=config.time_embedding_channels,
                 normalization_config=config.normalization_config,
                 activation_config=config.activation_config,

ai_edge_torch/generative/layers/unet/model_config.py

@@ -50,10 +50,12 @@ class ResidualBlock2DConfig:
   in_channels: int
   hidden_channels: int
   out_channels: int
+  hidden_channels: int
   normalization_config: layers_cfg.NormalizationConfig
   activation_config: layers_cfg.ActivationConfig
   # Optional time embedding channels if the residual block takes a time embedding context as input
   time_embedding_channels: Optional[int] = None
+  residual_out_channels: Optional[int] = None
 
 
 @dataclasses.dataclass
@@ -63,6 +65,7 @@ class AttentionBlock2DConfig:
   attention_config: layers_cfg.AttentionConfig
   enable_hlfb: bool = True
   attention_batch_size: int = 1
+  hidden_dim: Optional[int] = None
 
 
 @dataclasses.dataclass
@@ -101,6 +104,8 @@ class UpDecoderBlock2DConfig:
   normalization_config: layers_cfg.NormalizationConfig
   activation_config: layers_cfg.ActivationConfig
   num_layers: int
+  # The dimension of output channels of previous connected block
+  prev_out_channels: Optional[int] = None
   # Optional time embedding channels if the residual blocks take a time embedding as input
   time_embedding_channels: Optional[int] = None
   # Whether to add upsample operation after residual blocks
@@ -136,6 +141,8 @@ class SkipUpDecoderBlock2DConfig:
   transformer_block_config: Optional[TransformerBlock2DConfig] = None
   # Optional dimension of context tensor if context tensor is given as input.
   context_dim: Optional[int] = None
+  sub_block_channels: Optional[tuple] = None
+  hidden_channels: Optional[int] = None
 
 
 @dataclasses.dataclass
@@ -157,6 +164,7 @@ class DownEncoderBlock2DConfig:
   transformer_block_config: Optional[TransformerBlock2DConfig] = None
   # Optional dimension of context tensor if context tensor is given as input.
   context_dim: Optional[int] = None
+  hidden_channels: Optional[int] = None
 
 
 @dataclasses.dataclass

ai_edge_torch/generative/utilities/loader.py

@@ -157,6 +157,10 @@ class ModelLoader:
       converted_state["tok_embedding.weight"] = state.pop(
           f"{self._names.embedding}.weight"
       )
+      if model.config.embedding_use_bias:
+        converted_state["tok_embedding.bias"] = state.pop(
+            f"{self._names.embedding}.bias"
+        )
       if self._names.embedding_position is not None:
         converted_state["tok_embedding_position"] = state.pop(
             f"{self._names.embedding_position}"

ai_edge_torch/generative/utilities/verifier.py

@@ -228,7 +228,7 @@ def verify_reauthored_model(
     rtol: float = 1e-05,
     atol: float = 1e-05,
     continue_on_failure: bool = False,
-):
+) -> bool:
   """Verifies the reauthored model against the original model.
 
   It verifies the reauthored model with two methods:
@@ -237,7 +237,8 @@ def verify_reauthored_model(
   2. It compares the answer generated by the original and the reauthored model
      with a prompt.
 
-  It prints out "PASS" or "FAILED" to the console.
+  It prints out "PASS" or "FAILED" to the console. It returns True if all
+  verification passes, False otherwise.
 
   Args:
     original_model (ModelWrapper): The original model.
@@ -253,6 +254,8 @@ def verify_reauthored_model(
     continue_on_failure (bool): If True, it continues to verify the next prompt
       or input IDs even if a previous one fails.
   """
+  failure_count = 0
+
   for input_ids in forward_input_ids:
     logging.info("Verifying the reauthored model with input IDs: %s", input_ids)
     try:
@@ -261,8 +264,9 @@ def verify_reauthored_model(
       )
     except AssertionError as e:
       logging.error("*** FAILED *** verify with input IDs: %s", input_ids)
+      failure_count += 1
       if not continue_on_failure:
-        raise e
+        return False
     else:
       logging.info("*** PASSED *** verify with input IDs: %s", input_ids)
 
@@ -274,7 +278,15 @@ def verify_reauthored_model(
       )
     except AssertionError as e:
       logging.error("*** FAILED *** verify with prompts: %s", prompts)
+      failure_count += 1
       if not continue_on_failure:
-        raise e
+        return False
     else:
       logging.info("*** PASSED *** verify with prompts: %s", prompts)
+
+  if failure_count == 0:
+    logging.info("*** PASSED *** verify_reauthored_model")
+    return True
+  else:
+    logging.error("*** FAILED *** verify_reauthored_model")
+    return False

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241110.dist-info → ai_edge_torch_nightly-0.3.0.dev20241115.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20241110
+Version: 0.3.0.dev20241115
 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.dev20241110.dist-info → ai_edge_torch_nightly-0.3.0.dev20241115.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=518ziQ0TUxqum2qZXqlD8qr65pHPh8ZNLnwFC6zvK3k,4253
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=0kbL8PwrdMx4mw42_rj8uAYUeehe8jsFhw_tENefuGM,706
+ai_edge_torch/version.py,sha256=pp4KVtq0a8ju4UB5nOeiv7QDkmgpHmz5XUokSR86qfI,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=HwzfRx_DX5TLtPqwEH1_NOm38_INvHzHl4_mX67KOdQ,5448
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -12,7 +12,7 @@ ai_edge_torch/_convert/signature.py,sha256=rGpBNss3Y9FCCCcdBwDo16KqavJi8N5P0M_6W
 ai_edge_torch/_convert/to_channel_last_io.py,sha256=_31phf7TYgZY2ftpNbrdlB1RhDium1lz_BXEQ6IsMFc,2893
 ai_edge_torch/_convert/fx_passes/__init__.py,sha256=NVe-eGcm7j8jZpP2pcMhC8j5dVjgR1pPzyXhHdvKH4E,1267
 ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py,sha256=doaww8KqrgRTD5LotBVAIRFsEqzPn9R5lcGehBJOczA,9098
-ai_edge_torch/_convert/fx_passes/build_interpolate_composite_pass.py,sha256=XCVqWg_ask0Kb64PED0ZGAODsUuIgfyO2ZJM6aK-TXI,4283
+ai_edge_torch/_convert/fx_passes/build_interpolate_composite_pass.py,sha256=m_yj66V11LmWCYgA7yLtr__cy14IbC5WEJe0BE0_IPE,4339
 ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=IlZuK42kfVcRqAWZp4j2k_81T2uWo9T2558U_GPJAlU,2327
 ai_edge_torch/_convert/fx_passes/remove_non_user_outputs_pass.py,sha256=f1IUVWyhioOClsMiZzLyynoW2R17U83vA-7Q-3pGPM4,2126
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/__init__.py,sha256=lxnoH-WGLeiQIF8XjMGodjiZEFTxucl7g05N7MR9OPk,796
@@ -50,8 +50,8 @@ ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=RZDs6
 ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=oSbysiPvwp5efMbNYZop3HrxDMGiD15Tmz-HiQuTr2E,3315
 ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=RQFQDMEnIVp8PefcCTr7P0CvllKI7FVoIJLXbPLLIsc,9056
 ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=ip-Gmk4CI5f0GWSdAIdrectxQWJ0t328KCsA4nfHuGg,1736
-ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=K77k-JpdhIwm3tbBnzpw8HQsFRwAVyszxRo82fR6-q4,1762
-ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=sqltZbnyKemNvKqqi9d09i74gP-PPQFodRYfDfnhycQ,4933
+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
 ai_edge_torch/generative/examples/llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/llama/convert_to_tflite.py,sha256=P0-pByTM5tslE23ILgo7nd0nOGE25ciBRG5wKJj0bBk,2411
 ai_edge_torch/generative/examples/llama/llama.py,sha256=AMcCbuDBxEfbO-l3KiEXbUaXEJ3RLLwkHii7to7UhVo,6854
@@ -62,7 +62,9 @@ ai_edge_torch/generative/examples/openelm/openelm.py,sha256=sFakstoPDcOHSak0IGFE
 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/decoder.py,sha256=JSb9h3gcIh5oYrbLU6rI8OU8FzfWeTCFJT5XRWu4btE,3675
+ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=v19_EKALhAP9FjkINKqpv8JsVaQ6iH_7X5FpnhE6abw,5500
 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/phi/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py,sha256=rkbTtMaqSVG48cm-NTxR_LDgZmXAEBqayTm9O49oMXc,2171
 ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=3go690yX6PFeXMdpY7y4JZorAwxX0HT_b_pKZieauvk,2169
@@ -108,19 +110,19 @@ ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=7Bk8z033M-BCXJ299f
 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.py,sha256=zN3BQjA25Ej_aRU0rFnyx--K74xf5ykc02zGvUpYHeE,13295
 ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
 ai_edge_torch/generative/layers/builder.py,sha256=Z5LyzCEThgnYZeyViakaE3yJVzTGHtw13acHsAQR15U,5050
 ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
 ai_edge_torch/generative/layers/kv_cache.py,sha256=lbm-yJ1jGPtcgWS4C3FmSnB1IlxqDE7g0BLRh3PN4N4,6324
-ai_edge_torch/generative/layers/model_config.py,sha256=DdsdhTP5tZAtyWim-qW2m8HDBsYbs7boqSDb83vwgmE,6998
+ai_edge_torch/generative/layers/model_config.py,sha256=xqa7ZBEjgK4UWJAThRXb_VBFZ5KCGtDu-QaY5GXar9s,7366
 ai_edge_torch/generative/layers/normalization.py,sha256=eKAGst9rPuyRFExMcQFJO7R3iHdCtlmjeF_lITjLhwE,6498
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=CZqOoibLcHvUgrgaIIWAlmk3XgE2inzx340MN-npLoU,1347
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=gXxh3papKy4FBpGEX7VyZ7rZ1Js6aHK70Q6DKrVSckY,4154
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=JwndhL3Z31TvkdGlAoTL5PQzmKfHdRWaaE1EbaMI4Gs,27540
+ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=9jKzOfiBQ66bp1ZnVIAoREIifVNFx4aTlQeYMAx2_pA,29062
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
-ai_edge_torch/generative/layers/unet/model_config.py,sha256=raYm8Ol-EFi0zs5vNqmj2ZJCFsnQW2TfwhgDcClfwFA,9356
+ai_edge_torch/generative/layers/unet/model_config.py,sha256=pPDwLawc23pfMaPVyMJlYmxVVusjMvx-l8wBwOYOH-c,9692
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/quantize/example.py,sha256=1lfVNUd2cEyRUnoZ7BLbRJ9IN-FTKiWBtZNPFUzAiWE,1747
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=n1Fm8BFC8gJa_oiwwAOOghJyHtOXYZ4q-5ZRy4pHrIw,1957
@@ -137,12 +139,12 @@ ai_edge_torch/generative/test/test_quantize.py,sha256=8geJhKwYBU20m0mdGPD1BUFwQ0
 ai_edge_torch/generative/test/utils.py,sha256=YvEhO2HIj1LkBs5du1UxY-cGRW9HMyAYsOUhgsTrTpA,1796
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
 ai_edge_torch/generative/utilities/converter.py,sha256=17O83wVifH1vQJCI4WC3DaNiCIOyK2gys1GzohbLrRs,5554
-ai_edge_torch/generative/utilities/loader.py,sha256=b9iotIhVDX-Zc9XjIDUaLxnV395AyBnkQe3dV5YA7Co,13297
+ai_edge_torch/generative/utilities/loader.py,sha256=k5fjCokNomte4ymy9IJrEWAuCSMhsPCJfmv1y5s0ZEc,13452
 ai_edge_torch/generative/utilities/model_builder.py,sha256=89jt80UUfDzYBi-x077HBavWeuNJuYPXym9fiKCY1Tk,5278
 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=5C2cm54d9kwL7nGRX-YfnBIJny1ICNhiU-LB3IqJq2E,10075
+ai_edge_torch/generative/utilities/verifier.py,sha256=h5hGyIpYGyPZwvelbzpdkjy99Kpd4JkvhqWtQN9cm-M,10413
 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
@@ -189,8 +191,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.dev20241110.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241110.dist-info/METADATA,sha256=ECohBv1Uc5BzRcnT3r3yM8_sElMqIMmpYcnRP_nOp84,1897
-ai_edge_torch_nightly-0.3.0.dev20241110.dist-info/WHEEL,sha256=bFJAMchF8aTQGUgMZzHJyDDMPTO3ToJ7x23SLJa1SVo,92
-ai_edge_torch_nightly-0.3.0.dev20241110.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241110.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241115.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241115.dist-info/METADATA,sha256=epuuYZFnqVvLzIS0X27XMCFQpnc-dO8JJQ8DXVNv5IE,1897
+ai_edge_torch_nightly-0.3.0.dev20241115.dist-info/WHEEL,sha256=bFJAMchF8aTQGUgMZzHJyDDMPTO3ToJ7x23SLJa1SVo,92
+ai_edge_torch_nightly-0.3.0.dev20241115.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241115.dist-info/RECORD,,