ai-edge-torch-nightly 0.6.0.dev20250619__py3-none-any.whl → 0.6.0.dev20250620__py3-none-any.whl

ai_edge_torch/generative/examples/gemma3/image_encoder.py

@@ -24,26 +24,27 @@ import torch.nn.functional as F
 
 
 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",
+    ff_up_proj="siglip_vision_model.encoder_blocks.{}.mlp.fc1",
+    ff_down_proj="siglip_vision_model.encoder_blocks.{}.mlp.fc2",
     attn_query_proj=(
-        "vision_tower.vision_model.encoder.layers.{}.self_attn.q_proj"
+        "siglip_vision_model.encoder_blocks.{}.self_attn.q_proj"
     ),
     attn_key_proj=(
-        "vision_tower.vision_model.encoder.layers.{}.self_attn.k_proj"
+        "siglip_vision_model.encoder_blocks.{}.self_attn.k_proj"
     ),
     attn_value_proj=(
-        "vision_tower.vision_model.encoder.layers.{}.self_attn.v_proj"
+        "siglip_vision_model.encoder_blocks.{}.self_attn.v_proj"
     ),
     attn_output_proj=(
-        "vision_tower.vision_model.encoder.layers.{}.self_attn.out_proj"
+        "siglip_vision_model.encoder_blocks.{}.self_attn.o_proj"
     ),
-    pre_attn_norm="vision_tower.vision_model.encoder.layers.{}.layer_norm1",
-    embedding="vision_tower.vision_model.embeddings.patch_embedding",
+    pre_attn_norm="siglip_vision_model.encoder_blocks.{}.layer_norm1",
+    pre_ff_norm="siglip_vision_model.encoder_blocks.{}.layer_norm2",
+    embedding="siglip_vision_model.patch_embedding",
     embedding_position=(
-        "vision_tower.vision_model.embeddings.position_embedding.weight"
+        "siglip_vision_model.position_embedding.weight"
     ),
-    final_norm="vision_tower.vision_model.post_layernorm",
+    final_norm="siglip_vision_model.final_norm",
 )
 
 

ai_edge_torch/generative/layers/attention_utils.py

@@ -61,6 +61,7 @@ def build_causal_mask_cache(
     size: int,
     dtype: torch.dtype = torch.float32,
     device: torch.device = None,
+    mask_value: float = float('-inf'),
 ) -> torch.Tensor:
   """Build a cache for causal attention mask.
 
@@ -70,6 +71,8 @@ def build_causal_mask_cache(
         torch.float32.
       device (torch.device, optional): Output tensor's data type. Defaults to
         None in which case "cpu" is used.
+      mask_value (float, optional): The value to set the mask to. Defaults to
+        float('-inf').
 
   Returns:
       torch.Tensor: Causal attention mask.
@@ -77,7 +80,7 @@ def build_causal_mask_cache(
 
   if device is None:
     device = torch.device('cpu')
-  mask = torch.full((size, size), float('-inf'), dtype=dtype, device=device)
+  mask = torch.full((size, size), mask_value, dtype=dtype, device=device)
   return torch.triu(mask, diagonal=1).unsqueeze(0).unsqueeze(0)
 
 
@@ -86,6 +89,7 @@ def build_sliding_window_mask_cache(
     window_size: int,
     dtype: torch.dtype = torch.float32,
     device: torch.device = None,
+    mask_value: float = float('-inf'),
 ) -> torch.Tensor:
   """Build a cache for a sliding window mask.
 
@@ -96,18 +100,20 @@ def build_sliding_window_mask_cache(
         torch.float32.
       device (torch.device, optional): Output tensor's data type. Defaults to
         None in which case "cpu" is used.
+      mask_value (float, optional): The value to set the mask to. Defaults to
+        float('-inf').
 
   Returns:
       torch.Tensor: Causal attention mask.
   """
 
-  mask = build_causal_mask_cache(size, dtype, device)
+  mask = build_causal_mask_cache(size, dtype, device, mask_value)
   all_ones = torch.ones_like(mask)
   window_size = min(size, window_size)
   sliding_mask = torch.triu(all_ones, -1 * window_size + 1) * torch.tril(
       all_ones, window_size - 1
   )
-  return torch.where(sliding_mask == 1, mask, float('-inf'))
+  return torch.where(sliding_mask == 1, mask, mask_value)
 
 
 def relative_position_bucket(

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.dev20250619"
+__version__ = "0.6.0.dev20250620"

{ai_edge_torch_nightly-0.6.0.dev20250619.dist-info → ai_edge_torch_nightly-0.6.0.dev20250620.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.6.0.dev20250619
+Version: 0.6.0.dev20250620
 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.dev20250619.dist-info → ai_edge_torch_nightly-0.6.0.dev20250620.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=aCO6sA_1IPQGd5f8Ya-ce4ZKJE1EEt2BkypXJLQ3qvI,806
+ai_edge_torch/version.py,sha256=L-kCkN9vMaBYCuarY-Y8kGgttAEHZdyPWaupqKVMLiA,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
@@ -74,7 +74,7 @@ ai_edge_torch/generative/examples/gemma3/__init__.py,sha256=JaAnrFoXTl3RJX97Xspk
 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/image_encoder.py,sha256=OCMIAQfNmPR4uQUAtlYL6j4xkG0dw2Ays4-lnThcWqQ,5110
 ai_edge_torch/generative/examples/gemma3/verify_gemma3.py,sha256=v8oNXFICmVOtQxfO7IhZ8GnbvotEkDi9lzYHjoQyOso,2464
 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
@@ -171,7 +171,7 @@ ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=myGjal5A
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/attention.py,sha256=RaXENRRQo1MsLdt3U8h3kYTCmd6imHQ-aCXtmPXCh_o,13911
 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.py,sha256=2qfg7Tzk9ikKph5w3geOHC1I6EyOCdDsWXMr7F7IOZM,7630
 ai_edge_torch/generative/layers/attention_utils_test.py,sha256=22gQ1gcRPkwqFG3_p82GZfRKVE3udEssSy58wNOqv0w,2431
 ai_edge_torch/generative/layers/builder.py,sha256=2bUgkyowDkDznkF8XaHyZs4nowHr1QEHYLM7pMaFmIk,4921
 ai_edge_torch/generative/layers/einsum.py,sha256=EsZSWNVWUs0-1plp4TBnhP4ZhaRDBa2VlDO6hWpUAqU,1288
@@ -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.dev20250619.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.6.0.dev20250619.dist-info/METADATA,sha256=_nUnboHwt2qZ0ejpyggXZf5q5iuNSNUUboXq6e8uQGw,2074
-ai_edge_torch_nightly-0.6.0.dev20250619.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.6.0.dev20250619.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.6.0.dev20250619.dist-info/RECORD,,
+ai_edge_torch_nightly-0.6.0.dev20250620.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.6.0.dev20250620.dist-info/METADATA,sha256=W4cZLDBaywmznVjn7haIlLen5cKvXF7VYVIapnx4h0E,2074
+ai_edge_torch_nightly-0.6.0.dev20250620.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.6.0.dev20250620.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.6.0.dev20250620.dist-info/RECORD,,