ai-edge-torch-nightly 0.5.0.dev20250508__py3-none-any.whl → 0.5.0.dev20250510__py3-none-any.whl

ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py

@@ -20,7 +20,9 @@ from ai_edge_torch.generative.examples.gemma import gemma2
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
 
-flags = converter.define_conversion_flags("gemma2-2b")
+flags = converter.define_conversion_flags(
+    "gemma2-2b", default_mask_as_input=True, default_transpose_kv_cache=True
+)
 
 
 def main(_):

ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py

@@ -20,7 +20,9 @@ from ai_edge_torch.generative.examples.gemma3 import gemma3
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
 
-flags = converter.define_conversion_flags('gemma3-1b')
+flags = converter.define_conversion_flags(
+    'gemma3-1b', default_mask_as_input=True, default_transpose_kv_cache=True
+)
 
 _MODEL_SIZE = flags.DEFINE_string(
     'model_size',

ai_edge_torch/generative/examples/gemma3/decoder.py

@@ -119,9 +119,7 @@ class Decoder(nn.Module):
         config.vocab_size, config.embedding_dim, padding_idx=0
     )
     self.lm_head = nn.Linear(
-        config.embedding_dim,
-        config.vocab_size,
-        bias=config.lm_head_use_bias,
+        config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
     )
     # Gemma3 re-uses the embedding as the head projection layer.
     self.lm_head.weight.data = self.tok_embedding.weight.data
@@ -130,30 +128,13 @@ class Decoder(nn.Module):
         for idx in range(config.num_layers)
     )
     self.final_norm = builder.build_norm(
-        config.embedding_dim,
-        config.final_norm_config,
+        config.embedding_dim, config.final_norm_config
     )
     self.mask_cache = attn_utils.build_causal_mask_cache(
         size=config.kv_cache_max,
     )
-    # Gemma3 has same hyper parameters for each layer except for attention
-    # types. Use the first layer.
-    attn_config = config.block_config(0).attn_config
-    self.sliding_window_mask_cache = attn_utils.build_sliding_window_mask_cache(
-        size=config.kv_cache_max,
-        window_size=attn_config.sliding_window_size,
-    )
     self.config = config
 
-  def get_attention_mask(
-      self,
-      attn_type: cfg.AttentionType,
-      input_pos: torch.Tensor,
-  ) -> torch.Tensor:
-    if attn_type == cfg.AttentionType.LOCAL_SLIDING:
-      return self.sliding_window_mask_cache.index_select(2, input_pos)
-    return self.mask_cache.index_select(2, input_pos)
-
   def get_local_global_attention_mask(
       self,
       attention_mask: torch.Tensor,
@@ -200,9 +181,7 @@ class Decoder(nn.Module):
         sliding_mask_bool,
         torch.zeros_like(sliding_mask_bool, dtype=torch.float),
         torch.full_like(
-            sliding_mask_bool,
-            self.config.causal_mask_value,
-            dtype=torch.float,
+            sliding_mask_bool, self.config.causal_mask_value, dtype=torch.float
         ),
     )
 
@@ -261,7 +240,6 @@ class Decoder(nn.Module):
       pixel_mask = self.build_pixel_mask(image_indices)
     # RoPE parameters are the same for all blocks. Use the first layer.
     attn_config = self.config.block_config(0).attn_config
-    n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
     # Different rotary base for global and local attention
     # based on attention pattern
     rope = [
@@ -273,12 +251,8 @@ class Decoder(nn.Module):
         for i in range(self.config.num_layers)
     ]
     if mask is None:
-      mask = [
-          self.get_attention_mask(
-              self.config.block_config(i).attn_config.attn_type, input_pos
-          )
-          for i in range(self.config.num_layers)
-      ]
+      mask = self.mask_cache.index_select(2, input_pos)
+      mask = mask[:, :, :, : self.config.kv_cache_max]
 
     return self._forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache, pixel_mask, export_config
