ai-edge-torch-nightly 0.5.0.dev20250511__py3-none-any.whl → 0.5.0.dev20250513__py3-none-any.whl

ai_edge_torch/generative/examples/amd_llama_135m/amd_llama_135m.py

@@ -49,7 +49,9 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
       intermediate_size=2048,
   )
-  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True
+  )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,

ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py

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

ai_edge_torch/generative/examples/openelm/convert_to_tflite.py

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

ai_edge_torch/generative/examples/openelm/openelm.py

@@ -51,7 +51,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
     The model config for an OpenELM model.
   """
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, enable_hlfb=True
   )
   num_heads = [12] * 4 + [16] * 14 + [20] * 12 + [24] * 6
   num_query_groups = [3] * 4 + [4] * 14 + [5] * 12 + [6] * 6

ai_edge_torch/generative/quantize/quant_recipes.py

@@ -27,37 +27,49 @@ Typical usage example:
   )
 """
 
+from typing import Optional
+from ai_edge_torch.generative.layers import model_config
 from ai_edge_torch.generative.quantize import quant_recipe
 from ai_edge_torch.generative.quantize import quant_recipe_utils
 from ai_edge_torch.quantize import quant_config
 
 
-def full_int8_dynamic_recipe() -> quant_config.QuantConfig:
+def full_int8_dynamic_recipe(
+    mcfg: Optional[model_config.ModelConfig] = None,
+) -> quant_config.QuantConfig:
   return quant_config.QuantConfig(
       generative_recipe=quant_recipe.GenerativeQuantRecipe(
           default=quant_recipe_utils.create_layer_quant_int8_dynamic(),
+          _model_config=mcfg,
       )
   )
 
 
-def full_int8_weight_only_recipe() -> quant_config.QuantConfig:
+def full_int8_weight_only_recipe(
+    mcfg: Optional[model_config.ModelConfig] = None,
+) -> quant_config.QuantConfig:
   return quant_config.QuantConfig(
       generative_recipe=quant_recipe.GenerativeQuantRecipe(
           default=quant_recipe_utils.create_layer_quant_int8_weight_only(),
+          _model_config=mcfg,
       )
   )
 
 
-def full_fp16_recipe() -> quant_config.QuantConfig:
+def full_fp16_recipe(
+    mcfg: Optional[model_config.ModelConfig] = None,
+) -> quant_config.QuantConfig:
   return quant_config.QuantConfig(
       generative_recipe=quant_recipe.GenerativeQuantRecipe(
-          default=quant_recipe_utils.create_layer_quant_fp16()
+          default=quant_recipe_utils.create_layer_quant_fp16(),
+          _model_config=mcfg,
       )
   )
 
 
 def all_supported_int4_dynamic_block_recipe(
     block_size: int,
+    mcfg: Optional[model_config.ModelConfig] = None,
 ) -> quant_config.QuantConfig:
   return quant_config.QuantConfig(
       generative_recipe=quant_recipe.GenerativeQuantRecipe(
@@ -65,5 +77,6 @@ def all_supported_int4_dynamic_block_recipe(
               block_size
           ),
           embedding=quant_recipe_utils.create_layer_quant_int8_dynamic(),
+          _model_config=mcfg,
       )
   )

ai_edge_torch/generative/utilities/converter.py

@@ -26,6 +26,7 @@ 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 import export_config
+from ai_edge_torch.quantize import quant_config as qcfg
 import torch
 
 ExportConfig = export_config.ExportConfig
