ai-edge-torch-nightly 0.5.0.dev20250427__py3-none-any.whl → 0.5.0.dev20250429__py3-none-any.whl

ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py

@@ -21,7 +21,6 @@ from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("gemma-2b")
-ExportConfig = export_config.ExportConfig
 
 
 def main(_):
@@ -35,7 +34,7 @@ def main(_):
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=ExportConfig(),
+      export_config=export_config.get_from_flags(),
   )
 
 

ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py

@@ -21,7 +21,6 @@ from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("gemma2-2b")
-ExportConfig = export_config.ExportConfig
 
 
 def main(_):
@@ -35,7 +34,7 @@ def main(_):
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=ExportConfig(),
+      export_config=export_config.get_from_flags(),
   )
 
 

ai_edge_torch/generative/examples/gemma/gemma1.py

@@ -65,6 +65,7 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       type=cfg.NormalizationType.RMS_NORM,
       epsilon=1e-6,
       zero_centered=True,
+      enable_hlfb=True,
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -236,6 +236,7 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       type=cfg.NormalizationType.RMS_NORM,
       epsilon=1e-6,
       zero_centered=True,
+      enable_hlfb=True,
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.GATED,
@@ -314,5 +315,5 @@ def build_2b_model(checkpoint_path: str, **kwargs) -> nn.Module:
           tensor_names=tensor_names,
           model_class=Gemma2,
       )
-    except KeyError as ke:
+    except KeyError as _:
       continue

ai_edge_torch/generative/examples/gemma/verify_gemma2.py

@@ -18,6 +18,7 @@
 from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.gemma import verify_util
+import ai_edge_torch.generative.layers.kv_cache as kv_utils
 import kagglehub
 
 
@@ -31,12 +32,27 @@ _MAX_NEW_TOKENS = flags.DEFINE_integer(
     30,
     "The maximum size of the generated tokens.",
 )
-
+_MASK_AS_INPUT = flags.DEFINE_bool(
+    "mask_as_input",
+    True,
+    "Pass the causal self attention mask to the model.",
+)
+_TRANSPOSE_KV_CACHE = flags.DEFINE_bool(
+    "transpose_kv_cache",
+    True,
+    "Transpose the KV cache to reduce memory usage.",
+)
 
 def main(_):
   checkpoint = kagglehub.model_download("google/gemma-2/pyTorch/gemma-2-2b-it")
 
-  verify_util.verify_gemma2(checkpoint, _PROMPTS.value, _MAX_NEW_TOKENS.value)
+  verify_util.verify_gemma2(
+      checkpoint,
+      _PROMPTS.value,
+      _MAX_NEW_TOKENS.value,
+      _MASK_AS_INPUT.value,
+      kv_utils.KV_LAYOUT_TRANSPOSED if _TRANSPOSE_KV_CACHE.value else kv_utils.KV_LAYOUT_DEFAULT,
+  )
 
 
 if __name__ == "__main__":

ai_edge_torch/generative/examples/gemma/verify_util.py

@@ -21,6 +21,7 @@ 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
+import ai_edge_torch.generative.layers.kv_cache as kv_utils
 from ai_edge_torch.generative.utilities import verifier
 from gemma import config as gemma_config
 from gemma import model as gemma_model
@@ -108,6 +109,8 @@ def verify_reauthored_gemma_model(
     weight_filename: str = "model.ckpt",
     tokenizer_filename: str = "tokenizer.model",
     max_new_tokens: int = 20,
+    mask_as_input: bool = False,
+    kv_layout: kv_utils.KVLayout = kv_utils.KV_LAYOUT_DEFAULT,
     rtol: float = 1e-05,
     atol: float = 1e-05,
 ) -> bool:
