ai-edge-torch-nightly 0.6.0.dev20250610__py3-none-any.whl → 0.6.0.dev20250612__py3-none-any.whl

ai_edge_torch/generative/examples/amd_llama_135m/verify_util.py

@@ -60,6 +60,7 @@ def verify_amd_llama_135m(
   reauthored_model = amd_llama_135m.build_model(
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/examples/deepseek/verify_util.py

@@ -60,6 +60,7 @@ def verify_deepseek_r1_distill_1_5b(
   reauthored_model = deepseek.build_model(
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/examples/gemma/verify_gemma1.py

@@ -20,6 +20,7 @@ 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
 
 
@@ -39,7 +40,9 @@ def main(_):
   checkpoint = kagglehub.model_download("google/gemma/pyTorch/2b-it")
 
   logging.info("Building the reauthored model from: %s", checkpoint)
-  reauthored_model = gemma1.build_2b_model(checkpoint)
+  reauthored_model = gemma1.build_2b_model(
+      checkpoint, mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN
+  )
 
   verify_util.verify_reauthored_gemma_model(
       checkpoint=checkpoint,

ai_edge_torch/generative/examples/gemma/verify_util.py

@@ -62,8 +62,13 @@ class GemmaWrapper(verifier.ModelWrapper):
     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])
-    local_mask_cache = attn_utils.build_sliding_window_mask_cache(
-        tokens.shape[1], self.model.config.sliding_window_size)
+    local_mask_cache = (
+        attn_utils.build_sliding_window_mask_cache(
+            tokens.shape[1], self.model.config.sliding_window_size
+        )
+        if self.model.config.sliding_window_size
+        else None  # Do not use local mask cache if sliding window size is None.
+    )
     _, logits = self.model.forward(
         input_token_ids=tokens[0, :actual_input_len].unsqueeze(0),
         input_positions=input_pos,
@@ -75,6 +80,8 @@ class GemmaWrapper(verifier.ModelWrapper):
         top_ps=torch.tensor([1.0], dtype=torch.float),
         top_ks=torch.tensor([1], dtype=torch.long),
         local_mask=local_mask_cache.index_select(2, input_pos)
+        if local_mask_cache
+        else None,
     )
     return logits
 
@@ -143,9 +150,7 @@ def verify_reauthored_gemma_model(
   return verifier.verify_reauthored_model(
       original_model=GemmaWrapper(original_model),
       reauthored_model=verifier.ReauthoredModelWrapper(
-          reauthored_model,
-          mask_as_input=mask_as_input,
-          kv_layout=kv_layout,
+          reauthored_model, mask_as_input, kv_layout
       ),
       tokenizer=GemmaTokenizerWrapper(original_model.tokenizer),
       generate_prompts=generate_prompts,
@@ -171,7 +176,11 @@ def verify_gemma2(
   """
   checkpoint_path = os.path.join(checkpoint_dir, weight_filename)
   logging.info("Building the reauthored model from: %s", checkpoint_path)
-  reauthored_model = gemma2.build_2b_model(checkpoint_path, custom_loader)
+  reauthored_model = gemma2.build_2b_model(
+      checkpoint_path,
+      custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
+  )
 
   return verify_reauthored_gemma_model(
       checkpoint=checkpoint_dir,
@@ -193,7 +202,11 @@ def verify_gemma1_with_custom_loader(checkpoint_dir: str) -> bool:
   weight_filename = "gemma-2b-it.ckpt"
   checkpoint_path = os.path.join(checkpoint_dir, weight_filename)
   custom_loader = loader.get_custom_loader(checkpoint_path)
-  reauthored_model = gemma1.build_2b_model(checkpoint_path, custom_loader)
+  reauthored_model = gemma1.build_2b_model(
+      checkpoint_path,
+      custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
+  )
   return verify_reauthored_gemma_model(
       checkpoint=checkpoint_dir,
       variant="2b",

ai_edge_torch/generative/examples/gemma3/verify_util.py

@@ -93,21 +93,27 @@ class GemmaWrapper(verifier.ModelWrapper):
 class UnifiedGemma3Wrapper(verifier.ReauthoredModelWrapper):
   """Unified Gemma3 model wrapper for verification."""
 
-  def __init__(self, model: torch.nn.Module):
-    super().__init__(model, kv_layout=kv_utils.KV_LAYOUT_TRANSPOSED)
+  def __init__(
+      self,
+      model: torch.nn.Module,
+      kv_cache_max_len: int = verifier.DEFAULT_KV_CACHE_MAX_LEN,
+  ):
+    super().__init__(
+        model,
+        kv_layout=kv_utils.KV_LAYOUT_TRANSPOSED,
+        kv_cache_max_len=kv_cache_max_len,
+    )
 
   def _init_kv_cache(self):
     return kv_utils.KVCache.from_model_config(
-        self.model.model.config, kv_layout=self.kv_layout
+        self.kv_cache_max_len, self.model.model.config, kv_layout=self.kv_layout
     )
 
   def forward(
       self, tokens: torch.Tensor, pixel_values: torch.Tensor = None
   ) -> torch.Tensor:
     """Forwards the model."""
-    mask = attn_utils.build_causal_mask_cache(
-        self.model.model.config.kv_cache_max_len
-    )
+    mask = attn_utils.build_causal_mask_cache(self.kv_cache_max_len)
     input_pos = torch.arange(0, tokens.shape[1], dtype=torch.int)
     mask = mask.index_select(2, input_pos)
     output = self.model.model.forward(
@@ -127,9 +133,7 @@ class UnifiedGemma3Wrapper(verifier.ReauthoredModelWrapper):
     tokens = torch.tensor([input_ids])
     input_pos = torch.arange(0, tokens.shape[1], dtype=torch.int)
     kv_cache = self._init_kv_cache()
-    mask_cache = attn_utils.build_causal_mask_cache(
-        self.model.model.config.kv_cache_max_len
-    )
+    mask_cache = attn_utils.build_causal_mask_cache(self.kv_cache_max_len)
     for _ in range(max_new_tokens):
       mask = mask_cache.index_select(2, input_pos)
       output = self.model.model.forward(
@@ -245,7 +249,11 @@ def verify_gemma3(
 
   if variant == "1b":
     reauthored_model = UnifiedGemma3Wrapper(
-        gemma3.build_model_1b(gemma3_model_path, custom_loader)
+        gemma3.build_model_1b(
+            gemma3_model_path,
+            custom_loader,
+            mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
+        )
     )
   else:
     raise ValueError(f"Unsupported Gemma3 variant: {variant}")

ai_edge_torch/generative/examples/hammer/verify_util.py

@@ -66,6 +66,7 @@ def verify_hammer(
   reauthored_model = _BUILDER[model_size](
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/examples/llama/verify_util.py

@@ -65,6 +65,7 @@ def verify_llama_3_2(
   reauthored_model = _BUILDER[model_size](
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/examples/openelm/verify_util.py

@@ -60,6 +60,7 @@ def verify_openelm(
   reauthored_model = openelm.build_model(
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/examples/paligemma/verify.py

@@ -66,7 +66,9 @@ class ReauthoredPaliGemmaWrapper(verifier.ReauthoredModelWrapper):
   """Reauthored PaliGemma model wrapper."""
 
   def _init_kv_cache(self):
-    return kv_cache.KVCache.from_model_config(self.model.config.decoder_config)
+    return kv_cache.KVCache.from_model_config(
+        self.kv_cache_max_len, self.model.config.decoder_config
+    )
 
 
 def main(_):
@@ -88,7 +90,9 @@ def main(_):
 
   logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
   reauthored_model = paligemma.build_model(
-      reauthored_checkpoint, version=int(_VERSION.value)
+      reauthored_checkpoint,
+      version=int(_VERSION.value),
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
   wrapped_reauthored_model = ReauthoredPaliGemmaWrapper(reauthored_model)
 

ai_edge_torch/generative/examples/paligemma/verify_decoder.py

@@ -51,7 +51,9 @@ def main(_):
   )
   reauthored_checkpoint = pathlib.Path(cached_config_file).parent
   logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
-  reauthored_model = decoder.build_decoder(reauthored_checkpoint)
+  reauthored_model = decoder.build_decoder(
+      reauthored_checkpoint, mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN
+  )
 
   logging.info("Loading the tokenizer from: %s", checkpoint)
   # It works only when GemmaTokenizerFast is available. In some environments,

ai_edge_torch/generative/examples/paligemma/verify_decoder2.py

@@ -48,7 +48,9 @@ def main(_):
   original_language_model = original_full_model.eval().language_model
 
   logging.info("Building the reauthored model from: %s", checkpoint)
-  reauthored_model = decoder2.build_decoder2(checkpoint)
+  reauthored_model = decoder2.build_decoder2(
+      checkpoint, mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN
+  )
 
   logging.info("Loading the tokenizer from: %s", checkpoint)
   # It works only when GemmaTokenizerFast is available. In some environments,

ai_edge_torch/generative/examples/phi/verify_util.py

@@ -67,6 +67,7 @@ def verify_phi(
   reauthored_model = _BUILDER[version](
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/examples/qwen/verify_util.py

@@ -80,6 +80,7 @@ def verify_qwen(
   reauthored_model = _BUILDER[model_version][model_size](
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/examples/qwen_vl/verify.py

@@ -55,7 +55,9 @@ class ReauthoredQwenVLWrapper(verifier.ReauthoredModelWrapper):
   """Reauthored Qwen VL model wrapper."""
 
   def _init_kv_cache(self):
-    return kv_cache.KVCache.from_model_config(self.model.config.decoder_config)
+    return kv_cache.KVCache.from_model_config(
+        self.kv_cache_max_len, self.model.config.decoder_config
+    )
 
 
 def main(_):
@@ -73,7 +75,10 @@ def main(_):
   )
   reauthored_checkpoint = pathlib.Path(cached_config_file).parent
   logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
-  reauthored_model = qwen_vl.build_model(str(reauthored_checkpoint))
+  reauthored_model = qwen_vl.build_model(
+      str(reauthored_checkpoint),
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
+  )
   wrapped_reauthored_model = ReauthoredQwenVLWrapper(reauthored_model)
 
   logging.info("Loading the processor from: %s", checkpoint)

ai_edge_torch/generative/examples/qwen_vl/verify_decoder.py

@@ -52,7 +52,9 @@ def main(_):
   )
   reauthored_checkpoint = pathlib.Path(cached_config_file).parent
   logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
-  reauthored_model = decoder.build_decoder(reauthored_checkpoint)
+  reauthored_model = decoder.build_decoder(
+      reauthored_checkpoint, mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN
+  )
 
   # Verify the reauthored model only with input IDs because the original decoder
   # does not support generate() with prompts.

ai_edge_torch/generative/examples/smollm/verify_util.py

@@ -65,6 +65,7 @@ def verify_smollm_135m(
   reauthored_model = _BUILDER[model_version](
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/examples/tiny_llama/verify_util.py

@@ -61,6 +61,7 @@ def verify_tiny_llama(
   reauthored_model = tiny_llama.build_model(
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/utilities/verifier.py

@@ -24,6 +24,10 @@ import torch
 
 ExportConfig = export_config.ExportConfig
 
+# The default KV cache maximum length for the reauthored model for verification.
+# If desired, each model can override it when a model is initialized.
+DEFAULT_KV_CACHE_MAX_LEN = 256
+
 
 class ModelWrapper(torch.nn.Module):
   """A wrapper for the model to be verified.
@@ -92,7 +96,7 @@ class ReauthoredModelWrapper(ModelWrapper):
       model: torch.nn.Module,
       mask_as_input: bool = False,
       kv_layout: kv_utils.KVLayout = kv_utils.KV_LAYOUT_DEFAULT,
-      kv_cache_max_len: int = 1280,
+      kv_cache_max_len: int = DEFAULT_KV_CACHE_MAX_LEN,
   ):
     """Wraps a reauthored model with some options."""
     super().__init__(model)

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -331,6 +331,18 @@ def _aten_sym_size_int(lctx, x: ir.Value, dim: int):
   return stablehlo.get_dimension_size(x, dim)
 
 
+# Lowering for the subtraction operator (`-`).
+# Handles cases where one operand is an integer (scalar) and the other is a
+# tensor, broadcasting the scalar to the tensor's shape before subtraction.
+@lower(operator.sub)
+def _operator_sub(lctx, self: int | ir.Value, other: int | ir.Value):
+  if isinstance(self, int) and isinstance(other, ir.Value):
+    self = utils.splat(self, other.type.element_type, other.type.shape)
+  if isinstance(other, int) and isinstance(self, ir.Value):
+    other = utils.splat(other, self.type.element_type, self.type.shape)
+  return stablehlo.subtract(self, other)
+
+
 # Lowering for the multiplication operator (`*`).
 # Handles cases where one operand is an integer (scalar) and the other is a
 # tensor, broadcasting the scalar to the tensor's shape before multiplication.

ai_edge_torch/version.py

@@ -15,4 +15,4 @@
 
 # The next version of ai-edge-torch.
 # The minor version code should be bumped after every release.
-__version__ = "0.6.0.dev20250610"
+__version__ = "0.6.0.dev20250612"

{ai_edge_torch_nightly-0.6.0.dev20250610.dist-info → ai_edge_torch_nightly-0.6.0.dev20250612.dist-info}/METADATA

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

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=lemyLCNoGYRnJsmDuGZu7qOqLbLqG6CGDFtu3ue1syU,129
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=A7loFu8jE9CsXsfMmHYZ-KDFJiaD8Kkqwm_9d3IVzk0,5638
-ai_edge_torch/version.py,sha256=lAP6pGKB6xmTObkL1qwu8aAg_Q3qpirKapVhpDQ8Bg8,806
+ai_edge_torch/version.py,sha256=fUrxl7ESwOFUVUJ0CNVkVhUDStq6UHtLrAxcH5hdXzA,806
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=iQk3R-pLq4c1nfLqPB4xTRj78gghxPGzJCJtILLdg5o,6123
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -56,37 +56,37 @@ ai_edge_torch/generative/examples/amd_llama_135m/__init__.py,sha256=hHLluseD2R0H
 ai_edge_torch/generative/examples/amd_llama_135m/amd_llama_135m.py,sha256=Z2oUQ--pXPgNQyGLdWDNdyC90cSZzDN9YnSclRrPWAg,2703
 ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py,sha256=ze0MPZsUJpIBSCMJB4BrnJT5MFomRcffvuZEqKG20AM,1848
 ai_edge_torch/generative/examples/amd_llama_135m/verify.py,sha256=uyBg5-trxQEjEHDZMX4qojkcsZgERUiPqIgR9n0_AY4,1311
-ai_edge_torch/generative/examples/amd_llama_135m/verify_util.py,sha256=OJTjULdz-8LVd5DhmX8isxUI1PzhE-dQLj7--5DHfJ8,2961
+ai_edge_torch/generative/examples/amd_llama_135m/verify_util.py,sha256=J_FwH5KzS1tlyxcICSPir35gP2tWqRQ5wyVTsV9Uk8w,3018
 ai_edge_torch/generative/examples/deepseek/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
 ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py,sha256=DvoHDBNc0GPW3kOv0TbGg0PnNTKijW_6z9LfQMku3RY,1816
 ai_edge_torch/generative/examples/deepseek/deepseek.py,sha256=cm3rLw6aHFWQswwCqfvwxr05JG0maSzNVn1pBqD19qY,2816
 ai_edge_torch/generative/examples/deepseek/verify.py,sha256=HkvgEyGb-V_f6mWfyeN7Ai5uADAVQNzWvkygaKJiLAc,1344
-ai_edge_torch/generative/examples/deepseek/verify_util.py,sha256=jui_16J0C0VhICGBJDiO7Br5l5QCrhm-AohXYuUyyqQ,2971
+ai_edge_torch/generative/examples/deepseek/verify_util.py,sha256=xf2-jhv-ORwr_-Rs2UBpO2vcAwJNZIdmOxpcuKIQaww,3028
 ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=m5N3MxNX4Yu8-0vXRszkMgfVuFKN6gskVoUIGqonJFk,1814
 ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=fR4869w1UZIuUVGLDdPof0IDEMDhqn2ej5kl7QW9vNA,1882
 ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=TH9XQAp5p4S829XbaWbJQZBwB18WizDRIQMsUkKqj38,3377
 ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=E6jotWYYIx6SXUWqurKWjiZpbfj_M2jJrBc2rQ90z1s,11782
-ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=ip-Gmk4CI5f0GWSdAIdrectxQWJ0t328KCsA4nfHuGg,1736
+ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=USyist332qZkhCBof2tJwQSqtnKjTQsKAK_jCE_CO2U,1853
 ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=eAM7EVVMW-QCqjeZEss7TOkVKArgUs1La51LAC-5a9A,1962
-ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=kSGifZobCvvRVoNAb8OumwuMKea-Mar_wIlJ5S2qF0Y,8101
+ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=b12naCj4zZxOjkIKrd08qovtajYuX-Ba3fbrv6kkDZs,8410
 ai_edge_torch/generative/examples/gemma3/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
 ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=UEDNN3JmI31WfE2pvacxeJpqumKK86L2dEus3yTURaY,2114
 ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=1UVv9SFFg5degX3wf-Fefx7nor1AzJj2NWBVuo8bRnM,15540
 ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=fFMyIS8si3GdwW8EsdhYk1OKyg_27xDv1HTQ2Gv4N8E,6616
 ai_edge_torch/generative/examples/gemma3/image_encoder.py,sha256=tUOI99kdny33qcDM7-z0R6F-1aU1lZ24kG5zeLVdwow,5129
 ai_edge_torch/generative/examples/gemma3/verify_gemma3.py,sha256=v8oNXFICmVOtQxfO7IhZ8GnbvotEkDi9lzYHjoQyOso,2464
-ai_edge_torch/generative/examples/gemma3/verify_util.py,sha256=ZgoocQSTJqFEYbNNvBaO0Be4_bgSUEnihVm47mCTzTg,9864
+ai_edge_torch/generative/examples/gemma3/verify_util.py,sha256=5OmUwz38kVHYLA-v8U8evvDN9da2WioZtGo-XK6yq1o,10067
 ai_edge_torch/generative/examples/hammer/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/hammer/convert_to_tflite.py,sha256=BujBJml_ShWpKJyUgKAJqV4C6yd1uODjKg8fOgchPKI,2041
 ai_edge_torch/generative/examples/hammer/hammer.py,sha256=0fz6LmZLXrVyty6w5T6LWRoqvQ-5XYREH6gFzB_gOu0,3737
 ai_edge_torch/generative/examples/hammer/verify.py,sha256=iuKFMkI1VZc9_ESUQr32cHWHsDtYOAF702TdqhD3Ns4,1589
-ai_edge_torch/generative/examples/hammer/verify_util.py,sha256=OtmLYBd2AlliSqj_5rNZokfwoXt2pBwwjJAYRv_dKNg,2905
+ai_edge_torch/generative/examples/hammer/verify_util.py,sha256=vB5kiFcDBM2G-AkTByJarxur-kJT_J2wGCW4mzf5whc,2962
 ai_edge_torch/generative/examples/llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/llama/convert_to_tflite.py,sha256=K-hke37VBMWakuqCeRmpaZBujGWRykDIaUBKJQ0nI5g,2022
 ai_edge_torch/generative/examples/llama/llama.py,sha256=PpC3idSFSATjQzcLv8oXwo66_5zU-B_np2YVAs2Naog,6602
 ai_edge_torch/generative/examples/llama/verify.py,sha256=XoF_-kxdryjt0Bt_YeHnIbLfjwFxSVioTSEG75moDr8,1581
-ai_edge_torch/generative/examples/llama/verify_util.py,sha256=JIAkG7TEaR_5scibyM-zgS9G2KrIfUBgX4vHDJwmaBo,3019
+ai_edge_torch/generative/examples/llama/verify_util.py,sha256=YVwoiQNwnwc_yx2ra3KYMCmqY-PiWm1kYZQjcw2BXws,3076
 ai_edge_torch/generative/examples/moonshine/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/moonshine/convert_moonshine_to_tflite.py,sha256=_GkaSkregS3NWN38UGXxj4pED5gtQGaaPZx5_CZ0TVM,1657
 ai_edge_torch/generative/examples/moonshine/moonshine.py,sha256=nZ2b8u4TmsB5sgdClgAuH8E78bcTv9RCnF9666HqP2M,3394
@@ -94,16 +94,16 @@ ai_edge_torch/generative/examples/openelm/__init__.py,sha256=hHLluseD2R0Hh4W6XZR
 ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=C5KEMMle8v6yZWrJPYoN-1CAjeZv7GQEJvyL28DeEYs,1813
 ai_edge_torch/generative/examples/openelm/openelm.py,sha256=Zfoe0VzNHiCK3CIlp0UYzxxJuE9yZXrSCb84S24LMM0,4448
 ai_edge_torch/generative/examples/openelm/verify.py,sha256=kRoNEUEsrz51PFSeTPcrYsPBQRLtUmYM3t_-Jl0oFqM,1300
-ai_edge_torch/generative/examples/openelm/verify_util.py,sha256=arX-ZGWVFfsyZ7n_V8ivaw8fqpmK3CmcxfsS0b1pOBY,2913
+ai_edge_torch/generative/examples/openelm/verify_util.py,sha256=3Tay96oEiV3Gtz-K2Gu2-tov2Y4mMpztPhwXfs4c5EA,2970
 ai_edge_torch/generative/examples/paligemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=qwxnytRMcAYlSe2Ll3TvhB05YozEVnkARQhlaaQ2UHI,2278
 ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=XRilFiLrRoWboKM6ptowtRGscek7TqSOM9YyWNhpldA,5298
 ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=Y3CLRFs3W9JSNFYCMO-CYd7iyA57b2cUnitGhB0RAbo,5847
 ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=V0RrkocOe-y2EDvcg8DMcSpWzzHUruQAEofHn20Jw7M,5589
 ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=83MHNqbdza4j1CWWusVC6upKxJonzFfpBTqW0S6iBtg,6323
-ai_edge_torch/generative/examples/paligemma/verify.py,sha256=myHdeIAtVTOqb915h661CnvjvFkwmihy3Vp4UrKHb5I,6195
-ai_edge_torch/generative/examples/paligemma/verify_decoder.py,sha256=al5wMPWri4IRVWrLmCplPi6uoCzwh0vBHMGnCt-XUqo,2690
-ai_edge_torch/generative/examples/paligemma/verify_decoder2.py,sha256=tm-UfLr0YeBRVcQsWLBOMWI9JUzHmtPEbYK2vpITpqY,2534
+ai_edge_torch/generative/examples/paligemma/verify.py,sha256=Se2T5mBQAs2hSlWfyZGkERqYbAdmwcFQg8BPnHxtT10,6296
+ai_edge_torch/generative/examples/paligemma/verify_decoder.py,sha256=Z-Xx2mpygBcTcoPs1sY8Fd2eO95D-9RldmDGHK0vCx0,2751
+ai_edge_torch/generative/examples/paligemma/verify_decoder2.py,sha256=ryn3OpMFeVwUXoywl73XyypivjghSBojdLDLzvF3XI4,2595
 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=th7EYY2NmT4qLMtQsEIU697sGhwZi24V8OwB6xZRHEw,1802
@@ -115,7 +115,7 @@ ai_edge_torch/generative/examples/phi/phi4.py,sha256=n-LLniR-Q3Nyz4JCTvKZe9hBHVW
 ai_edge_torch/generative/examples/phi/verify.py,sha256=fIWgqypLQ3uOQ1u5uuklYiQSJPhKCTYIBACjrp7DbMA,1346
 ai_edge_torch/generative/examples/phi/verify_phi3.py,sha256=TwIu2xUPQyMUTFdz29E2y75wfq4c1fGJnT3QfA3eS1s,1347
 ai_edge_torch/generative/examples/phi/verify_phi4.py,sha256=2MlgQrfRkhE7Dya8MIixGwpqEZYdPjQkUGB47Mt1hSI,1343
-ai_edge_torch/generative/examples/phi/verify_util.py,sha256=30l5EJGcAR8u7yOPc7uQHzRxZaUIfBooZJoPdlMou9c,2980
+ai_edge_torch/generative/examples/phi/verify_util.py,sha256=te2ws_4ZsuMr_d3CkN-vCIAnWGq7g0eDJ0RLna3E2gw,3037
 ai_edge_torch/generative/examples/qwen/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/qwen/convert_to_tflite.py,sha256=XQee8T78I4CqFghZw6sSsrtryWcA_NlOzg2wMfHzHwY,2064
 ai_edge_torch/generative/examples/qwen/convert_v3_to_tflite.py,sha256=P10_5tnuvjXNYyWTp0Fzzj-8moFWyONDBhcGi2yvdp4,2070
@@ -123,21 +123,21 @@ ai_edge_torch/generative/examples/qwen/qwen.py,sha256=VrrTS83q1Fzp34SjUf0pFiCoMw
 ai_edge_torch/generative/examples/qwen/qwen3.py,sha256=cJcN5nH-yAglKtm5Yc5NsLMQvkJ4Drz_IsQJzS91Veo,5329
 ai_edge_torch/generative/examples/qwen/verify_qwen2.py,sha256=ry-c2QesH-0KnrSQygfjUFs6d4kOFvJz2ts_8mP156I,1659
 ai_edge_torch/generative/examples/qwen/verify_qwen3.py,sha256=hmE0gdyzgcDpEDcWiwOzKQcxt4XeAe9DPRspy_I-lc8,1628
-ai_edge_torch/generative/examples/qwen/verify_util.py,sha256=9XhDQyIHFoJ6BYota1lIzwsBrN7t_Mtugpgbjgv17Gg,3277
+ai_edge_torch/generative/examples/qwen/verify_util.py,sha256=dGHWzaGkuw_LLNe87GJBNGNROUbRLPCE9-U_juZpUAw,3334
 ai_edge_torch/generative/examples/qwen_vl/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
 ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py,sha256=aC7ONr9hoQMGtbvfq_JBqGhSD6UYHT6fmMYdKp6dkh8,2477
 ai_edge_torch/generative/examples/qwen_vl/decoder.py,sha256=xmRDK3fmyNnY5GlHStIymYJZDUNFyjH8OjPoPxNlAKk,4307
 ai_edge_torch/generative/examples/qwen_vl/image_encoder.py,sha256=nuAHb-RXqTffpwjwCHOd_2mCrSMwL6Q1z_yjsU64gmI,14992
 ai_edge_torch/generative/examples/qwen_vl/qwen_vl.py,sha256=ILxGsI1u4yYlrWTK63yEdsxomdssVrSv7XLxWIU9mqQ,7870
-ai_edge_torch/generative/examples/qwen_vl/verify.py,sha256=4WKgAFQNQzwmeJhC8ayI5vjGj9ko6VcU2HA3VAkhHug,5812
-ai_edge_torch/generative/examples/qwen_vl/verify_decoder.py,sha256=xPWoOBLh2eK12KEhELLYymfL7xvc0chmYC98c6x37oo,2602
+ai_edge_torch/generative/examples/qwen_vl/verify.py,sha256=m-SdwHlxjOHsw740_i3T6SlQTzMjQrruauw18ZdWpmI,5917
+ai_edge_torch/generative/examples/qwen_vl/verify_decoder.py,sha256=ZuVRPLnxpaJWAvZLvJa-u1hZIkqc1D6Ew8FxdXtLlyM,2663
 ai_edge_torch/generative/examples/qwen_vl/verify_image_encoder.py,sha256=PZ392nDoJG2OmHZ_7Jet3Zu1JkN6QErxKcDc7a-PPds,3126
 ai_edge_torch/generative/examples/smollm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/smollm/convert_to_tflite.py,sha256=CsXboOmzk2GB506a-Je-Js-XFlUmikVX7dhBFQAmiUs,2040
 ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py,sha256=2BFScI1RbpK3YqJcsvacGz9haKR_gwGsSNHfYmlOvNg,2045
 ai_edge_torch/generative/examples/smollm/smollm.py,sha256=3QbVefwSlFjhze7ju3P6xgMex2d1AJpF6Ii_qlbfJ34,3656
 ai_edge_torch/generative/examples/smollm/verify.py,sha256=tXiAnwOnqgwyoa8dI4tCBiGUXkOMfdE9MUkkY_Bc4Ig,1603
-ai_edge_torch/generative/examples/smollm/verify_util.py,sha256=_Z5jtNzQnULeWrXWC1LIpEawKwL_wA1_FtRc6A4-s4k,3019
+ai_edge_torch/generative/examples/smollm/verify_util.py,sha256=tHwmo8E474gLKAu4IrG7t2hmy-3Go_BJI1mgu0KkWOY,3076
 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=lSCRZsoLjH_kqasRMwCy5IogkhyJdwcHKsPEfyxsXCQ,6112
@@ -165,7 +165,7 @@ ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=hHLluseD2R0Hh4W6
 ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=dYqb9OhpZUBIdww1lUdMjFhEo3IYhpvitvrZfO0uGSs,1827
 ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=fRuPLHwZu8Am8h0Y3dT5_1vsCEnm6euGWB_KHc3xBXk,2735
 ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=qzUsW8tJlAD9mqRxDSAcz5xSUKNlBz_DykA3PwUHMwc,1336
-ai_edge_torch/generative/examples/tiny_llama/verify_util.py,sha256=AxJeLbpDfKN-ztHYTr2qjoTpWQ-n9yeqYKhorFEBzLs,2937
+ai_edge_torch/generative/examples/tiny_llama/verify_util.py,sha256=z6vPBXDWAL6gN45z-Wm1B2bnZuJf3ErruSuFaEPPEGY,2994
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=PFSMsA1vfBfrV9ssBCkYJNl8Hx_bLdWjN01iyjPM5jE,1094
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=myGjal5A8yIBoqgArd2k40rZmCgD1Ya369KR7182bhI,2129
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -218,7 +218,7 @@ ai_edge_torch/generative/utilities/t5_loader.py,sha256=tEsfy8-ymzbbjOIc-oesXF3yG
 ai_edge_torch/generative/utilities/test_utils.py,sha256=fhUMCMxoeMzxYbOCjNeX5wbQmF6Y88Hi52FtRiZYJAk,1147
 ai_edge_torch/generative/utilities/transformers_verifier.py,sha256=l54bmmhj613eB2oCoONIAKEHhf8TQOhC9Gwjp6lxHAE,1659
 ai_edge_torch/generative/utilities/types.py,sha256=gZI9hIPB3XAo4oecKIIoVDfiyibLaSNFhecPFx4VDTM,2913
-ai_edge_torch/generative/utilities/verifier.py,sha256=HQhEpOmUCN2uu6I47um1x2azAa5hy9-kPbEvFUov2zI,13499
+ai_edge_torch/generative/utilities/verifier.py,sha256=lSzVBFkfnVcFr3rHYQqf1OxfxgQrevkF6jGdqmZTJHA,13702
 ai_edge_torch/hlfb/__init__.py,sha256=sH4um75na-O8tzxN6chFyp6Y4xnexsE7kUQpZySv6dE,735
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=JsVmYrM_JEuN_smMHXUsRlo3Liapp7UyktbPpPARwDk,5386
 ai_edge_torch/hlfb/mark_pattern/fx_utils.py,sha256=YCtMgu-4w2BQ5fpnlpWC6IauKPf_tVqc7Ff91OTqlSw,1796
@@ -248,7 +248,7 @@ ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNi
 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=uJ-niilt1c-D6QJzLwgvCUf62le_JsxQTlqj_iP_Ps0,1009
-ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=fEWjIdEpDIqT1EYLZE13O9A41OuaNdbfBrv3vNxS9gI,11601
+ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=VWb5HEeVljnuXi1eecKp1ieOIcBrSLlu7YIZnxnrozU,12198
 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/_decomp_registry.py,sha256=ybOdoFE5HIJTkyiYcc73zpyUyUpioVnAca6k0wyJPs4,2572
@@ -268,8 +268,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.6.0.dev20250610.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.6.0.dev20250610.dist-info/METADATA,sha256=hEShiHfbg6-_F12iYg8Q-GgPPyun1JAEsDhR_UJBibI,2074
-ai_edge_torch_nightly-0.6.0.dev20250610.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.6.0.dev20250610.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.6.0.dev20250610.dist-info/RECORD,,
+ai_edge_torch_nightly-0.6.0.dev20250612.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.6.0.dev20250612.dist-info/METADATA,sha256=PMopQVfn8nbh8hkC9zHYzl0bF82XoNsGfOfff4K7hBk,2074
+ai_edge_torch_nightly-0.6.0.dev20250612.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.6.0.dev20250612.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.6.0.dev20250612.dist-info/RECORD,,