@@ -123,7 +124,8 @@ def define_conversion_flags(
 
 def get_quant_recipe_from_flag(
     quantize: str,
-) -> Optional[quant_recipes.QuantizationRecipe]:
+    model_config: cfg.ModelConfig,
+) -> Optional[qcfg.QuantConfig]:
   """Processes the quantization flag and returns the corresponding recipe.
 
   Args:
@@ -139,15 +141,19 @@ def get_quant_recipe_from_flag(
     case QuantizationName.NONE:
       return None
     case QuantizationName.DYNAMIC_INT8:
-      return quant_recipes.full_int8_dynamic_recipe()
+      return quant_recipes.full_int8_dynamic_recipe(mcfg=model_config)
     case QuantizationName.WEIGHT_ONLY_INT8:
-      return quant_recipes.full_int8_weight_only_recipe()
+      return quant_recipes.full_int8_weight_only_recipe(mcfg=model_config)
     case QuantizationName.FP16:
       return quant_recipes.full_fp16_recipe()
     case QuantizationName.DYNAMIC_INT4_BLOCK32:
-      return quant_recipes.full_int4_dynamic_block_recipe(32)
+      return quant_recipes.all_supported_int4_dynamic_block_recipe(
+          32, mcfg=model_config
+      )
     case QuantizationName.DYNAMIC_INT4_BLOCK128:
-      return quant_recipes.full_int4_dynamic_block_recipe(128)
+      return quant_recipes.all_supported_int4_dynamic_block_recipe(
+          128, mcfg=model_config
+      )
     case _:
       raise ValueError(f'Unsupported quantization flag: {quantize}')
 
@@ -351,8 +357,7 @@ def _export_helper(
       kv_layout=export_config.kvcache_layout,
   )
 
-  quant_config = get_quant_recipe_from_flag(quantize)
-  quant_config._model_config = config
+  quant_config = get_quant_recipe_from_flag(quantize, config)
 
   # For export, we create a module that captures any non-exportable,
   # arugments, e.g. the generation config object.

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.5.0.dev20250511.dist-info → ai_edge_torch_nightly-0.5.0.dev20250513.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250511
+Version: 0.5.0.dev20250513
 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.dev20250511.dist-info → ai_edge_torch_nightly-0.5.0.dev20250513.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=hwgGQ5rNXzjaW8x5d5_q1vreBcKdC0qd0Sd_5QYRF_o,706
+ai_edge_torch/version.py,sha256=Q2u2GS0KjqxWhznlOZBgkCi4NAQcdpjJzkUYdcGYQ5o,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
@@ -52,8 +52,8 @@ ai_edge_torch/generative/custom_ops/bmm_4d.py,sha256=JmVbZCujG_wuBchma8QF3DSBfVc
 ai_edge_torch/generative/custom_ops/dynamic_update_slice.py,sha256=ZGAq2CfWZsfef5mHulsWmyUx0dDWJX6J6xPjhBrjQdM,2097
 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=z5MWiZLnsQzhNYMiQbcI9i0ki-dtkbimCptkiTFZxwo,1586
+ai_edge_torch/generative/examples/amd_llama_135m/amd_llama_135m.py,sha256=XsDXx6k0kE_OYu_dr7GEC26jCepV1Kv39iH-kpuqA4M,2794
+ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py,sha256=hiuMFJ8QPymGMM6PiSQqQrfR4M1mblpPuDfjjabcr_w,1560
 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=l0OrPGmX8WscuG9MIgtd0sqR4BeReNAu7fADzyPbnZw,1580
@@ -86,8 +86,8 @@ ai_edge_torch/generative/examples/moonshine/__init__.py,sha256=hHLluseD2R0Hh4W6X
 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=Hgp31zIQdJsTweRMr0U3d2SKW1h2nWnqWt1FlmuQqiI,1551
-ai_edge_torch/generative/examples/openelm/openelm.py,sha256=sIJ8Ie1oxFrJM-1jvv2ukiJbQOTIUGuMEZvmwZbt3n0,4556
+ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=wRdT7bWbCX8g4TbzKbjcLx6vmKtuT5-g-ipg19hJW-M,1525
+ai_edge_torch/generative/examples/openelm/openelm.py,sha256=hPcXYHj-nBP56TOeQQejB3HRzv6yHSftHOx0OEPP5M8,4574
 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
@@ -180,7 +180,7 @@ ai_edge_torch/generative/quantize/example.py,sha256=1lfVNUd2cEyRUnoZ7BLbRJ9IN-FT
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=plMsd7JBi98r2NHsAdMdvS6TPTXAoRFLCwOXu8H3-24,2004
 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=5UkUAT0qsWzLtNAeX-M5hEMi-kqoLV70_F76QiXmVZ4,2424
+ai_edge_torch/generative/quantize/quant_recipes.py,sha256=45DJfcQXZ1FA1qI4LgYoYE4UD4yvfIYoY9LgYTeKFVw,2845
 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
@@ -192,7 +192,7 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=vQWmpzMkJ2hP
 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=LrBqxXVxkOWh4abcHfY4QXRpYxjjfEYd4ifrpGGbebI,14441
+ai_edge_torch/generative/utilities/converter.py,sha256=OMBy_nos9mEGMQOAD8o0on-gAkRk-kliodFSTthD5BE,14612
 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=IG-88o7nWI9XrNDnwnQ-MoilsuqJ7KwrnbP3bn2EY9U,6334
@@ -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.dev20250511.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250511.dist-info/METADATA,sha256=PmlXdlLctkno1gMu-BWqW8CjHcSargbvVYhYycNMKTs,2074
-ai_edge_torch_nightly-0.5.0.dev20250511.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250511.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250511.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250513.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250513.dist-info/METADATA,sha256=qh5r3x7C0ksa3D2WriWd0yePFgxK8urh9aSsJCC_gjY,2074
+ai_edge_torch_nightly-0.5.0.dev20250513.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250513.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250513.dist-info/RECORD,,