ai-edge-torch-nightly 0.4.0.dev20250402__py3-none-any.whl → 0.4.0.dev20250405__py3-none-any.whl

ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py

@@ -15,14 +15,14 @@
 
 """Example of converting AMD-Llama-135m model to multi-signature tflite model."""
 
-import os
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.amd_llama_135m import amd_llama_135m
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("amd-llama-135m")
+ExportConfig = export_config.ExportConfig
+
 
 def main(_):
   pytorch_model = amd_llama_135m.build_model(

ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py

@@ -15,16 +15,13 @@
 
 """Example of converting DeepSeek R1 distilled models to tflite model."""
 
-import os
-import pathlib
-
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.deepseek import deepseek
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("deepseek")
+ExportConfig = export_config.ExportConfig
 
 def main(_):
   pytorch_model = deepseek.build_model(

ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py

@@ -15,13 +15,14 @@
 
 """Example of converting a Gemma1 model to multi-signature tflite model."""
 
-import os
 from absl import app
 from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("gemma-2b")
+ExportConfig = export_config.ExportConfig
+
 
 def main(_):
   pytorch_model = gemma1.build_2b_model(

ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py

@@ -15,14 +15,14 @@
 
 """Example of converting a Gemma2 model to multi-signature tflite model."""
 
-import os
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.gemma import gemma2
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("gemma2-2b")
+ExportConfig = export_config.ExportConfig
+
 
 def main(_):
   pytorch_model = gemma2.build_2b_model(

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -23,6 +23,7 @@ from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
@@ -151,7 +152,7 @@ class Gemma2(nn.Module):
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
       mask: Optional[torch.Tensor] = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     _, seq_len = tokens.size()
     assert self.config.max_seq_len >= seq_len, (
@@ -184,7 +185,7 @@ class Gemma2(nn.Module):
       mask: torch.Tensor | List[torch.Tensor],
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     """Forwards the model with input embeddings."""
     assert len(self.transformer_blocks) == len(kv_cache.caches), (

ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py

@@ -15,16 +15,16 @@
 
 """Example of converting a Gemma3 model to multi-signature tflite model."""
 
-import os
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.gemma3 import gemma3
 from ai_edge_torch.generative.layers.experimental import kv_cache
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 import torch
 
 flags = converter.define_conversion_flags('gemma3-1b')
+ExportConfig = export_config.ExportConfig
+
 
 _MODEL_SIZE = flags.DEFINE_string(
     'model_size',

ai_edge_torch/generative/examples/gemma3/decoder.py

@@ -23,6 +23,7 @@ from ai_edge_torch.generative.layers.experimental import attention
 from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
@@ -244,7 +245,7 @@ class Decoder(nn.Module):
       input_embeds: Optional[torch.Tensor] = None,
       mask: Optional[torch.Tensor] = None,
       image_indices: Optional[torch.Tensor] = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCacheBase]:
 
     pixel_mask = None
@@ -288,7 +289,7 @@ class Decoder(nn.Module):
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCacheBase,
       pixel_mask: Optional[torch.Tensor] = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCacheBase]:
     """Forwards the model with input embeddings."""
     assert len(self.transformer_blocks) == len(kv_cache.caches), (

ai_edge_torch/generative/examples/gemma3/gemma3.py

@@ -24,7 +24,7 @@ from ai_edge_torch.generative.examples.gemma3 import image_encoder
 from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
-from ai_edge_torch.generative.utilities import model_builder
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
 from torch import nn
@@ -83,7 +83,7 @@ class Gemma3MM(nn.Module):
       image_indices: Optional[torch.Tensor] = None,
       image_feat_indices: Optional[torch.Tensor] = None,
       pixel_values: torch.Tensor = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     _, seq_len = tokens.size()
     assert self.config.decoder_config.max_seq_len >= seq_len, (
@@ -150,6 +150,7 @@ class Gemma3MM(nn.Module):
         export_config=export_config,
     )
 
+
 def get_fake_model_config(**kwargs) -> Gemma3MMConfig:
   return Gemma3MMConfig(
       image_encoder_config=image_encoder.get_fake_image_encoder_config(),

ai_edge_torch/generative/examples/llama/convert_to_tflite.py

@@ -15,15 +15,15 @@
 
 """Example of converting Llama 3.2 1B model to multi-signature tflite model."""
 
-import os
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.llama import llama
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 
 
 flags = converter.define_conversion_flags('llama')
+ExportConfig = export_config.ExportConfig
+
 
 _MODEL_SIZE = flags.DEFINE_enum(
     'model_size',

ai_edge_torch/generative/examples/openelm/convert_to_tflite.py

@@ -15,14 +15,14 @@
 
 """Example of converting OpenELM model to multi-signature tflite model."""
 
-import os
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.openelm import openelm
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("openelm")
+ExportConfig = export_config.ExportConfig
+
 
 def main(_):
   pytorch_model = openelm.build_model(

ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py

@@ -15,15 +15,15 @@
 
 """Example of converting a PaliGemma model to multi-signature tflite model."""
 
-import os
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.paligemma import paligemma
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 import torch
 
 flags = converter.define_conversion_flags('paligemma2-3b-224')
+ExportConfig = export_config.ExportConfig
+
 
 _VERSION = flags.DEFINE_enum(
     'version',

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -19,6 +19,7 @@ from typing import Optional
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
@@ -54,7 +55,7 @@ class Decoder(model_builder.DecoderOnlyModel):
       kv_cache: kv_utils.KVCache,
       input_embeds: torch.Tensor = None,
       mask: Optional[torch.Tensor] = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if input_embeds is None:
       return super().forward(

ai_edge_torch/generative/examples/paligemma/decoder2.py

@@ -20,6 +20,7 @@ from typing import Optional
 from ai_edge_torch.generative.examples.gemma import gemma2
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
@@ -57,7 +58,7 @@ class Decoder2(gemma2.Gemma2):
       kv_cache: kv_utils.KVCache,
       input_embeds: torch.Tensor = None,
       mask: Optional[torch.Tensor] = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if input_embeds is None:
       return super().forward(tokens, input_pos, kv_cache, mask, export_config)

ai_edge_torch/generative/examples/paligemma/paligemma.py

@@ -23,7 +23,7 @@ from ai_edge_torch.generative.examples.paligemma import decoder2
 from ai_edge_torch.generative.examples.paligemma import image_encoder
 import ai_edge_torch.generative.layers.kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
-from ai_edge_torch.generative.utilities import model_builder
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
 from torch import nn
@@ -71,7 +71,7 @@ class PaliGemma(nn.Module):
       kv_cache: kv_utils.KVCache,
       mask: Optional[torch.Tensor] = None,
       pixel_values: torch.Tensor = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if pixel_values is None:
       return self.decoder(

ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py

@@ -15,14 +15,14 @@
 
 """Example of converting a Phi-3.5 model to multi-signature tflite model."""
 
-import os
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.phi import phi3
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("phi3")
+ExportConfig = export_config.ExportConfig
+
 
 def main(_):
   pytorch_model = phi3.build_model(

ai_edge_torch/generative/examples/phi/convert_phi4_to_tflite.py

@@ -15,14 +15,14 @@
 
 """Example of converting a Phi-4 model to multi-signature tflite model."""
 
-import os
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.phi import phi4
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("phi4")
+ExportConfig = export_config.ExportConfig
+
 
 def main(_):
   pytorch_model = phi4.build_model(

ai_edge_torch/generative/examples/phi/convert_to_tflite.py

@@ -15,16 +15,15 @@
 
 """Example of converting a Phi-2 model to multi-signature tflite model."""
 
-import os
-import pathlib
 
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.phi import phi2
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("phi2")
+ExportConfig = export_config.ExportConfig
+
 
 def main(_):
   pytorch_model = phi2.build_model(

ai_edge_torch/generative/examples/qwen/convert_to_tflite.py

@@ -15,14 +15,14 @@
 
 """Example of converting Qwen 2.5 models to multi-signature tflite model."""
 
-import os
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.qwen import qwen
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags('qwen')
+ExportConfig = export_config.ExportConfig
+
 
 _MODEL_SIZE = flags.DEFINE_enum(
     'model_size',

ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py

@@ -15,14 +15,14 @@
 
 """Example of converting a Qwen 2.5 VL model to multi-signature tflite model."""
 
-import os
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.qwen_vl import qwen_vl
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags('qwen_vl')
+ExportConfig = export_config.ExportConfig
+
 
 _IMAGE_HEIGHT = flags.DEFINE_integer(
     'image_height',

ai_edge_torch/generative/examples/qwen_vl/decoder.py

@@ -19,6 +19,7 @@ from typing import Optional, Tuple
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 from ai_edge_torch.generative.utilities import model_builder
 import torch
 
@@ -41,7 +42,7 @@ class Decoder(model_builder.DecoderOnlyModel):
       input_embeds: torch.Tensor = None,
       rope: Tuple[torch.Tensor, torch.Tensor] = None,
       mask: Optional[torch.Tensor] = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if input_embeds is None:
       _, seq_len = tokens.size()

ai_edge_torch/generative/examples/qwen_vl/qwen_vl.py

@@ -22,7 +22,7 @@ from ai_edge_torch.generative.examples.qwen_vl import decoder
 from ai_edge_torch.generative.examples.qwen_vl import image_encoder
 import ai_edge_torch.generative.layers.kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
-from ai_edge_torch.generative.utilities import model_builder
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
 from torch import nn
@@ -61,7 +61,7 @@ class QwenVL(nn.Module):
       kv_cache: kv_utils.KVCache,
       mask: Optional[torch.Tensor] = None,
       pixel_values: torch.Tensor = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if pixel_values is None:
       return self.decoder(

ai_edge_torch/generative/examples/smollm/convert_to_tflite.py

@@ -20,7 +20,7 @@ from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.smollm import smollm
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import model_builder
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 
 flags = converter.define_conversion_flags('smollm')
 
@@ -42,7 +42,7 @@ def main(_):
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=model_builder.ExportConfig(
+      export_config=export_cfg.ExportConfig(
           decode_batch_size=_DECODE_BATCH_SIZE.value
       ),
   )

ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py

@@ -15,12 +15,10 @@
 
 """Example of converting SmolLM2 model to multi-signature tflite model."""
 
-import os
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.smollm import smollm
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import model_builder
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 
 flags = converter.define_conversion_flags('smollm2')
 
@@ -43,7 +41,7 @@ def main(_):
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=model_builder.ExportConfig(
+      export_config=export_cfg.ExportConfig(
           decode_batch_size=_DECODE_BATCH_SIZE.value
       ),
   )

ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py

@@ -22,11 +22,12 @@ from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.model_config as cfg
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 import torch
 from torch import nn
 
 RoPECache = Tuple[torch.Tensor, torch.Tensor]
+ExportConfig = export_config.ExportConfig
 
 
 class ToyModelWithKVCache(torch.nn.Module):

ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py

@@ -15,14 +15,14 @@
 
 """Example of converting TinyLlama model to multi-signature tflite model."""
 
-import os
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("tiny_llama")
+ExportConfig = export_config.ExportConfig
+
 
 def main(_):
   pytorch_model = tiny_llama.build_model(

ai_edge_torch/generative/test/test_custom_dus.py

@@ -15,6 +15,7 @@
 
 """A suite of tests to validate the Dynamic Update Slice Custom Op."""
 
+from ai_edge_torch.generative.custom_ops.dynamic_update_slice import dynamic_update_slice
 import torch
 from torch import nn
 

ai_edge_torch/generative/utilities/converter.py

@@ -23,9 +23,11 @@ from ai_edge_torch._convert import converter as converter_utils
 from ai_edge_torch.generative.layers import lora as lora_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.quantize import quant_recipes
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 import torch
 
+ExportConfig = export_config.ExportConfig
+
 
 class ExportableModule(torch.nn.Module):
 

ai_edge_torch/generative/utilities/export_config.py

@@ -0,0 +1,38 @@
+# Copyright 2024 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.
+# ==============================================================================
+
+"""Config for customizing model export process."""
+import dataclasses
+from typing import List, Optional
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+import torch
+
+dataclass = dataclasses.dataclass
+
+
+@dataclass
+class ExportConfig:
+  """Model generating configuration settings."""
+
+  # On prefill signatures, should the model produce logit output?
+  # When False, only decode signatures will produce output.
+  output_logits_on_prefill: bool = False
+  # Attention masks given as inputs to the model.
+  prefill_mask: Optional[torch.Tensor | List[torch.Tensor]] = None
+  decode_mask: Optional[torch.Tensor | List[torch.Tensor]] = None
+  # The KV Cache class for K and V buffers in attention.
+  kvcache_cls: type = kv_utils.KVCache
+  # The batch size of the decode signature.
+  decode_batch_size: int = 1

ai_edge_torch/generative/utilities/model_builder.py

@@ -16,8 +16,7 @@
 """Utilities to be used for re-authoring transformer models."""
 
 import copy
-from dataclasses import dataclass
-from typing import List, Optional, Tuple
+from typing import Optional, Tuple
 
 from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import builder
@@ -25,6 +24,7 @@ from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.generative.layers import lora as lora_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
 from torch import nn
@@ -48,22 +48,6 @@ TENSOR_NAMES_WITH_SEPARATE_LM_HEAD = copy.copy(TENSOR_NAMES)
 TENSOR_NAMES_WITH_SEPARATE_LM_HEAD.lm_head = "lm_head"
 
 
-@dataclass
-class ExportConfig:
-  """Model generating configuration settings."""
-
-  # On prefill signatures, should the model produce logit output?
-  # When False, only decode signatures will produce output.
-  output_logits_on_prefill: bool = False
-  # Attention masks given as inputs to the model.
-  prefill_mask: Optional[torch.Tensor | List[torch.Tensor]] = None
-  decode_mask: Optional[torch.Tensor | List[torch.Tensor]] = None
-  # The KV Cache class for K and V buffers in attention.
-  kvcache_cls: type = kv_utils.KVCache
-  # The batch size of the decode signature.
-  decode_batch_size: int = 1
-
-
 class DecoderOnlyModel(nn.Module):
   """A simple decoder-only transformer model built from the Edge Generative API.
 
@@ -107,7 +91,7 @@ class DecoderOnlyModel(nn.Module):
       kv_cache: kv_utils.KVCache,
       mask: Optional[torch.Tensor] = None,
       lora: Optional[lora_utils.LoRA] = None,
-      export_config: Optional[ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     _, seq_len = tokens.size()
     assert self.config.max_seq_len >= seq_len, (
@@ -139,7 +123,7 @@ class DecoderOnlyModel(nn.Module):
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
       lora: Optional[lora_utils.LoRA] = None,
-      export_config: Optional[ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     """Forwards the model with input embeddings."""
     assert len(self.transformer_blocks) == len(kv_cache.caches), (

ai_edge_torch/generative/utilities/verifier.py

@@ -19,9 +19,11 @@ import logging
 from typing import Any, List, Optional
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 import torch
 
+ExportConfig = export_config.ExportConfig
+
 
 class ModelWrapper(torch.nn.Module):
   """A wrapper for the model to be verified.

ai_edge_torch/odml_torch/lowerings/utils.py

@@ -14,10 +14,10 @@
 # ==============================================================================
 """Utilities for building MLIR lowerings."""
 
+from collections.abc import Callable
 import functools
 import numbers
-from typing import Any
-from typing import Optional
+from typing import Any, Optional, Union
 from ai_edge_torch.odml_torch import export_utils
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import hlo as stablehlo
@@ -222,3 +222,48 @@ def convert_int_to_float(t: ir.Value) -> ir.Value:
     return stablehlo.convert(
         ir.RankedTensorType.get(t.type.shape, ir.F64Type.get()), t
     )
+
+
+# IR Helpers
+IrValues = Union[ir.Value, tuple[ir.Value, ...]]
+
+
+# Non-canonicalized dtype to IR type mapping.
+_numpy_dtype_to_ir_type: dict[np.dtype, Callable[[], ir.Type]] = {
+    np.dtype(np.bool_): functools.partial(ir.IntegerType.get_signless, 1),
+    np.dtype(np.int8): functools.partial(ir.IntegerType.get_signless, 8),
+    np.dtype(np.int16): functools.partial(ir.IntegerType.get_signless, 16),
+    np.dtype(np.int32): functools.partial(ir.IntegerType.get_signless, 32),
+    np.dtype(np.int64): functools.partial(ir.IntegerType.get_signless, 64),
+    np.dtype(np.uint8): functools.partial(ir.IntegerType.get_unsigned, 8),
+    np.dtype(np.uint16): functools.partial(ir.IntegerType.get_unsigned, 16),
+    np.dtype(np.uint32): functools.partial(ir.IntegerType.get_unsigned, 32),
+    np.dtype(np.uint64): functools.partial(ir.IntegerType.get_unsigned, 64),
+    np.dtype(np.float16): ir.F16Type.get,
+    np.dtype(np.float32): ir.F32Type.get,
+    np.dtype(np.float64): ir.F64Type.get,
+    np.dtype(np.complex64): lambda: ir.ComplexType.get(ir.F32Type.get()),
+    np.dtype(np.complex128): lambda: ir.ComplexType.get(ir.F64Type.get()),
+}
+
+
+def numpy_dtype_to_ir_type(dtype: np.dtype | np.generic) -> ir.Type:
+  assert isinstance(dtype, (np.dtype, np.generic)), type(dtype)
+  dtype = np.dtype(dtype)
+  try:
+    ir_type_factory = _numpy_dtype_to_ir_type[dtype]
+  except KeyError as err:
+    raise TypeError(
+        f"No numpy_dtype_to_ir_type handler for dtype: {dtype}"
+    ) from err
+  return ir_type_factory()
+
+
+def numpy_array_constant(x: np.ndarray | np.generic) -> IrValues:
+  element_type = numpy_dtype_to_ir_type(x.dtype)
+  shape = x.shape
+  if x.dtype == np.bool_:
+    x = np.packbits(x, bitorder="little")  # type: ignore
+  x = np.ascontiguousarray(x)
+  attr = ir.DenseElementsAttr.get(x, type=element_type, shape=shape)  # type: ignore
+  return stablehlo.constant(attr)

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.4.0.dev20250402"
+__version__ = "0.4.0.dev20250405"

{ai_edge_torch_nightly-0.4.0.dev20250402.dist-info → ai_edge_torch_nightly-0.4.0.dev20250405.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.4.0.dev20250402
+Version: 0.4.0.dev20250405
 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.4.0.dev20250402.dist-info → ai_edge_torch_nightly-0.4.0.dev20250405.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=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=WqI1n7ZtdJC36jS18ODSCCHgcMQ12kAmRNk-TTrOeAo,706
+ai_edge_torch/version.py,sha256=dHwQF4IL6SQLgcDQtrJ9QSvkuyMetKYrnxGGO9rTHtE,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=gpXQnifODU-mWxkUZw_3ov1lEYBw1SPVIcqj5k7pTGo,5550
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -50,52 +50,52 @@ ai_edge_torch/generative/custom_ops/dynamic_update_slice.py,sha256=ZGAq2CfWZsfef
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/amd_llama_135m/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/amd_llama_135m/amd_llama_135m.py,sha256=urNif89PyCXbdXT5spOeDvdM5luJ-a5HaXHM86v4JnU,2766
-ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py,sha256=aHWsVo_oKpbkEfpQpyPb5DiE9OsUNPBnyjGRHRsVQvs,1589
+ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py,sha256=z5MWiZLnsQzhNYMiQbcI9i0ki-dtkbimCptkiTFZxwo,1586
 ai_edge_torch/generative/examples/amd_llama_135m/verify.py,sha256=o13NkFlBgawBsjdJup05VMUjAPvDRAmig6VyEkX8q6U,2426
 ai_edge_torch/generative/examples/deepseek/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
-ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py,sha256=0RY9zQJX2RFD0HdagcfACYN2tYmQACJqHafgbm2vJZU,1573
+ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py,sha256=r6Pb5_LRKvw2QrOMn3PzunrVxPB-LSdyU2H1XORZo9A,1553
 ai_edge_torch/generative/examples/deepseek/deepseek.py,sha256=AOAJ7ltXwY5IbmcCP2nVHW9FmRwexzfNxnoDlR-sW9c,2885
 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=ylXh2lyu9fRADCV1AEc63dwu_AhoW3MBDrSOdPSzWfQ,1532
-ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=BPtCVR4aaRxUeBHsp--nZH1pzmIrisajiIM1pDIB6gw,1556
+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=lR-T25GkjCfd_sN8mAKY_0XNA0MEkMgsj4ZBQnnytHo,11465
+ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=SyyRae8vWLn0WxduxtahzVRbdSq4T2k5-7t8PfCR_k8,11534
 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/gemma3/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
-ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=i_3kpz2wT2La9047ISoOQigUw7JOhpg2QMHrXp0J6Ic,2955
-ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=4Vf1zA94qLyNzj9iLU0jrd3kzFFZXft4uiItoIBjKyM,15632
-ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=NQzqZ55cmC8tGlZ1SKkDeD0Su8mZ79KiazCS8X08xUY,6473
+ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=6Dkxi7Vs8xBaqMif00ATQSr_hTPhYXMdDqHwzOsAzq8,2952
+ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=__kpzArZ0mLfX7IzpHPmYFuhKTP9uI_9Lrzk_EfFDlE,15701
+ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=5PEt0aWJ5wkUBvMoWFOJ-C48ZhG7uCVb8PCKQtZ8Fvw,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=u30qiZu3HJCTt5noWqtf9PgGLKQ87ke4Zpa4cpG6-As,8883
 ai_edge_torch/generative/examples/llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/llama/convert_to_tflite.py,sha256=euNMhQ-O7PUQ0v2ootyCbh7yra378toBdCP3TRjNMiQ,1765
+ai_edge_torch/generative/examples/llama/convert_to_tflite.py,sha256=A4uLUdqvU1NKo3seqZlWSS3fqYahnEKqNBQBJO6yXvE,1762
 ai_edge_torch/generative/examples/llama/llama.py,sha256=UKvMO85_5z1vEY5MVu6QBW_vpQYA8LWHbJI4Yx6BrCc,6592
 ai_edge_torch/generative/examples/llama/verify.py,sha256=X7oKQi85M789ugBrOlMvzk8eSRR3Kf1Mprfl-U-WIpo,2842
 ai_edge_torch/generative/examples/moonshine/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/moonshine/convert_moonshine_to_tflite.py,sha256=7m3rYRzThRDYb-7pGnpLr3ACi4PWX07Mg20Q98ArPc4,1714
 ai_edge_torch/generative/examples/moonshine/moonshine.py,sha256=nZ2b8u4TmsB5sgdClgAuH8E78bcTv9RCnF9666HqP2M,3394
 ai_edge_torch/generative/examples/openelm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=tlBBAQ33T4CHgZOclPf735gLUxnkTkXpruLObu-_RDg,1554
+ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=Hgp31zIQdJsTweRMr0U3d2SKW1h2nWnqWt1FlmuQqiI,1551
 ai_edge_torch/generative/examples/openelm/openelm.py,sha256=sIJ8Ie1oxFrJM-1jvv2ukiJbQOTIUGuMEZvmwZbt3n0,4556
 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=H8PKHHxeoVSF5KmDVcdfnoFU4loMQXrpvdvXt1ATSOg,1959
-ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=z658dW_D0Iqvo6xnh4vG7_o17-Fufndyis8Rq5yafJY,5439
-ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=GZa0Ou_DvOijB2nTL_jRvGbn0_dvJPosQAPf47yqicw,5988
+ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=2GLE4empjc8IANssR02ECFUqhdUNJV_OVHCf1UXKL8Y,1956
+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/image_encoder.py,sha256=SvuR97sjkBtfkerH7Hu1UXB8kCFLpEATNbPfCbNAyfo,5614
-ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=x1mgRtVLxkCTvlkPow3y7ADoGTjUh5uc5pF46mxatLw,6099
+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
 ai_edge_torch/generative/examples/paligemma/verify_decoder.py,sha256=al5wMPWri4IRVWrLmCplPi6uoCzwh0vBHMGnCt-XUqo,2690
 ai_edge_torch/generative/examples/paligemma/verify_decoder2.py,sha256=tm-UfLr0YeBRVcQsWLBOMWI9JUzHmtPEbYK2vpITpqY,2534
 ai_edge_torch/generative/examples/paligemma/verify_image_encoder.py,sha256=vNm-wTT8BD6zbX6GocfP1QrVoHl0zSvuVxoXN36eeiU,3540
 ai_edge_torch/generative/examples/phi/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py,sha256=OyblvnMpEPmHQ-ifBPiQUbECbB-3BkoiHebwL_6bB_k,1543
-ai_edge_torch/generative/examples/phi/convert_phi4_to_tflite.py,sha256=3xwSUl-eA1-2hJo2bgh6BRDWUz9YTcDIZ_3h8CsW5ys,1541
-ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=M2m359q2jNXzBIlAn1Jf9knGoYT3XNAYzzUQFf6LzMI,1557
+ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py,sha256=Y2qaObMJeh9UABkUI7FBm4sCGi2YMQhsj0CSOS2fYek,1540
+ai_edge_torch/generative/examples/phi/convert_phi4_to_tflite.py,sha256=TuGW_FPMs0pV7ZBe46FfaDrlfte4Dz75vGHmBOCFfww,1538
+ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=VZe7OQ54dgOGWe74XT2W7zZBm5uJaeIF8ZuNakkL0iA,1539
 ai_edge_torch/generative/examples/phi/phi2.py,sha256=c6PYCky7yJn6MVIYOCTx8S_CH27kOPmJbRZcI95nbZs,3477
 ai_edge_torch/generative/examples/phi/phi3.py,sha256=ddo52Inl5ub81q460cEyKhnsC3txellRErut-_qtBbM,6949
 ai_edge_torch/generative/examples/phi/phi4.py,sha256=OkMwLGe8l2JEAgOFi19AdbNBl1xp1djZBZo8MJP58ho,5732
@@ -103,20 +103,20 @@ ai_edge_torch/generative/examples/phi/verify.py,sha256=YPFCdbnfmvq38fbpBNr0kHPfS
 ai_edge_torch/generative/examples/phi/verify_phi3.py,sha256=kVYaBVvddfQng0IyZGxyTJEzhiPO0G4VFJm2WOc2Q94,2360
 ai_edge_torch/generative/examples/phi/verify_phi4.py,sha256=BoCa5kUBRHtMQ-5ql6yD4pG4xHJMyUiQlpMOWVx-JgY,2356
 ai_edge_torch/generative/examples/qwen/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/qwen/convert_to_tflite.py,sha256=vx4gsuqsHTUiME68Q6Xkqs4lwvhBy8Cm9dQWVmqZEbY,1805
+ai_edge_torch/generative/examples/qwen/convert_to_tflite.py,sha256=om3lXL1RnA87PkfU_cRfP6RnPgXrCmaB-cK98H-nqbA,1802
 ai_edge_torch/generative/examples/qwen/qwen.py,sha256=Zi_qiQ1JPokXZ95jgSEnQp3F-LKzFCvWvFLKhJjnASo,4199
 ai_edge_torch/generative/examples/qwen/verify.py,sha256=9_AyEJTeUfvhhID64Rto2bflFPyXMFokdQLsseLUMiI,2775
 ai_edge_torch/generative/examples/qwen_vl/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
-ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py,sha256=5r3jOXMYooS5pfc9-RtceUa4o5FFvDEztK-i_y1WLZE,1998
-ai_edge_torch/generative/examples/qwen_vl/decoder.py,sha256=0x4iDg2cBe3PFnjVce3nj7g2rjagGHcKqRCfbASNxA8,4402
+ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py,sha256=OcE2-8lqAukoK5hM1sqdgfXU37kxWQ84racweNAdjyk,1995
+ai_edge_torch/generative/examples/qwen_vl/decoder.py,sha256=7RFM25tDj_b0FkpSv8RUWir8K8v9p2jMtwZmP4VAUhw,4474
 ai_edge_torch/generative/examples/qwen_vl/image_encoder.py,sha256=nHzBe_YSPnUe1d5i09v4bePQomVifzJNeUjRfprmxC0,14878
-ai_edge_torch/generative/examples/qwen_vl/qwen_vl.py,sha256=rcYHkpO-NbF4F1Da7q2xNiTng9NHiLx59HyuOgQX5W0,7753
+ai_edge_torch/generative/examples/qwen_vl/qwen_vl.py,sha256=mfLFrT8NPEPh9CqlJYHwh-I2y6ST7hH_vEmbZYartHQ,7764
 ai_edge_torch/generative/examples/qwen_vl/verify.py,sha256=JUwHoC_zvcC3RC3wZ3e3e6fGmrhbgdoztjK8HGSUG8I,5044
 ai_edge_torch/generative/examples/qwen_vl/verify_decoder.py,sha256=xPWoOBLh2eK12KEhELLYymfL7xvc0chmYC98c6x37oo,2602
 ai_edge_torch/generative/examples/qwen_vl/verify_image_encoder.py,sha256=PZ392nDoJG2OmHZ_7Jet3Zu1JkN6QErxKcDc7a-PPds,3126
 ai_edge_torch/generative/examples/smollm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/smollm/convert_to_tflite.py,sha256=bGeS_emY_TTYzBKwe4BIzUQRTf85pJHXNUIIFUtwMks,1737
-ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py,sha256=5MZha6sFuIOG27h7Af44I0TvEpHfUjcWkHTgUyqeBiU,1743
+ai_edge_torch/generative/examples/smollm/convert_to_tflite.py,sha256=IjV0jriRKlF9aV5yLjtONjACb4_VxNIAGk9w1sr_hmc,1748
+ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py,sha256=Wa_3OWXcM389iySwS5E47uCYZaTj6h-4RTP_Xi2-1aE,1721
 ai_edge_torch/generative/examples/smollm/smollm.py,sha256=3uUltb6D3Q1aHpndcYTJrsWM_RBwLAraKDniH8ZZous,3779
 ai_edge_torch/generative/examples/smollm/verify.py,sha256=KpYxVz_lv61YWy6HLfwT68n0owZMvty5Rr3W7ZNWWSw,2702
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -141,9 +141,9 @@ ai_edge_torch/generative/examples/t5/t5_attention.py,sha256=l01oYyJo77INzRwN4xqX
 ai_edge_torch/generative/examples/test_models/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/test_models/convert_toy_model.py,sha256=6-WaNHckq_LlXMVTh8x90MGWeWq2bu_T_XQd3w9FnGg,3261
 ai_edge_torch/generative/examples/test_models/toy_model.py,sha256=Crpj-vOwSViHpblXOrRJmsIn4DrHyuB3XZ8kHifb7LA,5203
-ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=Ab_N9xc-4DImA-Pvevr-nnnslBXScXVo4Pw7L3_OlhI,4732
+ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=-z5tkQzGHbo37eAl9sDAJuT1Egxm8xI9CZmYLcmqIfU,4761
 ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=YwI3JWoTFcsBQ3nLLbi3ig94p50sJOMguyguvBCpE-o,1568
+ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=LPxg7mAJ_aAUIx6eE5bxixPA8Ep9Vul0CWJoNcrD5oE,1565
 ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=mhJ18rb9sxrYRzv1YSzhbNs97oUZck99avZDcUO2oV8,2800
 ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=LRu6PSw7Lqu6HGbv1tO2i0nUCqe-VkRgboA10VZ7KNg,2431
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=4rFrppMRKlTwwZeX1ON_cdp4yUqoTOES161IZQkJF6c,1143
@@ -176,7 +176,7 @@ ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=4fgmP_GgeiFUOkIaC
 ai_edge_torch/generative/quantize/quant_recipes.py,sha256=0Kvr_o7pbMnE8VMe6Ml0FBxkHM6RJ3C14B2I1mjItjc,2030
 ai_edge_torch/generative/quantize/supported_schemes.py,sha256=FjdycEOvxRgBmQdZVufetPvkDoD7rUowIOSKV9oV5Kk,1418
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/test/test_custom_dus.py,sha256=ifgnUCWihT59eFdLrlc5_j9sWygEKclU6Iqw6zdlgeI,3177
+ai_edge_torch/generative/test/test_custom_dus.py,sha256=MjIhTvkTko872M35XMciobvICcDWTcIDJ3rociko-wM,3267
 ai_edge_torch/generative/test/test_kv_cache.py,sha256=MBPS-0bDXB0tQSKHa1XwDQeVIfabRbc8JQA99h9fzlQ,5961
 ai_edge_torch/generative/test/test_loader.py,sha256=9mQUeeZKOVApOWSWl2cN9c10axZjMKM1-0Zd823CCS4,3449
 ai_edge_torch/generative/test/test_lora.py,sha256=6QIM6RLTc2HrodGpp_aS3OxM9Rco2KAzEnYgotkg41M,5310
@@ -185,14 +185,15 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=-v2Vj7Qdd3Gy
 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=Ol-CCztmjAaSeRGkUaKxA9-5ihATYuLS5kwZo1hpA-A,9577
+ai_edge_torch/generative/utilities/converter.py,sha256=ycXDcd3ZE-EdjksDjHi4ru3JpfhtrfOompg_990qvWI,9607
+ai_edge_torch/generative/utilities/export_config.py,sha256=-UuukWqUUj8RM8lTtMCa_PD6SqCZv97i4BMiJA2zBPg,1491
 ai_edge_torch/generative/utilities/loader.py,sha256=7p__m2JryWphGlYOuRxdoT4id4_tWJEVOV7y2X4H-Ak,13737
-ai_edge_torch/generative/utilities/model_builder.py,sha256=eY3qAcBhupIn955YnWuzUi9hoWYvl4ntRWA6PBudzMo,6888
+ai_edge_torch/generative/utilities/model_builder.py,sha256=ZYX1TxpFdj573du2QCyHJlFjx4q1m12R74fp4Gwl92A,6343
 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/transformers_verifier.py,sha256=8sp9m_FMcXn7nqOrochtu2jIANkJKhnhIBUmH0ZTDR4,1549
-ai_edge_torch/generative/utilities/verifier.py,sha256=6lnBU9Cy5GanB8JWK3-2_VU3PxqunDWGe-SgSLba5Yw,12065
+ai_edge_torch/generative/utilities/verifier.py,sha256=gsBGv-WgeP73Bmao29CiTIQOC1YZ43IUJcGzytpcZyM,12095
 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
@@ -232,7 +233,7 @@ ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQD
 ai_edge_torch/odml_torch/lowerings/_rand.py,sha256=g6SuqDkuC6hD35lyP1-5H7ASDIzPSmKukeNT5naZSv8,4133
 ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
 ai_edge_torch/odml_torch/lowerings/registry.py,sha256=Tp2h11l5uTifO0aIkuUOWAF_ibEjmd65Xx99w3EXuGE,1924
-ai_edge_torch/odml_torch/lowerings/utils.py,sha256=tIyZiSy2Rbtea9ZlRYUfTaYE5vW_lAU6itT6_rUp8Qg,7028
+ai_edge_torch/odml_torch/lowerings/utils.py,sha256=-TzK1igPgR38oZkU1iPh-DZhlKVwuBtGWVC-y81PXzY,8935
 ai_edge_torch/odml_torch/passes/__init__.py,sha256=AVwIwUTMx7rXacKjGy4kwrtMd3XB2v_ncdc40KOjUqQ,1245
 ai_edge_torch/quantize/__init__.py,sha256=aB5dXot04bqyUhpsDFvxt9CIi15QAC4euvqOndJ0XLU,714
 ai_edge_torch/quantize/pt2e_quantizer.py,sha256=CKIEhs9jCcna64qj1jFH9zEbMbRdyeGV_TmSqEBPjes,15741
@@ -242,8 +243,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=dguMa-aEi-WDPnmGBUs2IPdEmt2IVmHOELH19uiJ1uU,3014
 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.4.0.dev20250402.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.4.0.dev20250402.dist-info/METADATA,sha256=Dm1OfGpmUSSXIT6gnxQlAYu5qm8e6WbLdq-ocbXEoBU,1966
-ai_edge_torch_nightly-0.4.0.dev20250402.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.4.0.dev20250402.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.4.0.dev20250402.dist-info/RECORD,,
+ai_edge_torch_nightly-0.4.0.dev20250405.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.4.0.dev20250405.dist-info/METADATA,sha256=er5v3ZpkS5bJW9tI2dxbtPqaRVnaV99gv8vlMuV_AfI,1966
+ai_edge_torch_nightly-0.4.0.dev20250405.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.4.0.dev20250405.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.4.0.dev20250405.dist-info/RECORD,,