ai-edge-torch-nightly 0.5.0.dev20250502__py3-none-any.whl → 0.5.0.dev20250503__py3-none-any.whl

ai_edge_torch/generative/examples/gemma3/decoder.py

@@ -199,7 +199,11 @@ class Decoder(nn.Module):
     sliding_mask = torch.where(
         sliding_mask_bool,
         torch.zeros_like(sliding_mask_bool, dtype=torch.float),
-        torch.full_like(sliding_mask_bool, float("-inf"), dtype=torch.float),
+        torch.full_like(
+            sliding_mask_bool,
+            self.config.get_causal_mask_value(),
+            dtype=torch.float,
+        ),
     )
 
     return sliding_mask
@@ -215,7 +219,7 @@ class Decoder(nn.Module):
       mask = torch.logical_and(mask, pixel_mask)
     else:
       mask = torch.logical_or(mask, pixel_mask)
-    mask = torch.where(mask, 0, float("-inf"))
+    mask = torch.where(mask, 0, self.config.get_causal_mask_value())
     return mask
 
   def build_pixel_mask(self, image_indices: torch.Tensor):

ai_edge_torch/generative/examples/hammer/convert_to_tflite.py

@@ -17,15 +17,10 @@
 
 from absl import app
 from ai_edge_torch.generative.examples.hammer import hammer