@@ -305,7 +279,7 @@ class Decoder(nn.Module):
     if pixel_mask is None:
       mask = [
           self.get_local_global_attention_mask(
-              mask,
+              mask[i] if isinstance(mask, list) else mask,
               self.config.block_config(i).attn_config.attn_type,
               input_pos,
               self.config.block_config(i).attn_config.sliding_window_size,
@@ -316,7 +290,7 @@ class Decoder(nn.Module):
       pixel_mask = pixel_mask.index_select(2, input_pos)
       mask = [
           self.compose_mask(
-              mask[i],
+              mask[i] if isinstance(mask, list) else mask,
               pixel_mask,
               self.config.block_config(i).attn_config.attn_type,
           )
@@ -330,6 +304,7 @@ class Decoder(nn.Module):
       if kv_entry:
         updated_kv_entries.append(kv_entry)
     updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entries))
+
     if export_config is not None:
       if (
           torch.numel(input_pos) > 1

ai_edge_torch/generative/quantize/quant_recipe.py

@@ -16,9 +16,12 @@
 from dataclasses import dataclass
 from typing import Optional, Union
 
+from ai_edge_torch.generative.layers import model_config
 from ai_edge_torch.generative.quantize import quant_attrs
 from ai_edge_torch.generative.quantize import supported_schemes
 
+ModelConfig = model_config.ModelConfig
+
 
 @dataclass
 class LayerQuantRecipe:
@@ -52,7 +55,7 @@ class LayerQuantRecipe:
         f'w:{self.weight_dtype.name}, '
         f'{self.mode.name}, '
         f'{self.algorithm.name}, '
-        f'{self.granularity.name}'
+        f'{self.granularity.name}, '
         f'{self.block_size}'
     )
     return f'{base_str})'
@@ -133,6 +136,7 @@ class GenerativeQuantRecipe:
   feedforward: Union[
       Optional[LayerQuantRecipe], Optional[dict[int, LayerQuantRecipe]]
   ] = None
+  _model_config: Optional[ModelConfig] = None
 
   def __str__(self):
     return f"""GenerativeQuantRecipe(

ai_edge_torch/generative/quantize/quant_recipes.py

@@ -63,6 +63,7 @@ def all_supported_int4_dynamic_block_recipe(
       generative_recipe=quant_recipe.GenerativeQuantRecipe(
           default=quant_recipe_utils.create_layer_quant_int4_dynamic_block(
               block_size
-          )
+          ),
+          embedding=quant_recipe_utils.create_layer_quant_int8_dynamic(),
       )
   )

ai_edge_torch/generative/test/test_quantize.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 
 import ai_edge_torch
-from ai_edge_torch import config
 from ai_edge_torch.generative.examples.test_models import toy_model  # NOQA
 from ai_edge_torch.generative.quantize import quant_recipe
 from ai_edge_torch.generative.quantize import quant_recipe_utils

ai_edge_torch/generative/utilities/converter.py

@@ -15,6 +15,7 @@
 
 """Common utility functions for model conversion."""
 
+import enum
 import os
 import pathlib
 from typing import Optional, Union
@@ -42,7 +43,32 @@ class ExportableModule(torch.nn.Module):
     return self.module(*export_args, **full_kwargs)
 
 
-def define_conversion_flags(model_name: str):
+class QuantizationName(str, enum.Enum):
+  """Strings for all supported quantization recipes.
+
+  none: No quantization.
+  dynamic_int8: Dynamic range quantization with int8 weights.
+  weight_only_int8: Weight only quantization with int8 weights.
+  fp16: Float16 quantization.
+  dynamic_int4_block32: Dynamic range quantization with int4 weights and block
+  size of 32, better model quality but slower inference.
+  dynamic_int4_block128: Dynamic range quantization with int4 weights and block
+  size of 128, faster inference but worse model quality.
+  """
+
+  NONE = 'none'
+  DYNAMIC_INT8 = 'dynamic_int8'
+  WEIGHT_ONLY_INT8 = 'weight_only_int8'
+  FP16 = 'fp16'
+  DYNAMIC_INT4_BLOCK32 = 'dynamic_int4_block32'
+  DYNAMIC_INT4_BLOCK128 = 'dynamic_int4_block128'
+
+
+def define_conversion_flags(
+    model_name: str,
+    default_mask_as_input: bool = False,
+    default_transpose_kv_cache: bool = False,
+):
   """Defines common flags used for model conversion."""
 
   flags.DEFINE_string(
@@ -70,10 +96,10 @@ def define_conversion_flags(model_name: str):
       1280,
       'The maximum size of KV cache buffer, including both prefill and decode.',
   )
-  flags.DEFINE_bool(
+  flags.DEFINE_string(
       'quantize',
-      True,
-      'Whether the model should be quantized.',
+      'dynamic_int8',
+      'How the model should be quantized.',
   )
   flags.DEFINE_multi_integer(
       'lora_ranks',
@@ -83,18 +109,78 @@ def define_conversion_flags(model_name: str):
   )
   flags.DEFINE_bool(
       'mask_as_input',
-      False,
+      default_mask_as_input,
       'If true, the mask will be passed in as input. Otherwise, mask will be '
       'built by the model internally.',
   )
   flags.DEFINE_bool(
       'transpose_kv_cache',
-      False,
+      default_transpose_kv_cache,
       'If true, the model will be converted with transposed KV cache.',
   )
   return flags
 
 
+def get_quant_recipe_from_flag(
+    quantize: str,
+) -> Optional[quant_recipes.QuantizationRecipe]:
+  """Processes the quantization flag and returns the corresponding recipe.
+
+  Args:
+      quantize: The quantization type.
+
+  Returns:
+      The quantization recipe, or None if no quantization is needed.
+
+  Raises:
+      ValueError: If the quantization type is not supported.
+  """
+  match quantize:
+    case QuantizationName.NONE:
+      return None
+    case QuantizationName.DYNAMIC_INT8:
+      return quant_recipes.full_int8_dynamic_recipe()
+    case QuantizationName.WEIGHT_ONLY_INT8:
+      return quant_recipes.full_int8_weight_only_recipe()
+    case QuantizationName.FP16:
+      return quant_recipes.full_fp16_recipe()
+    case QuantizationName.DYNAMIC_INT4_BLOCK32:
+      return quant_recipes.full_int4_dynamic_block_recipe(32)
+    case QuantizationName.DYNAMIC_INT4_BLOCK128:
+      return quant_recipes.full_int4_dynamic_block_recipe(128)
+    case _:
+      raise ValueError(f'Unsupported quantization flag: {quantize}')
+
+
+def create_quantize_suffix(quantize: str) -> str:
+  """Creates a suffix for the output file name based on the quantization type.
+
+  Args:
+      quantize: The quantization type.
+
+  Returns:
+      A string representing the quantization suffix.
+
+  Raises:
+      ValueError: If the quantization type is not supported.
+  """
+  match quantize:
+    case QuantizationName.NONE:
+      return 'f32'
+    case QuantizationName.DYNAMIC_INT8:
+      return 'q8'
+    case QuantizationName.WEIGHT_ONLY_INT8:
+      return 'q8_wo'
+    case QuantizationName.FP16:
+      return 'fp16'
+    case QuantizationName.DYNAMIC_INT4_BLOCK32:
+      return 'q4_block32'
+    case QuantizationName.DYNAMIC_INT4_BLOCK128:
+      return 'q4_block128'
+    case _:
+      raise ValueError(f'Unsupported quantization flag: {quantize}')
+
+
 def _build_mask(mask_len, kv_cache_max_len, causal_mask_value) -> torch.Tensor:
   if isinstance(mask_len, list):
     return [
@@ -114,7 +200,7 @@ def convert_to_tflite(
     prefill_seq_len: Union[int, list[int]],
     pixel_values_size: torch.Size = None,
     pixel_seq_len: int = 0,
-    quantize: bool = True,
+    quantize: str = 'dynamic_int8',
     config: cfg.ModelConfig = None,
     lora_ranks: Optional[list[int]] = None,
     export_config: ExportConfig = None,
@@ -160,8 +246,8 @@ def convert_to_tflite(
         embeddings generated by the image encoder with pixel values. The actual
         length of prefill_seq_len will be added by pixel_seq_len when pixel
         values are passed.
-      quantize (bool, optional): Whether the model should be quanized. Defaults
-        to True.
+      quantize (str, optional): The quantization type. Defaults to
+        'dynamic_int8'.
       config (cfg.ModelConfig, optional): The model config used to configure KV
         cache. If None, it uses the config of the pytorch_model.
       lora_ranks (list[int], optional): The ranks of the LORA layers. If None,
@@ -182,7 +268,7 @@ def convert_to_tflite(
       lora = lora_utils.LoRA.zeros(rank, config)
       loras.append(lora)
 
-  quant_suffix = 'q8' if quantize else 'f32'
+  quant_suffix = create_quantize_suffix(quantize)
   kv_size = config.kv_cache_max_len
   lora_suffix = (
       '' if not lora_ranks else f'_lora{",".join(map(str, lora_ranks))}'
@@ -216,7 +302,7 @@ def _export_helper(
     prefill_seq_lens: list[int],
     pixel_values_size: torch.Size,
     pixel_seq_len: int,
-    quantize: bool,
+    quantize: str,
     config: cfg.ModelConfig,
     loras: list[None | lora_utils.LoRA],
     export_config: ExportConfig,
@@ -265,7 +351,8 @@ def _export_helper(
       kv_layout=export_config.kvcache_layout,
   )
 
-  quant_config = quant_recipes.full_int8_dynamic_recipe() if quantize else None
+  quant_config = get_quant_recipe_from_flag(quantize)
+  quant_config._model_config = config
 
   # For export, we create a module that captures any non-exportable,
   # arugments, e.g. the generation config object.
@@ -330,5 +417,7 @@ def _export_helper(
         sample_kwargs=sample_kwargs,
     )
 
-  edge_model = converter.convert(quant_config=quant_config)
+  edge_model = converter.convert(
+      quant_config=quant_config,
+  )
   edge_model.export(output_file)

ai_edge_torch/generative/utilities/model_builder.py

@@ -75,8 +75,7 @@ class DecoderOnlyModel(nn.Module):
         for idx in range(config.num_layers)
     )
     self.final_norm = builder.build_norm(
-        config.embedding_dim,
-        config.final_norm_config,
+        config.embedding_dim, config.final_norm_config
     )
     self.mask_cache = attn_utils.build_causal_mask_cache(
         size=config.kv_cache_max,

ai_edge_torch/lowertools/_shim.py

@@ -50,7 +50,7 @@ def exported_programs_to_tflite(
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
     _tfl_converter_flags: Optional[dict[str, Any]] = None,
-    _saved_model_dir: Optional[str] = None
+    _saved_model_dir: Optional[str] = None,
 ):
   """Converts a list of ExportedProgram to a TFLite model.
 

ai_edge_torch/lowertools/translate_recipe.py

@@ -29,6 +29,8 @@ _IDX_TRANSFORMER_BLOCKS_REGEX_STR = 'transformer_blocks\[{}\]'
 _ATTENTION_REGEX_STR = 'ai_edge_torch.generative.layers.attention'
 _FEEDFORWARD_REGEX_STR = 'ai_edge_torch.generative.layers.feed_forward'
 _EMBEDDING_REGEX_STR = 'Embedding_tok_embedding'
+# TODO: b/415833584 - Improve the regex for pre-softmax layer.
+_DECODE_LOGITS_REGEX_STR = 'StatefulPartitionedCall'
 _ANY_TWO_DIGITS_REGEX_STR = '\d{1,2}'
 
 
@@ -95,10 +97,11 @@ def _set_quant_config(
     rm: quantizer.recipe_manager.RecipeManager,
     layer_recipe: quant_recipe.LayerQuantRecipe,
     regex: str,
+    operation_name: _OpName = _OpName.ALL_SUPPORTED,
 ):
   rm.add_quantization_config(
       regex=regex,
-      operation_name=_OpName.ALL_SUPPORTED,
+      operation_name=operation_name,
       op_config=_OpQuantConfig(
           weight_tensor_config=_TensorQuantConfig(
               num_bits=_get_nbits_from_dtype(layer_recipe.weight_dtype),
@@ -126,6 +129,16 @@ def translate_to_ai_edge_recipe(
 
   if recipe.embedding is not None:
     _set_quant_config(rm, recipe.embedding, _EMBEDDING_REGEX_STR)
+    if (
+        recipe._model_config is not None
+        and recipe._model_config.lm_head_share_weight_with_embedding
+    ):
+      _set_quant_config(
+          rm,
+          recipe.embedding,
+          _DECODE_LOGITS_REGEX_STR,
+          _OpName.FULLY_CONNECTED,
+      )
 
   if recipe.attention is not None:
     if isinstance(recipe.attention, dict):

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.5.0.dev20250508.dist-info → ai_edge_torch_nightly-0.5.0.dev20250510.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250508
+Version: 0.5.0.dev20250510
 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
@@ -22,6 +22,7 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: absl-py
 Requires-Dist: numpy
 Requires-Dist: scipy
 Requires-Dist: safetensors

{ai_edge_torch_nightly-0.5.0.dev20250508.dist-info → ai_edge_torch_nightly-0.5.0.dev20250510.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=7lrbHHeWyBpqJdwFYYooOGJss4Rvg3UAdFSo9K0uzek,706
+ai_edge_torch/version.py,sha256=03QthwiMre1vVY49We8vVzXhxe0zkOzzsTMQZv3hDrk,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
@@ -61,15 +61,15 @@ ai_edge_torch/generative/examples/deepseek/deepseek.py,sha256=yhS_i2kR0GJWpWciCt
 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=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/convert_gemma2_to_tflite.py,sha256=9ozSw2-xuf5Wfh1HeLDTP3wJxxUZmrD3An1njJPMpdI,1594
 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=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=xGxeNKQvgyrENmUQMu0uKymL3qthvbdoxdMbAzwiLz0,15725
+ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=JLXXn2mFEBs4DlHH_O6hpEG9KInJqsCdWy3DrgUjT1c,1827
+ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=v0ZcKrIAvERQLb1wK1Vc_ewWWVZgJFUdRTyoVY0Lfus,14955
 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
@@ -178,9 +178,9 @@ ai_edge_torch/generative/layers/unet/model_config.py,sha256=pPDwLawc23pfMaPVyMJl
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/quantize/example.py,sha256=1lfVNUd2cEyRUnoZ7BLbRJ9IN-FTKiWBtZNPFUzAiWE,1747
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=plMsd7JBi98r2NHsAdMdvS6TPTXAoRFLCwOXu8H3-24,2004
-ai_edge_torch/generative/quantize/quant_recipe.py,sha256=3xT4N5tfggXJqgwKW4ntIkwsrNVtkG2SIUHeiSF5yOs,5579
+ai_edge_torch/generative/quantize/quant_recipe.py,sha256=CEW-ewHxwb59x_GISx4jr7WMihvn-jKWVcBonllzDS4,5724
 ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=h3k_na6rbR08Ip79-2JbkeH8RDk_rrnEGiytuzFDhqc,2678
-ai_edge_torch/generative/quantize/quant_recipes.py,sha256=a71KFHVbjJdBDpYshbUI69NxGNOmPuqp_NZvNSrf00c,2349
+ai_edge_torch/generative/quantize/quant_recipes.py,sha256=5UkUAT0qsWzLtNAeX-M5hEMi-kqoLV70_F76QiXmVZ4,2424
 ai_edge_torch/generative/quantize/supported_schemes.py,sha256=TwR2FpQuBEORy6FshEyHNBMKARWlA2MVtTfX9tXV5aE,1488
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/test/test_custom_dus.py,sha256=MjIhTvkTko872M35XMciobvICcDWTcIDJ3rociko-wM,3267
@@ -189,13 +189,13 @@ ai_edge_torch/generative/test/test_loader.py,sha256=9mQUeeZKOVApOWSWl2cN9c10axZj
 ai_edge_torch/generative/test/test_lora.py,sha256=6QIM6RLTc2HrodGpp_aS3OxM9Rco2KAzEnYgotkg41M,5310
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=mhNJikLnGVGi9NKmXB8FhnqeDy9gtrvC3yEbrTABZ4Y,6163
 ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=vQWmpzMkJ2hPmWpg41ZMWwBsngTykRVzRPHtpbkwiLM,12811
-ai_edge_torch/generative/test/test_quantize.py,sha256=TG6vTF9yOZWe2wW7v8-hmuaQoODwJC1Z-2d5xv3zgfI,7389
+ai_edge_torch/generative/test/test_quantize.py,sha256=kKJ01wscTC2t_Ylr7huO5gNKES01gm3dT1gx52z15PA,7356
 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=d0JOWN5l2vbvt8RzFFiRoulkWiejyEZ21xKv5LdLIyc,11675
+ai_edge_torch/generative/utilities/converter.py,sha256=LrBqxXVxkOWh4abcHfY4QXRpYxjjfEYd4ifrpGGbebI,14441
 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/model_builder.py,sha256=IG-88o7nWI9XrNDnwnQ-MoilsuqJ7KwrnbP3bn2EY9U,6334
 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
@@ -209,12 +209,12 @@ ai_edge_torch/hlfb/mark_pattern/pattern.py,sha256=Ui6BrehF3zJJN7uTxKwbO2yCY9mYjb
 ai_edge_torch/hlfb/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/hlfb/test/test_mark_pattern.py,sha256=5kmOJWCc7sU1Hrqr1y17BtShUrssTfaV1sMyUvdMbsg,5573
 ai_edge_torch/lowertools/__init__.py,sha256=A8WBXvWtuFYYWtNTqPD7waVntLaSVAnSMwx5ugjZBIw,761
-ai_edge_torch/lowertools/_shim.py,sha256=Mbg16tnCVK0YsHowfbpqpNX1qySuMLvpGI_-I5SIrG0,3276
+ai_edge_torch/lowertools/_shim.py,sha256=rEZkOdHiz7CPvPL0WATIYnH4K6wF1YBtcv3oFEx2ZeQ,3277
 ai_edge_torch/lowertools/common_utils.py,sha256=4HQtquPZ6oiId8vR_1ykW_uK4ELnyo5zo3MlX1QYW4c,4513
 ai_edge_torch/lowertools/odml_torch_utils.py,sha256=QRuS7S5lULRWEh3J1sWIsnKh-rbX7rd9tt6JJHbMPfo,8317
 ai_edge_torch/lowertools/test_utils.py,sha256=mdxTlhqHABZEQ_GEmPFCL8LIAWtqRtYZUGdSY1ieZjw,1949
 ai_edge_torch/lowertools/torch_xla_utils.py,sha256=1EytIw2R6dthhLhf69wN1L9BaQTeybCD0wga-PhHcMI,9518
-ai_edge_torch/lowertools/translate_recipe.py,sha256=kUVCe69_DzvfbNYVB0MY2rCZwWaN8t3NoNu8Vh4x5bQ,5849
+ai_edge_torch/lowertools/translate_recipe.py,sha256=JNsRc1Jmpj5W6PBww8KRMkbtxcv7ssl8Rr1R3x5_7to,6283
 ai_edge_torch/odml_torch/__init__.py,sha256=S8jOzE9nLof-6es3XDiGJRN-9H_XTxsVm9dE7lD3RWo,812
 ai_edge_torch/odml_torch/_torch_future.py,sha256=jSYHf1CMTJzMizPMbu2b39hAt0ZTR6gQLq67GMe9KTo,2336
 ai_edge_torch/odml_torch/_torch_library.py,sha256=Lw1gqL2HWNRspdTwNhIkYAHDyafHedHtkXyKKxn-Wss,805
@@ -251,8 +251,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.dev20250508.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250508.dist-info/METADATA,sha256=GGDJl2Fya8gLr9RIfSLCmm1K1xA3qzBrrEOy1hwR2dQ,2051
-ai_edge_torch_nightly-0.5.0.dev20250508.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250508.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250508.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250510.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250510.dist-info/METADATA,sha256=1fA2DwpzLkPWBS-gV86ik7v9m39lO_RUaU4k7qAEvkM,2074
+ai_edge_torch_nightly-0.5.0.dev20250510.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250510.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250510.dist-info/RECORD,,