ai-edge-torch-nightly 0.6.0.dev20250611__py3-none-any.whl → 0.6.0.dev20250613__py3-none-any.whl

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
 

ai_edge_torch/generative/utilities/verifier.py

@@ -351,6 +351,7 @@ def verify_reauthored_model(
         )
       except AssertionError as e:
         logging.error("*** FAILED *** verify with input IDs: %s", input_ids)
+        logging.error("*** Assertion Error: %s", e)
         failure_count += 1
         if not continue_on_failure:
           return False
@@ -366,6 +367,7 @@ def verify_reauthored_model(
         )
       except AssertionError as e:
         logging.error("*** FAILED *** verify with prompts: %s", prompts)
+        logging.error("*** Assertion Error: %s", e)
         failure_count += 1
         if not continue_on_failure:
           return False

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.dev20250611"
+__version__ = "0.6.0.dev20250613"

{ai_edge_torch_nightly-0.6.0.dev20250611.dist-info → ai_edge_torch_nightly-0.6.0.dev20250613.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.6.0.dev20250611
+Version: 0.6.0.dev20250613
 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.dev20250611.dist-info → ai_edge_torch_nightly-0.6.0.dev20250613.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=03SgV_6zF-xAWJAi3O3nHJh9aG0F05fy5C5nLfoj5-Q,806
+ai_edge_torch/version.py,sha256=YTxxjYgksS0xe-Vy2RMt-V4C9dddyuKoIcLfmaTZBh4,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
@@ -67,9 +67,9 @@ ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=m5N3M
 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=vFhb3mOxB0Pps-Nx8FQWskh-Zly7N83n_EPyG22R-oM,8204
+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
@@ -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=lSzVBFkfnVcFr3rHYQqf1OxfxgQrevkF6jGdqmZTJHA,13702
+ai_edge_torch/generative/utilities/verifier.py,sha256=HmpA5-q7s1a5WOoO5gMlfdlPp5oLeKwnx56n3DUOBBM,13806
 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.dev20250611.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.6.0.dev20250611.dist-info/METADATA,sha256=EPsBf0cXltLduemjqkavDrVuu4orHWptAtRfnnDGnoM,2074
-ai_edge_torch_nightly-0.6.0.dev20250611.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.6.0.dev20250611.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.6.0.dev20250611.dist-info/RECORD,,
+ai_edge_torch_nightly-0.6.0.dev20250613.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.6.0.dev20250613.dist-info/METADATA,sha256=WVeg8CPHQeHC3GNBou1kvlzF8zXyPtpSadLTayhpmew,2074
+ai_edge_torch_nightly-0.6.0.dev20250613.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.6.0.dev20250613.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.6.0.dev20250613.dist-info/RECORD,,