-from ai_edge_torch.generative.layers import kv_cache
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import export_config as export_cfg
-import torch
-
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags('hammer')
-ExportConfig = export_cfg.ExportConfig
-
 
 _MODEL_SIZE = flags.DEFINE_enum(
     'model_size',
@@ -40,35 +35,6 @@ _BUILDER = {
 }
 
 
-def _create_mask(mask_len, kv_cache_max_len):
-  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
-
-
-def _create_export_config(
-    prefill_seq_lens: list[int], kv_cache_max_len: int
-) -> ExportConfig:
-  """Creates the export config for the model."""
-  export_config = ExportConfig()
-  if isinstance(prefill_seq_lens, list):
-    prefill_mask = [_create_mask(i, kv_cache_max_len) for i in prefill_seq_lens]
-  else:
-    prefill_mask = _create_mask(prefill_seq_lens, kv_cache_max_len)
-
-  export_config.prefill_mask = prefill_mask
-
-  decode_mask = torch.full(
-      (1, kv_cache_max_len), float('-inf'), dtype=torch.float32
-  )
-  decode_mask = torch.triu(decode_mask, diagonal=1).unsqueeze(0).unsqueeze(0)
-  export_config.decode_mask = decode_mask
-  export_config.kvcache_layout = kv_cache.KV_LAYOUT_TRANSPOSED
-  return export_config
-
-
 def main(_):
   pytorch_model = _BUILDER[_MODEL_SIZE.value](
       flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
@@ -80,11 +46,7 @@ def main(_):
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=_create_export_config(
-          flags.FLAGS.prefill_seq_lens, flags.FLAGS.kv_cache_max_len
-      )
-      if flags.FLAGS.transpose_kv_cache
-      else ExportConfig(),
+      export_config=export_config.get_from_flags(),
   )
 
 

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -75,7 +75,7 @@ class Decoder(model_builder.DecoderOnlyModel):
     if mask is None:
       embeds_len = input_embeds.shape[1]
       mask = torch.zeros(embeds_len, self.config.kv_cache_max)
-      mask[:, embeds_len:] = float("-inf")
+      mask[:, embeds_len:] = attn_config.causal_mask_value
 
     return self._forward_with_embeds(
         input_embeds,

ai_edge_torch/generative/examples/paligemma/decoder2.py

@@ -75,7 +75,7 @@ class Decoder2(gemma2.Gemma2):
       # By default, don't mask image embeds with a diagonal causal mask.
       embeds_len = input_embeds.shape[1]
       mask = torch.zeros(embeds_len, self.config.kv_cache_max)
-      mask[:, embeds_len:] = float("-inf")
+      mask[:, embeds_len:] = attn_config.causal_mask_value
 
     return self._forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache, export_config

ai_edge_torch/generative/layers/model_config.py

@@ -116,6 +116,8 @@ class AttentionConfig:
   attn_type: Optional[AttentionType] = None
   # The size of the sliding window used for local attention.
   sliding_window_size: Optional[int] = None
+  # The default causal mask value used by attention layer.
+  causal_mask_value: float = float("-inf")
 
 
 @dataclasses.dataclass
@@ -247,3 +249,7 @@ class ModelConfig:
           f"Index {idx} is out of range for layer configs: {self.block_configs}"
       )
     return self.block_configs[idx]
+
+  @property
+  def get_causal_mask_value(self) -> float:
+    return self.block_config(0).attn_config.causal_mask_value

ai_edge_torch/generative/layers/scaled_dot_product_attention.py

@@ -160,7 +160,7 @@ def scaled_dot_product_attention_transposed(
   Args:
     query: Query tensor, with shape [B, T, N, H].
     key: Key tensor, with shape [B, T, KV_LEN, H].
-    value: Value tensor, with shape [B, T, KV_LEN, H].
+    value: Value tensor, with shape [B, T, H, KV_LEN].
     head_size (int): head dimension.
     mask (torch.Tensor): the optional mask tensor.
     scale (float): the optional scale factor.

ai_edge_torch/generative/layers/scaled_dot_product_attention_test.py

@@ -0,0 +1,87 @@
+# 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 import odml_torch
+from ai_edge_torch.generative.layers import scaled_dot_product_attention
+import torch
+
+from absl.testing import absltest as googletest
+
+
+class ScaledDotProductAttentionTest(googletest.TestCase):
+
+  def test_scaled_dot_product_attention(self):
+    query = torch.randn(1, 16, 16, 128, dtype=torch.float32)
+    key = torch.randn(1, 16, 16, 128, dtype=torch.float32)
+    value = torch.randn(1, 16, 16, 128, dtype=torch.float32)
+    mask = torch.ones((1, 1, 1, 16), dtype=torch.float32)
+    output = scaled_dot_product_attention.scaled_dot_product_attention(
+        query, key, value, head_size=128, mask=mask, scale=1.0, softcap=10.0
+    )
+    self.assertEqual(output.shape, (1, 16, 16, 128))
+
+  def test_scaled_dot_product_attention_transposed(self):
+    query = torch.randn(1, 16, 16, 128, dtype=torch.float32)
+    key = torch.randn(1, 16, 16, 128, dtype=torch.float32)
+    value = torch.randn(1, 16, 128, 16, dtype=torch.float32)
+    mask = torch.ones((1, 1, 1, 16), dtype=torch.float32)
+    output = (
+        scaled_dot_product_attention.scaled_dot_product_attention_transposed(
+            query, key, value, head_size=128, mask=mask, scale=1.0, softcap=10.0
+        )
+    )
+    self.assertEqual(output.shape, (1, 16, 16, 128))
+
+  def test_scaled_dot_product_attention_with_hlfb(self):
+    query = torch.randn(1, 16, 16, 128, dtype=torch.float32)
+    key = torch.randn(1, 16, 16, 128, dtype=torch.float32)
+    value = torch.randn(1, 16, 16, 128, dtype=torch.float32)
+    mask = torch.ones((1, 1, 1, 16), dtype=torch.float32)
+    output = (
+        scaled_dot_product_attention.scaled_dot_product_attention_with_hlfb(
+            query, key, value, head_size=128, mask=mask, scale=1.0, softcap=10.0
+        )
+    )
+    self.assertEqual(output.shape, (1, 16, 16, 128))
+
+    def model_to_mlir(model, args):
+      ep = torch.export.export(model, args)
+      mlir = odml_torch.export.exported_program_to_mlir(ep)
+      return mlir.get_text()
+
+    class SDPAModule(torch.nn.Module):
+
+      def __init__(self):
+        super().__init__()
+
+      def forward(self, query, key, value, mask):
+        return (
+            scaled_dot_product_attention.scaled_dot_product_attention_with_hlfb(
+                query,
+                key,
+                value,
+                head_size=128,
+                mask=mask,
+                scale=1.0,
+                softcap=10.0,
+            )
+        )
+
+    ir_text = model_to_mlir(SDPAModule().eval(), (query, key, value, mask))
+    self.assertEqual(ir_text.count("stablehlo.custom_call @mark_tensor"), 5)
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_torch/generative/utilities/converter.py

@@ -95,6 +95,18 @@ def define_conversion_flags(model_name: str):
   return flags
 
 
+def _build_mask(mask_len, kv_cache_max_len, causal_mask_value) -> torch.Tensor:
+  if isinstance(mask_len, list):
+    return [
+        _build_mask(i, kv_cache_max_len, causal_mask_value) for i in mask_len
+    ]
+
+  mask = torch.full(
+      (mask_len, kv_cache_max_len), causal_mask_value, dtype=torch.float32
+  )
+  return torch.triu(mask, diagonal=1).unsqueeze(0).unsqueeze(0)
+
+
 def convert_to_tflite(
     pytorch_model: torch.nn.Module,
     output_path: str,
@@ -229,14 +241,15 @@ def _export_helper(
           torch.arange(0, seq_len + pixel_seq_len, dtype=torch.int)
       )
 
-  if export_config.prefill_mask is None:
-    prefill_masks = None
-  elif isinstance(export_config.prefill_mask, torch.Tensor):
-    prefill_masks = [export_config.prefill_mask]
-  elif isinstance(export_config.prefill_mask, list):
-    prefill_masks = export_config.prefill_mask
-  else:
-    raise ValueError('Prefill masks unrecognized.')
+  prefill_masks = None
+  if flags.FLAGS.mask_as_input:
+    prefill_masks = [
+        _build_mask(
+            flags.FLAGS.prefill_seq_lens,
+            flags.FLAGS.kv_cache_max_len,
+            config.get_causal_mask_value(),
+        )
+    ]
 
   if prefill_masks:
     assert len(prefill_masks) == len(prefill_seq_lens)
@@ -299,8 +312,17 @@ def _export_helper(
         'input_pos': decode_input_pos,
         'kv_cache': decode_kv,
     }
-    if export_config.decode_mask is not None:
-      sample_kwargs['mask'] = export_config.decode_mask
+    if flags.FLAGS.mask_as_input:
+      # 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)
+      #
+      sample_kwargs['mask'] = _build_mask(
+          1, flags.FLAGS.kv_cache_max_len, config.get_causal_mask_value()
+      )
     if lora is not None:
       sample_kwargs['lora'] = lora
 

ai_edge_torch/generative/utilities/export_config.py

@@ -33,6 +33,8 @@ class ExportConfig:
   # When False, only decode signatures will produce output.
   output_logits_on_prefill: bool = False
   # Attention masks given as inputs to the model.
+  # Note that `prefill_mask`, `decode_mask`, and `kvcache_cls` are deprecated
+  # and will be removed in a future version.
   prefill_mask: Optional[torch.Tensor | List[torch.Tensor]] = None
   decode_mask: Optional[torch.Tensor | List[torch.Tensor]] = None
   # The KV Cache layout for K and V buffers in attention.
@@ -43,33 +45,10 @@ class ExportConfig:
   decode_batch_size: int = 1
 
 
-def _build_mask(mask_len, kv_cache_max_len) -> torch.Tensor:
-  if isinstance(mask_len, list):
-    return [_build_mask(i, kv_cache_max_len) for i in mask_len]
-
-  mask = torch.full(
-      (mask_len, kv_cache_max_len), float('-inf'), dtype=torch.float32
-  )
-  return torch.triu(mask, diagonal=1).unsqueeze(0).unsqueeze(0)
-
-
 def get_from_flags() -> ExportConfig:
   """Builds an export config according to the commandline flags."""
   export_config = ExportConfig()
 
-  if flags.FLAGS.mask_as_input:
-    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:
     export_config.kvcache_layout = kv_utils.KV_LAYOUT_TRANSPOSED
 

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.5.0.dev20250502.dist-info → ai_edge_torch_nightly-0.5.0.dev20250503.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250502
+Version: 0.5.0.dev20250503
 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.dev20250502.dist-info → ai_edge_torch_nightly-0.5.0.dev20250503.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=xg9p8FPpu9FCdvOFGsZZCz9bNlzIBYx6sFmrLF_JO7s,706
+ai_edge_torch/version.py,sha256=rjhPV_Qh8FDlHQTy8wAJvuXSNGcntZerhf-8FTEjuWI,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
@@ -66,13 +66,13 @@ ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=jhiyinOqPt5ZZjEa
 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/decoder.py,sha256=fzLpuJO5JseQLA38Li-i9Xdnh9I4zdBWQEOeNbUEfjI,15737
 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=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/convert_to_tflite.py,sha256=9r8LXyaoBXYIIhhe1WQgEIjaxALQPE1dO2N6qopyWCk,1753
 ai_edge_torch/generative/examples/hammer/hammer.py,sha256=76INcjffvaNCQ02fzXcxJUW_6EKHs4sg3q1nDBbEpHE,3431
 ai_edge_torch/generative/examples/hammer/verify.py,sha256=MkzAGkbPy4LKRhyCDm1cw-9jUt4VUxLPdwK_25fCGSE,2705
 ai_edge_torch/generative/examples/llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -88,8 +88,8 @@ ai_edge_torch/generative/examples/openelm/openelm.py,sha256=sIJ8Ie1oxFrJM-1jvv2u
 ai_edge_torch/generative/examples/openelm/verify.py,sha256=4W26ZtPF5Cb9mpHYuRM4b2QB_4W76zf4WV36KzexVjs,2446
 ai_edge_torch/generative/examples/paligemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=7HHXkC-IIu7ieBvBI4RlXs_oITz7R8a6YVYQskAs_Uk,2023
-ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=G1dwtWp_v77AI3uyIY-8g6qRP2tRH3CIKjJTeYNqFPU,5511
-ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=Z-SKdb0dd8uWT1d-FRwFx5-tJEqpdrQwiIZnFRhOtVo,6060
+ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=-EYUZp55dfRY1E-N0Pr3b9i5c7Tt1XvYxvsRixguVS8,5527
+ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=WB8r-e_Crog1ItBq3Zse_nUG-foFyBcJsuEG26r_Ji8,6076
 ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=SvuR97sjkBtfkerH7Hu1UXB8kCFLpEATNbPfCbNAyfo,5614
 ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=CFIjOmrn4a4Udki7l3im0JR4zTC_NttnsIr9_qWjKTY,6110
 ai_edge_torch/generative/examples/paligemma/verify.py,sha256=zrCNz_QSQU6BbaFtx-J-MqxXWcNlsAlquaHpKodsyW4,5350
@@ -162,10 +162,11 @@ ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-m
 ai_edge_torch/generative/layers/feed_forward_test.py,sha256=8ZGy79BBpsyS6yKKDEKrDt249G5Mz-8VKWW7_WHx0u4,1655
 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/model_config.py,sha256=dRZUMa71ADaEllu7TfXUWTMHRCcMgvkFMYMzmeJi4G8,8576
 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/scaled_dot_product_attention.py,sha256=2_AgwENsaOgaxgiSqgoj0V0JzQ09dFtP_nBhX-lJK2g,5648
+ai_edge_torch/generative/layers/scaled_dot_product_attention_test.py,sha256=c6JBMQsq9XeMmR1XvGEIidNsoh-YIvichXo2LwVHgr4,3301
 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
@@ -188,8 +189,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=8A1MvU8SbJQkn2SIhF-73TXbI_i6nrloCdkpw83P2xQ,10953
-ai_edge_torch/generative/utilities/export_config.py,sha256=yGkfdN8Qrp8b_K8e5H0qaYmDrg0Dx_eb75JLhOnlygQ,2827
+ai_edge_torch/generative/utilities/converter.py,sha256=K1gZWPq5f3Z7f9USeJ_PphctO1dyYTNrWSJQ-cztgKA,11658
+ai_edge_torch/generative/utilities/export_config.py,sha256=5IvR3grlMd4mWO5c_Y4x9Fk1b1xa57MzlYNE8XUaN28,2049
 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
@@ -247,8 +248,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.dev20250502.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250502.dist-info/METADATA,sha256=31xPCLYZskB9H70ijwBN7fjINXcq9ljgnlAjieGcUDM,2051
-ai_edge_torch_nightly-0.5.0.dev20250502.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250502.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250502.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250503.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250503.dist-info/METADATA,sha256=RU3caJRTJFodq-s8HxE5j7uo74dScWYcYMMAtqJVsD4,2051
+ai_edge_torch_nightly-0.5.0.dev20250503.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250503.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250503.dist-info/RECORD,,