ai-edge-torch-nightly 0.6.0.dev20250528__py3-none-any.whl → 0.6.0.dev20250529__py3-none-any.whl

ai_edge_torch/generative/layers/model_config.py

@@ -268,6 +268,12 @@ class ModelConfig:
   # export.
   use_mask_cache: bool = True
 
+  # An interleaved sequence of the attention types used in the model.
+  # E.g. [AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING,
+  # AttentionType.GLOBAL] means that the model has an attention pattern of 2
+  # local attentions followed by a global attention in a repeated pattern.
+  attention_patterns: Optional[Sequence[AttentionType]] = None
+
   @property
   def kv_cache_max(self) -> int:
     if self.kv_cache_max_len > 0:
@@ -286,3 +292,19 @@ class ModelConfig:
   @property
   def causal_mask_value(self) -> float:
     return self.block_config(0).attn_config.causal_mask_value
+
+  def check_if_global_attention_layer(self, layer_idx: int) -> bool:
+    """Returns True if the layer is a global attention layer."""
+    if self.attention_patterns is None:
+      # If attention_patterns is not set, we assume the model has global
+      # attention.
+      return True
+    assert layer_idx >= 0 and layer_idx < self.num_layers, (
+        "Layer index {layer_idx} is out of range for num_layers:"
+        f" {self.num_layers}"
+    )
+
+    return (
+        self.block_config(layer_idx).attn_config.attn_type
+        == AttentionType.GLOBAL
+    )

ai_edge_torch/generative/utilities/test_utils.py

@@ -0,0 +1,29 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Test utils for generative layers."""
+from typing import Sequence
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+import torch
+
+
+def initialize_kv_cache_all_zeros(
+    kv_shape: Sequence[int],
+    layout: kv_utils.KVLayout = kv_utils.KV_LAYOUT_DEFAULT,
+) -> kv_utils.KVCacheEntry:
+  return kv_utils.KVCacheEntry(
+      k_cache=torch.zeros(kv_shape, dtype=torch.float32),
+      v_cache=torch.zeros(kv_shape, dtype=torch.float32),
+      kv_layout=layout,
+  )

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.dev20250528"
+__version__ = "0.6.0.dev20250529"

{ai_edge_torch_nightly-0.6.0.dev20250528.dist-info → ai_edge_torch_nightly-0.6.0.dev20250529.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.6.0.dev20250528
+Version: 0.6.0.dev20250529
 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.dev20250528.dist-info → ai_edge_torch_nightly-0.6.0.dev20250529.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=wxjSFq_rBSxSqbUE8E8EJTCkgvgaRLjq_ZuAM-IZpCU,5606
-ai_edge_torch/version.py,sha256=MuW0AEZVV7KlferCv485Nb_a1fonWf_MSQEeft5h9yU,806
+ai_edge_torch/version.py,sha256=j84SLcYNp64OJ8Smn3CjFhF-ojE6L2pi7LJltVN_0Hg,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
@@ -180,7 +180,7 @@ ai_edge_torch/generative/layers/feed_forward.py,sha256=_GmtHxwL068l9gh_F_WFcFk7L
 ai_edge_torch/generative/layers/feed_forward_test.py,sha256=Y5l1eC9NgfYixHcfIfE1W4FGh7oC-9UGGyHdKS9tQKc,1880
 ai_edge_torch/generative/layers/kv_cache.py,sha256=b-7shzDaKexmvQF7P3SiAmIz4ZofjYWv3m5u71GojsA,10460
 ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
-ai_edge_torch/generative/layers/model_config.py,sha256=0FH3UJPVnEhgBO4eUlNaHuQBDo_OKH17ChG5-Ybj2T4,9895
+ai_edge_torch/generative/layers/model_config.py,sha256=UH7vQPUCLZi6NRALiA28tYhoG3O2XW_P_QtjTj2r0Ts,10808
 ai_edge_torch/generative/layers/normalization.py,sha256=ijwCpi22NLX-Sygwy5sK9l9WjGvbPIhZvVwoBAonWAo,7014
 ai_edge_torch/generative/layers/normalization_test.py,sha256=zwurZly-TgFxdgVVdpzu9vCpcLbd5RYt_gKg9Lfg1jI,2248
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIrILJ7blceAcxrNqX1ZCN0ECKG1gz-bV8,2655
@@ -215,6 +215,7 @@ ai_edge_torch/generative/utilities/model_builder.py,sha256=tBfOcsI_NcneggHqkCSyd
 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/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=ETO2ShU5KXG7MLP8eVOWuzuRLCUtapafYHcZ6TZHIkw,13061
@@ -267,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.dev20250528.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.6.0.dev20250528.dist-info/METADATA,sha256=4Zdcdi5qzXkx7bVEuHMUdtsdN0UljrDl-8ud8-Q1hQQ,2074
-ai_edge_torch_nightly-0.6.0.dev20250528.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.6.0.dev20250528.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.6.0.dev20250528.dist-info/RECORD,,
+ai_edge_torch_nightly-0.6.0.dev20250529.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.6.0.dev20250529.dist-info/METADATA,sha256=uBCPe1_F0-0_ZyaQRalg_9Qs_iTHnpWr8p4jr-DCLMk,2074
+ai_edge_torch_nightly-0.6.0.dev20250529.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.6.0.dev20250529.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.6.0.dev20250529.dist-info/RECORD,,