@@ -126,7 +129,11 @@ def verify_reauthored_gemma_model(
 
   return verifier.verify_reauthored_model(
       original_model=GemmaWrapper(original_model),
-      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      reauthored_model=verifier.ReauthoredModelWrapper(
+          reauthored_model,
+          mask_as_input=mask_as_input,
+          kv_layout=kv_layout,
+      ),
       tokenizer=GemmaTokenizerWrapper(original_model.tokenizer),
       generate_prompts=generate_prompts,
       max_new_tokens=max_new_tokens,
@@ -137,7 +144,11 @@ def verify_reauthored_gemma_model(
 
 
 def verify_gemma2(
-    gemma2_model_path: str, prompts: List[str], max_new_tokens: int
+    gemma2_model_path: str,
+    prompts: List[str],
+    max_new_tokens: int,
+    mask_as_input: bool = False,
+    kv_layout: kv_utils.KVLayout = kv_utils.KV_LAYOUT_DEFAULT,
 ) -> bool:
   """Verifies the reauthored Gemma2 model.
 
@@ -153,5 +164,7 @@ def verify_gemma2(
       generate_prompts=prompts,
       forward_input_ids=[[2, 651, 9456, 576, 573, 3520, 3858, 603, 235248]],
       max_new_tokens=max_new_tokens,
+      mask_as_input=mask_as_input,
+      kv_layout=kv_layout,
       atol=1e-04,
   )

ai_edge_torch/generative/examples/gemma3/verify_util.py

@@ -22,7 +22,7 @@ from typing import List, Optional, Tuple
 from ai_edge_torch.generative.examples.gemma3 import gemma3
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
-from ai_edge_torch.generative.utilities.experimental import verifier
+from ai_edge_torch.generative.utilities import verifier
 from gemma import config as gemma_config
 from gemma import model as gemma_model
 import torch
@@ -92,10 +92,12 @@ 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_kv_cache(self):
-    """Returns an initialized KV cache."""
     return kv_utils.KVCache.from_model_config(
-        self.model.model.config, kv_layout=kv_utils.KV_LAYOUT_TRANSPOSED
+        self.model.model.config, kv_layout=self.kv_layout
     )
 
   def forward(

ai_edge_torch/generative/layers/attention_test.py

@@ -0,0 +1,153 @@
+# 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.
+# ==============================================================================
+
+from ai_edge_torch.generative.layers import attention
+from ai_edge_torch.generative.layers import model_config as cfg
+import torch
+
+from absl.testing import absltest as googletest
+from absl.testing import parameterized
+
+
+class AttentionTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="local_causal_self_attention",
+          attn_type=cfg.AttentionType.LOCAL_SLIDING,
+          expected_shape=(1, 10, 16),
+      ),
+      dict(
+          testcase_name="global_causal_self_attention",
+          attn_type=cfg.AttentionType.GLOBAL,
+          expected_shape=(1, 10, 16),
+      ),
+  )
+  def test_causal_self_attention(
+      self, attn_type: cfg.AttentionType, expected_shape: tuple[int, ...]
+  ):
+    norm_config = cfg.NormalizationConfig(
+        type=cfg.NormalizationType.RMS_NORM,
+        epsilon=1e-6,
+        zero_centered=True,
+        enable_hlfb=True,
+    )
+    attn_config = cfg.AttentionConfig(
+        num_heads=2,
+        head_dim=8,
+        num_query_groups=1,
+        rotary_base=100,
+        rotary_percentage=1.0,
+        qkv_transpose_before_split=True,
+        query_norm_config=norm_config,
+        key_norm_config=norm_config,
+        logit_softcap=None,
+        sliding_window_size=16,
+        attn_type=attn_type,
+    )
+    self_atten = attention.CausalSelfAttention(
+        dim=16,
+        config=attn_config,
+        enable_hlfb=True,
+    )
+    x = torch.randn(1, 10, 16)
+    attn_mask = torch.ones((1, 1, 10, 10), dtype=torch.float32)
+    out = self_atten(x, rope=None, mask=attn_mask)
+    self.assertEqual(out.shape, expected_shape)
+
+  def test_cross_attention(self):
+    norm_config = cfg.NormalizationConfig(
+        type=cfg.NormalizationType.RMS_NORM,
+        epsilon=1e-6,
+        zero_centered=True,
+        enable_hlfb=True,
+    )
+    attn_config = cfg.AttentionConfig(
+        num_heads=2,
+        head_dim=8,
+        num_query_groups=1,
+        rotary_base=100,
+        rotary_percentage=1.0,
+        qkv_transpose_before_split=True,
+        query_norm_config=norm_config,
+        key_norm_config=norm_config,
+        logit_softcap=None,
+        sliding_window_size=16,
+        attn_type=cfg.AttentionType.GLOBAL,
+    )
+    cross_atten = attention.CrossAttention(
+        query_dim=16,
+        cross_dim=16,
+        hidden_dim=16,
+        output_dim=16,
+        config=attn_config,
+        enable_hlfb=True,
+    )
+    x = torch.randn(1, 10, 16)
+    y = torch.randn(1, 10, 16)
+    out = cross_atten(x, y, rope=None)
+    self.assertEqual(out.shape, (1, 10, 16))
+
+  def test_transformer_block(self):
+    norm_config = cfg.NormalizationConfig(
+        type=cfg.NormalizationType.RMS_NORM,
+        epsilon=1e-6,
+        zero_centered=True,
+        enable_hlfb=True,
+    )
+    attn_config = cfg.AttentionConfig(
+        num_heads=2,
+        head_dim=8,
+        num_query_groups=1,
+        rotary_base=100,
+        rotary_percentage=1.0,
+        qkv_transpose_before_split=True,
+        query_norm_config=norm_config,
+        key_norm_config=norm_config,
+        logit_softcap=None,
+        sliding_window_size=16,
+        attn_type=cfg.AttentionType.GLOBAL,
+    )
+    ff_config = cfg.FeedForwardConfig(
+        type=cfg.FeedForwardType.GATED,
+        activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
+        intermediate_size=32,
+    )
+    block_config = cfg.TransformerBlockConfig(
+        attn_config=attn_config,
+        ff_config=ff_config,
+        post_attention_norm_config=norm_config,
+        parallel_residual=True,
+    )
+    model_config = cfg.ModelConfig(
+        vocab_size=100,
+        embedding_dim=16,
+        enable_hlfb=True,
+        num_layers=1,
+        max_seq_len=10,
+        block_configs=[block_config],
+    )
+    transformer_block = attention.TransformerBlock(
+        config=block_config,
+        model_config=model_config,
+    )
+    x = torch.randn(1, 10, 16)
+    attn_mask = torch.ones((1, 1, 10, 10), dtype=torch.float32)
+    out = transformer_block(x, rope=None, mask=attn_mask)
+    self.assertEqual(out.shape, (1, 10, 16))
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_torch/generative/layers/attention_utils_test.py

@@ -0,0 +1,64 @@
+# 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.
+# ==============================================================================
+
+from ai_edge_torch.generative.layers import attention_utils
+import torch
+
+from absl.testing import absltest as googletest
+
+
+class AttentionUtilsTest(googletest.TestCase):
+
+  def test_build_causal_mask_cache(self):
+    mask = attention_utils.build_causal_mask_cache(3)
+    self.assertEqual(mask.shape, (1, 1, 3, 3))
+    self.assertEqual(mask[0, 0, 0, 0], 0)
+    self.assertEqual(mask[0, 0, 0, 1], float("-inf"))
+    self.assertEqual(mask[0, 0, 0, 2], float("-inf"))
+    self.assertEqual(mask[0, 0, 1, 0], 0)
+    self.assertEqual(mask[0, 0, 1, 1], 0)
+    self.assertEqual(mask[0, 0, 1, 2], float("-inf"))
+    self.assertEqual(mask[0, 0, 2, 0], 0)
+    self.assertEqual(mask[0, 0, 2, 1], 0)
+    self.assertEqual(mask[0, 0, 2, 2], 0)
+
+  def test_build_sliding_window_mask_cache(self):
+    mask = attention_utils.build_sliding_window_mask_cache(3, 2)
+    self.assertEqual(mask.shape, (1, 1, 3, 3))
+    self.assertEqual(mask[0, 0, 0, 0], 0)
+    self.assertEqual(mask[0, 0, 0, 1], float("-inf"))
+    self.assertEqual(mask[0, 0, 0, 2], float("-inf"))
+    self.assertEqual(mask[0, 0, 1, 0], 0)
+    self.assertEqual(mask[0, 0, 1, 1], 0)
+    self.assertEqual(mask[0, 0, 1, 2], float("-inf"))
+    self.assertEqual(mask[0, 0, 2, 0], float("-inf"))
+    self.assertEqual(mask[0, 0, 2, 1], 0)
+    self.assertEqual(mask[0, 0, 2, 2], 0)
+
+  def test_build_relative_position_buckets(self):
+    buckets = attention_utils.build_relative_position_buckets(
+        query_length=3, key_length=3, bidirectional=True, num_buckets=4
+    )
+    print(buckets)
+    self.assertEqual(buckets.shape, (1, 1, 3, 3))
+    self.assertTrue(
+        torch.equal(
+            buckets, torch.tensor([[[[0, 3, 3], [1, 0, 3], [1, 1, 0]]]])
+        )
+    )
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_torch/generative/layers/kv_cache.py

@@ -18,7 +18,7 @@
 import dataclasses
 from typing import Any, List, Tuple
 
-from ai_edge_torch.generative.custom_ops.dynamic_update_slice import dynamic_update_slice
+import ai_edge_torch.generative.custom_ops.dynamic_update_slice as dus_utils
 from ai_edge_torch.generative.layers import model_config
 from ai_edge_torch.generative.utilities import types
 import torch
@@ -266,8 +266,78 @@ def _update_kv_impl(
   k_slice_indices = _get_slice_indices(input_pos)
   v_slice_indices = _get_slice_indices(input_pos)
 
-  k = dynamic_update_slice(cache.k_cache, k_slice, k_slice_indices)
-  v = dynamic_update_slice(cache.v_cache, v_slice, v_slice_indices)
+  k = dus_utils.dynamic_update_slice(cache.k_cache, k_slice, k_slice_indices)
+  v = dus_utils.dynamic_update_slice(cache.v_cache, v_slice, v_slice_indices)
 
   updated_cache = KVCacheEntry(k, v, cache.kv_layout)
   return updated_cache
+
+
+def update_transposed(
+    cache: KVCacheEntry,
+    input_pos: torch.Tensor,
+    k_slice: torch.Tensor,
+    v_slice: torch.Tensor,
+) -> KVCacheEntry:
+  """Out of place update of Cache buffer.
+
+  Args:
+      cache (KVCacheEntry): The original cache buffer.
+      input_pos (torch.Tensor): The update slice positions.
+      k_slice (torch.Tensor): The K slice to be updated in the new cache.
+      v_slice (torch.Tensor): The V slice to be updated in the new cache.
+
+  Returns:
+      KVCacheEntry: The updated KVCacheBase entry based on the passed
+      inputs.
+  """
+  assert (
+      cache.kv_layout == KV_LAYOUT_TRANSPOSED
+  ), "KV entry must have transposed layout."
+  return _update_kv_impl_transposed(cache, input_pos, k_slice, v_slice)
+
+
+def _get_slice_indices_transposed(
+    positions: torch.Tensor, cache_dim: int, ts_idx: int
+) -> torch.Tensor:
+  """Returns the slice indices."""
+  positions = positions.float()[0].reshape(
+      1,
+  )
+
+  zeros = torch.zeros((1,), dtype=torch.float32)
+  indices = []
+  for i in range(cache_dim):
+    if i == ts_idx:
+      indices.append(positions)
+    else:
+      indices.append(zeros)
+  slice_indices = torch.cat(indices, dim=0)
+  slice_indices = slice_indices.int()
+  return slice_indices
+
+
+def _update_kv_impl_transposed(
+    cache: KVCacheEntry,
+    input_pos: torch.Tensor,
+    k_slice: torch.Tensor,
+    v_slice: torch.Tensor,
+) -> KVCacheEntry:
+  """Updates the cache buffer with High Level Function Boundary annotation."""
+  cache_dim = 4
+  k_ts_idx = 2
+  v_ts_idx = 3
+  positions = input_pos.clone()
+  k_slice_indices = _get_slice_indices_transposed(
+      positions, cache_dim, k_ts_idx
+  )
+  v_slice_indices = _get_slice_indices_transposed(
+      positions, cache_dim, v_ts_idx
+  )
+  k = dus_utils.dynamic_update_slice(
+      cache.k_cache, k_slice, [x for x in k_slice_indices]
+  )
+  v = dus_utils.dynamic_update_slice(
+      cache.v_cache, v_slice, [x for x in v_slice_indices]
+  )
+  return KVCacheEntry(k, v, cache.kv_layout)

ai_edge_torch/generative/layers/sdpa_with_kv_update.py

@@ -19,7 +19,6 @@ from typing import Tuple
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.generative.layers import scaled_dot_product_attention as sdpa
-from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils_experimental
 import ai_edge_torch.generative.layers.model_config as cfg
 import torch
 
@@ -68,7 +67,7 @@ def _sdpa_with_kv_update_transposed(
       1, -1, config.head_dim, seq_len
   )  # 1, bk, h, s
 
-  kv = kv_utils_experimental.update(kv, input_pos, key, value)
+  kv = kv_utils.update_transposed(kv, input_pos, key, value)
   key, value = kv.k_cache, kv.v_cache
 
   sdpa_out = sdpa.scaled_dot_product_attention_transposed(

ai_edge_torch/generative/utilities/converter.py

@@ -57,7 +57,7 @@ def define_conversion_flags(model_name: str):
   )
   flags.DEFINE_string(
       'output_name_prefix',
-      model_name,
+      f'{model_name}',
       'The prefix of the output tflite model name.',
   )
   flags.DEFINE_multi_integer(

ai_edge_torch/generative/utilities/export_config.py

@@ -50,8 +50,7 @@ def _build_mask(mask_len, kv_cache_max_len) -> torch.Tensor:
   mask = torch.full(
       (mask_len, kv_cache_max_len), float('-inf'), dtype=torch.float32
   )
-  mask = torch.triu(mask, diagonal=1).unsqueeze(0).unsqueeze(0)
-  return mask
+  return torch.triu(mask, diagonal=1).unsqueeze(0).unsqueeze(0)
 
 
 def get_from_flags() -> ExportConfig:
@@ -62,6 +61,13 @@ def get_from_flags() -> ExportConfig:
     export_config.prefill_mask = _build_mask(
         flags.FLAGS.prefill_seq_lens, flags.FLAGS.kv_cache_max_len
     )
+    # Note that the decode mask is not a correct causal mask, but it is okay
+    # for the conversion purpose because only the shape matters in conversion.
+    # A correct causal mask of decode for a given token position of decode, it
+    # should be built like:
+    #
+    #  torch.triu(mask, diagonal=decode_position).unsqueeze(0).unsqueeze(0)
+    #
     export_config.decode_mask = _build_mask(1, flags.FLAGS.kv_cache_max_len)
 
   if flags.FLAGS.transpose_kv_cache:

ai_edge_torch/generative/utilities/verifier.py

@@ -85,14 +85,35 @@ class ModelWrapper(torch.nn.Module):
 class ReauthoredModelWrapper(ModelWrapper):
   """A wrapper for the model reauthored with ai_edge_torch Generative API."""
 
+  def __init__(
+      self,
+      model: torch.nn.Module,
+      mask_as_input: bool = False,
+      kv_layout: kv_utils.KVLayout = kv_utils.KV_LAYOUT_DEFAULT,
+  ):
+    """Wraps a reauthored model with some options."""
+    super().__init__(model)
+    self.mask_as_input = mask_as_input
+    self.kv_layout = kv_layout
+
   def _init_kv_cache(self):
     """Returns an initialized KV cache."""
-    return kv_utils.KVCache.from_model_config(self.model.config)
+    return kv_utils.KVCache.from_model_config(
+        self.model.config, kv_layout=self.kv_layout
+    )
 
   def _get_extra_args_for_forward(self) -> dict[str, Any]:
     """Returns extra arguments for the forward() method."""
     return {}
 
+  def _build_mask(self, input_pos: torch.Tensor) -> torch.Tensor:
+    """Builds a mask for the model."""
+    kv_cache_max_len = self.model.config.kv_cache_max_len
+    mask = torch.full(
+        (len(input_pos), kv_cache_max_len), float("-inf"), dtype=torch.float32
+    )
+    return torch.triu(mask, diagonal=input_pos[0] + 1).unsqueeze(0).unsqueeze(0)
+
   def _forward_with_kv_cache(
       self,
       tokens: torch.Tensor,
@@ -119,6 +140,8 @@ class ReauthoredModelWrapper(ModelWrapper):
       extra_args["export_config"] = self.export_config
     if pixel_values is not None:
       extra_args["pixel_values"] = pixel_values
+    if self.mask_as_input:
+      extra_args["mask"] = self._build_mask(input_pos)
     output = self.model.forward(tokens, input_pos, kv_cache, **extra_args)
     return output["logits"], output["kv_cache"]
 

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.5.0.dev20250427"
+__version__ = "0.5.0.dev20250429"

{ai_edge_torch_nightly-0.5.0.dev20250427.dist-info → ai_edge_torch_nightly-0.5.0.dev20250429.dist-info}/METADATA

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

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,120
 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=RhNMNIs4sG78K3SOLk6zxuILeS_S2vhG7FJJOrV4cLM,706
+ai_edge_torch/version.py,sha256=I820JmIf90_QKTKyhmQGVjX9U-WMGUVEo9_N-Q_aQuk,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=QVugYVfbyaeBgSKKbhFzHG5oXA7t3M-40JcpcdSu6W8,5436
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -57,20 +57,20 @@ ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py,sha256=l0OrPGmX8
 ai_edge_torch/generative/examples/deepseek/deepseek.py,sha256=yhS_i2kR0GJWpWciCt4p9Z9nHYh6A5uJ8Ycy2ebFN9w,2909
 ai_edge_torch/generative/examples/deepseek/verify.py,sha256=iYldze-pvZGvPkkqr6zA7EmitPnH9sXkzjNVx353IcE,2403
 ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=tSEtGeS-Ndcc_cTm7c4CT4FqRiwrHedEv1oJk4Y_zYU,1552
-ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=k2BUkf2cciItc3gFAyzWqcWZhlVFrD3TVikTmLXq04c,1553
-ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=w8oWYibZzvEvCDyp39EYyAWmjgJljhzdYPyFCfAWxZA,3497
-ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=SyyRae8vWLn0WxduxtahzVRbdSq4T2k5-7t8PfCR_k8,11534
+ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=RRilUl2Ui08R9gy1Ua0jnaXNCrIJJb-oztgP62G3mX4,1526
+ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=7IlF-4NEfZAzIfkOUHR-HeCSLSUGEu7wnO52UtERCa4,1527
+ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=6ImjTzJcq6JoKz2Z-z8pjv5BsRu5nUeEsTK3IPs3xgI,3521
+ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=JQLLiHNVBM9jOrZqUF0EmgAwtDD0yTRlmIbLaWM7qTg,11557
 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
+ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=jhiyinOqPt5ZZjEadDRZt_wY5fiLSCpMo54PcxFaL_Q,1789
+ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=n7f2nF6Lin_tDvPs0JVldsuaBzo7pAwi5YAHAhlIxQg,6139
 ai_edge_torch/generative/examples/gemma3/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
 ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=MjkQDVynaw9C5z9ODzKfb85xW5JfxHUWBJ_Aco05FHo,1760
 ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=eXWE5CSX0KeUMsPevgsYOfvyajl9F1RFF4DCWhHcYPA,15646
 ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=GACDBI_MsFowR8A3wAWrpzradPYe-AUgB9ZjXaVBG-s,6485
 ai_edge_torch/generative/examples/gemma3/image_encoder.py,sha256=uRoLoBWzFtQz5wFZfPCxbkvZsgPAqSkUUsV3977GbYc,5184
 ai_edge_torch/generative/examples/gemma3/verify_gemma3.py,sha256=v8oNXFICmVOtQxfO7IhZ8GnbvotEkDi9lzYHjoQyOso,2464
-ai_edge_torch/generative/examples/gemma3/verify_util.py,sha256=nEv0qQ0l6gSXKxP5mNwkd2lRGxpFfD4e7FNV3V76zhw,8915
+ai_edge_torch/generative/examples/gemma3/verify_util.py,sha256=KnE9ME3mrpQkAxFlBOJLsqcQkjsdDL1ClNhJahX5K5I,8960
 ai_edge_torch/generative/examples/hammer/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/hammer/convert_to_tflite.py,sha256=946mchDmvUhMsv1kzslp4LHtCIuHn4qjimHYQ-XnxMo,2962
 ai_edge_torch/generative/examples/hammer/hammer.py,sha256=76INcjffvaNCQ02fzXcxJUW_6EKHs4sg3q1nDBbEpHE,3431
@@ -154,18 +154,18 @@ ai_edge_torch/generative/fx_passes/__init__.py,sha256=PFSMsA1vfBfrV9ssBCkYJNl8Hx
 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
 ai_edge_torch/generative/layers/attention.py,sha256=uK1ih2kxPZherwi-pGSm8B--NNWnQ8npEAfgcjMIkEY,12964
+ai_edge_torch/generative/layers/attention_test.py,sha256=9v8v96TLyFPdqxEylU1JOAeRFAp2s0YoDHZN83SFJJA,4764
 ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
+ai_edge_torch/generative/layers/attention_utils_test.py,sha256=22gQ1gcRPkwqFG3_p82GZfRKVE3udEssSy58wNOqv0w,2431
 ai_edge_torch/generative/layers/builder.py,sha256=LXGuSHIx6QZAzLFm7aJvlzoMPgQwbXLFchGEKYwOOUA,5090
 ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
-ai_edge_torch/generative/layers/kv_cache.py,sha256=dDeirtuo9AnlN1tYoLbFi_pKhIDmn35FQY1m6X28hSY,8468
+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=nLXvTkDAIHJQ0PTaWODF8oxJQoJ-K8D10cKR9229SAw,8355
 ai_edge_torch/generative/layers/normalization.py,sha256=MbwH-n80Fob5YvjBzdqDjBizMHLzSJGYRDdbD-rL5C0,6174
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIrILJ7blceAcxrNqX1ZCN0ECKG1gz-bV8,2655
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=efqqGRZPJ55hKn1MQJ-cXfrJD85uS1v7W_juyGyts58,5648
-ai_edge_torch/generative/layers/sdpa_with_kv_update.py,sha256=Hn8Zw-jiB9GH2uZ-yaRMcDdpmjECcW4uCy-YNH9zV8c,3693
-ai_edge_torch/generative/layers/experimental/__init__.py,sha256=nz-K0h8DfiATHzR6s1_bCw2akUmHWffU1bDRSkIzSqI,592
-ai_edge_torch/generative/layers/experimental/kv_cache.py,sha256=zgpFVftOfllvjh9-UEBSvUbm152SnQETn29rUMMMvAM,2978
+ai_edge_torch/generative/layers/sdpa_with_kv_update.py,sha256=iw7D_46CFe9iRvU0UumbkIoqWQEhDroxm9ABcK-CLlM,3600
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=ZteHZXK6HKyxYji49DQ46sA9aIy7U3Jnz0HZp6hfevY,28996
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
@@ -187,8 +187,8 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=vQWmpzMkJ2hP
 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=4RNNl7vk3WN_JG5EZajofiRSqtPnUNCYosxTacdEOto,10948
-ai_edge_torch/generative/utilities/export_config.py,sha256=maUVt0T5FsLpHO5H-BZ-O0FRBZO_ejKwGhPR9Qq8ViM,2490
+ai_edge_torch/generative/utilities/converter.py,sha256=8A1MvU8SbJQkn2SIhF-73TXbI_i6nrloCdkpw83P2xQ,10953
+ai_edge_torch/generative/utilities/export_config.py,sha256=yGkfdN8Qrp8b_K8e5H0qaYmDrg0Dx_eb75JLhOnlygQ,2827
 ai_edge_torch/generative/utilities/loader.py,sha256=7p__m2JryWphGlYOuRxdoT4id4_tWJEVOV7y2X4H-Ak,13737
 ai_edge_torch/generative/utilities/model_builder.py,sha256=ZYX1TxpFdj573du2QCyHJlFjx4q1m12R74fp4Gwl92A,6343
 ai_edge_torch/generative/utilities/moonshine_loader.py,sha256=_RpFabSqtGH5PHiP3_1f6QfO14qMADUxr_HGRlVDFB0,4891
@@ -196,7 +196,7 @@ ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=dqPD9qRXEWt
 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/types.py,sha256=gZI9hIPB3XAo4oecKIIoVDfiyibLaSNFhecPFx4VDTM,2913
-ai_edge_torch/generative/utilities/verifier.py,sha256=RSMQ8eda63VHM-5KmquKfogmTPyhGvGnqkoz9i4bppY,12270
+ai_edge_torch/generative/utilities/verifier.py,sha256=ETO2ShU5KXG7MLP8eVOWuzuRLCUtapafYHcZ6TZHIkw,13061
 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
@@ -246,8 +246,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.5.0.dev20250427.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250427.dist-info/METADATA,sha256=g53PvQrw8WP7McVXcoMYSEF9lmh7VWexPnfQLGOTVJg,2051
-ai_edge_torch_nightly-0.5.0.dev20250427.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250427.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250427.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250429.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250429.dist-info/METADATA,sha256=05nMBPcVBVJcZhDI9SzsjryW3d4vpeeH_9H07RaA-PI,2051
+ai_edge_torch_nightly-0.5.0.dev20250429.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250429.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250429.dist-info/RECORD,,

ai_edge_torch/generative/layers/experimental/__init__.py

@@ -1,14 +0,0 @@
-# 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.
-#

ai_edge_torch/generative/layers/experimental/kv_cache.py

@@ -1,90 +0,0 @@
-# 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.
-# ==============================================================================
-
-"""Utility functions for KV Cache.
-
-This is an experimental implementation and is subject to change at any time.
-"""
-
-from ai_edge_torch.generative.custom_ops import dynamic_update_slice as dus_utils
-from ai_edge_torch.generative.layers import kv_cache as kv_utils
-import torch
-
-
-def update(
-    cache: kv_utils.KVCacheEntry,
-    input_pos: torch.Tensor,
-    k_slice: torch.Tensor,
-    v_slice: torch.Tensor,
-) -> kv_utils.KVCacheEntry:
-  """Out of place update of Cache buffer.
-
-  Args:
-      cache (kv_utils.KVCacheEntry): The original cache buffer.
-      input_pos (torch.Tensor): The update slice positions.
-      k_slice (torch.Tensor): The K slice to be updated in the new cache.
-      v_slice (torch.Tensor): The V slice to be updated in the new cache.
-
-  Returns:
-      kv_utils.KVCacheEntry: The updated KVCacheBase entry based on the passed
-      inputs.
-  """
-  assert (
-      cache.kv_layout == kv_utils.KV_LAYOUT_TRANSPOSED
-  ), "KV entry must have transposed layout."
-  update_kv_cache = _update_kv_impl_transposed
-  return update_kv_cache(cache, input_pos, k_slice, v_slice)
-
-
-def _get_slice_indices(
-    positions: torch.Tensor, cache_dim: int, ts_idx: int
-) -> torch.Tensor:
-  """Returns the slice indices."""
-  positions = positions.float()[0].reshape(
-      1,
-  )
-
-  zeros = torch.zeros((1,), dtype=torch.float32)
-  indices = []
-  for i in range(cache_dim):
-    if i == ts_idx:
-      indices.append(positions)
-    else:
-      indices.append(zeros)
-  slice_indices = torch.cat(indices, dim=0)
-  slice_indices = slice_indices.int()
-  return slice_indices
-
-
-def _update_kv_impl_transposed(
-    cache: kv_utils.KVCacheEntry,
-    input_pos: torch.Tensor,
-    k_slice: torch.Tensor,
-    v_slice: torch.Tensor,
-) -> kv_utils.KVCacheEntry:
-  """Update the cache buffer with High Level Function Boundary annotation."""
-  cache_dim = 4
-  k_ts_idx = 2
-  v_ts_idx = 3
-  positions = input_pos.clone()
-  k_slice_indices = _get_slice_indices(positions, cache_dim, k_ts_idx)
-  v_slice_indices = _get_slice_indices(positions, cache_dim, v_ts_idx)
-  k = dus_utils.dynamic_update_slice(
-      cache.k_cache, k_slice, [x for x in k_slice_indices]
-  )
-  v = dus_utils.dynamic_update_slice(
-      cache.v_cache, v_slice, [x for x in v_slice_indices]
-  )
-  return kv_utils.KVCacheEntry(k, v, cache.kv_layout)