ai-edge-torch-nightly 0.3.0.dev20240919__py3-none-any.whl → 0.3.0.dev20240920__py3-none-any.whl

ai_edge_torch/_convert/test/test_convert.py

@@ -23,12 +23,12 @@ from ai_edge_torch import config
 from ai_edge_torch._convert import conversion_utils
 from ai_edge_torch.testing import model_coverage
 import numpy as np
-import tensorflow as tf
 import torch
 from torch import nn
 import torchvision
 
 from absl.testing import absltest as googletest
+from ai_edge_litert import interpreter as tfl_interpreter  # pylint: disable=g-direct-tensorflow-import
 
 
 @dataclasses.dataclass
@@ -466,7 +466,9 @@ class TestConvert(googletest.TestCase):
     np.testing.assert_almost_equal(edge_output["y_data_2_0"], args[1])
     np.testing.assert_almost_equal(edge_output["y_data_2_1"], args[2])
 
-    interpreter = tf.lite.Interpreter(model_content=edge_model._tflite_model)
+    interpreter = tfl_interpreter.Interpreter(
+        model_content=edge_model._tflite_model
+    )
     runner = interpreter.get_signature_runner("serving_default")
     output_details = runner.get_output_details()
     self.assertIn("x", output_details.keys())
@@ -477,7 +479,9 @@ class TestConvert(googletest.TestCase):
   def _compare_tflite_torch_args_kwargs(self, model, args, kwargs, flat_inputs):
     model.eval()
     edge_model = ai_edge_torch.convert(model, args, kwargs)
-    interpreter = tf.lite.Interpreter(model_content=edge_model._tflite_model)
+    interpreter = tfl_interpreter.Interpreter(
+        model_content=edge_model._tflite_model
+    )
     runner = interpreter.get_signature_runner("serving_default")
     input_details = runner.get_input_details()
     self.assertEqual(input_details.keys(), flat_inputs.keys())

ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py

@@ -30,17 +30,17 @@ _CHECKPOINT_PATH = flags.DEFINE_string(
 )
 _TFLITE_PATH = flags.DEFINE_string(
     'tflite_path',
-    '/tmp/gemma2_q8_seq512_ekv1024.tflite',
+    '/tmp/',
     'The tflite file path to export.',
 )
 _PREFILL_SEQ_LEN = flags.DEFINE_integer(
     'prefill_seq_len',
-    512,
+    1024,
     'The maximum size of prefill input tensor.',
 )
 _KV_CACHE_MAX_LEN = flags.DEFINE_integer(
     'kv_cache_max_len',
-    1024,
+    1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
 _QUANTIZE = flags.DEFINE_bool(
@@ -54,9 +54,11 @@ def main(_):
   pytorch_model = gemma2.build_2b_model(
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'gemma2_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
-      tflite_path=_TFLITE_PATH.value,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
       quantize=_QUANTIZE.value,
   )

ai_edge_torch/generative/examples/gemma/convert_to_tflite.py

@@ -30,17 +30,17 @@ _CHECKPOINT_PATH = flags.DEFINE_string(
 )
 _TFLITE_PATH = flags.DEFINE_string(
     'tflite_path',
-    '/tmp/gemma_q8_seq512_ekv1024.tflite',
+    '/tmp/',
     'The tflite file path to export.',
 )
 _PREFILL_SEQ_LEN = flags.DEFINE_integer(
     'prefill_seq_len',
-    512,
+    1024,
     'The maximum size of prefill input tensor.',
 )
 _KV_CACHE_MAX_LEN = flags.DEFINE_integer(
     'kv_cache_max_len',
-    1024,
+    1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
 _QUANTIZE = flags.DEFINE_bool(
@@ -54,9 +54,11 @@ def main(_):
   pytorch_model = gemma.build_2b_model(
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'gemma_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
-      tflite_path=_TFLITE_PATH.value,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
       quantize=_QUANTIZE.value,
   )

ai_edge_torch/generative/examples/openelm/convert_to_tflite.py

@@ -30,17 +30,17 @@ _CHECKPOINT_PATH = flags.DEFINE_string(
 )
 _TFLITE_PATH = flags.DEFINE_string(
     'tflite_path',
-    '/tmp/openelm_q8_seq512_ekv1024.tflite',
+    '/tmp/',
     'The tflite file path to export.',
 )
 _PREFILL_SEQ_LEN = flags.DEFINE_integer(
     'prefill_seq_len',
-    512,
+    1024,
     'The maximum size of prefill input tensor.',
 )
 _KV_CACHE_MAX_LEN = flags.DEFINE_integer(
     'kv_cache_max_len',
-    1024,
+    1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
 _QUANTIZE = flags.DEFINE_bool(
@@ -54,9 +54,11 @@ def main(_):
   pytorch_model = openelm.build_model(
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'openelm_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
-      tflite_path=_TFLITE_PATH.value,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
       quantize=_QUANTIZE.value,
   )

ai_edge_torch/generative/examples/openelm/verify.py

@@ -33,8 +33,10 @@ _PROMPTS = flags.DEFINE_multi_string(
 def main(_):
   checkpoint = "apple/OpenELM-3B"
   verifier.log_msg("Loading the original model from", checkpoint)
-  original_model = transformers.AutoModelForCausalLM.from_pretrained(
-      checkpoint, trust_remote_code=True
+  wrapper_model = verifier.ModelWrapper(
+      model=transformers.AutoModelForCausalLM.from_pretrained(
+          checkpoint, trust_remote_code=True
+      ),
   )
 
   # Locate the cached dir.
@@ -50,7 +52,7 @@ def main(_):
   tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_checkpoint)
 
   verifier.verify_reauthored_model(
-      original_model=original_model,
+      original_model=wrapper_model,
       reauthored_model=reauthored_model,
       tokenizer=tokenizer,
       prompts=_PROMPTS.value,

ai_edge_torch/generative/examples/phi/convert_to_tflite.py

@@ -30,17 +30,17 @@ _CHECKPOINT_PATH = flags.DEFINE_string(
 )
 _TFLITE_PATH = flags.DEFINE_string(
     'tflite_path',
-    '/tmp/phi2_q8_seq512_ekv1024.tflite',
+    '/tmp/',
     'The tflite file path to export.',
 )
 _PREFILL_SEQ_LEN = flags.DEFINE_integer(
     'prefill_seq_len',
-    512,
+    1024,
     'The maximum size of prefill input tensor.',
 )
 _KV_CACHE_MAX_LEN = flags.DEFINE_integer(
     'kv_cache_max_len',
-    1024,
+    1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
 _QUANTIZE = flags.DEFINE_bool(
@@ -54,9 +54,11 @@ def main(_):
   pytorch_model = phi2.build_model(
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'phi2_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
-      tflite_path=_TFLITE_PATH.value,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
       quantize=_QUANTIZE.value,
   )

ai_edge_torch/generative/examples/phi/verify.py

@@ -24,15 +24,25 @@ import transformers
 
 _PROMPTS = flags.DEFINE_multi_string(
     "prompts",
-    "What is the meaning of life?",
+    "Instruct: Write an email about the weather Output:",
     "The input prompts to generate answers.",
 )
 
+_MAX_NEW_TOKENS = flags.DEFINE_integer(
+    "max_new_tokens",
+    30,
+    "The maximum size of the generated tokens.",
+)
 
 def main(_):
   checkpoint = kagglehub.model_download("Microsoft/phi/transformers/2")
   verifier.log_msg("Loading the original model from", checkpoint)
-  original_model = transformers.AutoModelForCausalLM.from_pretrained(checkpoint)
+  generation_config = transformers.GenerationConfig.from_pretrained(checkpoint)
+  generation_config.max_new_tokens = _MAX_NEW_TOKENS.value
+  wrapper_model = verifier.ModelWrapper(
+      model=transformers.AutoModelForCausalLM.from_pretrained(checkpoint),
+      hf_generation_config=generation_config,
+  )
 
   verifier.log_msg("Building the reauthored model from", checkpoint)
   reauthored_model = phi2.build_model(checkpoint)
@@ -41,7 +51,7 @@ def main(_):
   tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
 
   verifier.verify_reauthored_model(
-      original_model=original_model,
+      original_model=wrapper_model,
       reauthored_model=reauthored_model,
       tokenizer=tokenizer,
       prompts=_PROMPTS.value,

ai_edge_torch/generative/examples/smollm/convert_to_tflite.py

@@ -30,17 +30,17 @@ _CHECKPOINT_PATH = flags.DEFINE_string(
 )
 _TFLITE_PATH = flags.DEFINE_string(
     'tflite_path',
-    '/tmp/smollm_q8_seq512_ekv1024.tflite',
+    '/tmp/',
     'The tflite file path to export.',
 )
 _PREFILL_SEQ_LEN = flags.DEFINE_integer(
     'prefill_seq_len',
-    512,
+    1024,
     'The maximum size of prefill input tensor.',
 )
 _KV_CACHE_MAX_LEN = flags.DEFINE_integer(
     'kv_cache_max_len',
-    1024,
+    1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
 _QUANTIZE = flags.DEFINE_bool(
@@ -54,9 +54,11 @@ def main(_):
   pytorch_model = smollm.build_model(
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'smollm_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
-      tflite_path=_TFLITE_PATH.value,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
       quantize=_QUANTIZE.value,
   )

ai_edge_torch/generative/examples/smollm/verify.py

@@ -33,8 +33,9 @@ _PROMPTS = flags.DEFINE_multi_string(
 def main(_):
   checkpoint = "HuggingFaceTB/SmolLM-135M"
   verifier.log_msg("Loading the original model from", checkpoint)
-  original_model = transformers.AutoModelForCausalLM.from_pretrained(checkpoint)
-
+  wrapper_model = verifier.ModelWrapper(
+      model=transformers.AutoModelForCausalLM.from_pretrained(checkpoint),
+  )
   # Locate the cached dir.
   cached_config_file = transformers.utils.cached_file(
       checkpoint, transformers.utils.CONFIG_NAME
@@ -47,7 +48,7 @@ def main(_):
   tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
 
   verifier.verify_reauthored_model(
-      original_model=original_model,
+      original_model=wrapper_model,
       reauthored_model=reauthored_model,
       tokenizer=tokenizer,
       prompts=_PROMPTS.value,

ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py

@@ -30,17 +30,17 @@ _CHECKPOINT_PATH = flags.DEFINE_string(
 )
 _TFLITE_PATH = flags.DEFINE_string(
     'tflite_path',
-    '/tmp/tiny_llama_q8_seq512_ekv1024.tflite',
+    '/tmp/',
     'The tflite file path to export.',
 )
 _PREFILL_SEQ_LEN = flags.DEFINE_integer(
     'prefill_seq_len',
-    512,
+    1024,
     'The maximum size of prefill input tensor.',
 )
 _KV_CACHE_MAX_LEN = flags.DEFINE_integer(
     'kv_cache_max_len',
-    1024,
+    1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
 _QUANTIZE = flags.DEFINE_bool(
@@ -54,9 +54,11 @@ def main(_):
   pytorch_model = tiny_llama.build_model(
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'tinyllama_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
-      tflite_path=_TFLITE_PATH.value,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
       quantize=_QUANTIZE.value,
   )

ai_edge_torch/generative/examples/tiny_llama/verify.py

@@ -33,10 +33,11 @@ _PROMPTS = flags.DEFINE_multi_string(
 def main(_):
   checkpoint = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
   verifier.log_msg("Loading the original model from", checkpoint)
-  original_model = transformers.AutoModelForCausalLM.from_pretrained(
-      checkpoint, trust_remote_code=True
+  wrapper_model = verifier.ModelWrapper(
+      model=transformers.AutoModelForCausalLM.from_pretrained(
+          checkpoint, trust_remote_code=True
+      ),
   )
-
   # Locate the cached dir.
   cached_config_file = transformers.utils.cached_file(
       checkpoint, transformers.utils.CONFIG_NAME
@@ -49,7 +50,7 @@ def main(_):
   tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
 
   verifier.verify_reauthored_model(
-      original_model=original_model,
+      original_model=wrapper_model,
       reauthored_model=reauthored_model,
       tokenizer=tokenizer,
       prompts=_PROMPTS.value,

ai_edge_torch/generative/test/test_model_conversion.py

@@ -25,7 +25,7 @@ import numpy as np
 import torch
 
 from absl.testing import absltest as googletest
-from tensorflow.lite.python import interpreter
+from ai_edge_litert import interpreter
 
 
 class TestModelConversion(googletest.TestCase):

ai_edge_torch/generative/test/test_model_conversion_large.py

@@ -28,7 +28,7 @@ import numpy as np
 import torch
 
 from absl.testing import absltest as googletest
-from tensorflow.lite.python import interpreter
+from ai_edge_litert import interpreter
 
 
 class TestModelConversion(googletest.TestCase):

ai_edge_torch/generative/utilities/verifier.py

@@ -16,17 +16,66 @@
 """Common utility functions to verify the reauthored models."""
 
 import datetime
-from typing import List
+from typing import List, Optional, Union
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import numpy as np
 import torch
+import transformers
 
 
 def log_msg(*args):
   print("[%s]" % datetime.datetime.now(), *args)
 
 
+class ModelWrapper(torch.nn.Module):
+  """A wrapper for the model to be verified, this could be a HuggingFace model
+
+  or a regular PyTorch model.
+  """
+
+  def __init__(
+      self,
+      model: torch.nn.Module,
+      model_format: str = "huggingface",
+      hf_generation_config: Optional[transformers.GenerationConfig] = None,
+  ):
+    """Initializes the wrapper.
+
+    Args:
+      model (torch.nn.Module): The original model. This could be a model built
+        from HuggingFace transformers, or a regular PyTorch model.
+      model_format (str): The format of the model. It should be either
+        "huggingface" or "pytorch".
+      hf_generation_config (transformers.GenerationConfig): The HuggingFace
+        generation config. This config will only be used if the underlying model
+        is built from HuggingFace transformers.
+    """
+    super().__init__()
+    self.model = model
+    self.model_format = model_format
+    self.hf_generation_config = hf_generation_config
+
+  def generate(
+      self, inputs: torch.Tensor
+  ) -> Union[transformers.utils.ModelOutput, torch.LongTensor]:
+    if self.model_format == "huggingface":
+      return self.model.generate(
+          inputs=inputs, generation_config=self.hf_generation_config
+      )
+    else:
+      raise NotImplementedError(
+          "generate() is not implemented for model format: %s"
+          % self.model_format
+      )
+
+  def forward(
+      self,
+      inputs: torch.Tensor,
+  ):
+    return self.model.forward(inputs)
+
+
 def forward(
     model: torch.nn.Module,
     tokens: torch.Tensor,
@@ -75,7 +124,7 @@ def generate(
 
 
 def verify_with_input_ids(
-    original_model: torch.nn.Module,
+    original_model: ModelWrapper,
     reauthored_model: torch.nn.Module,
     input_ids: torch.Tensor = torch.from_numpy(np.array([[1, 2, 3, 4]])).int(),
     kv_cache_max_len: int = 1024,
@@ -87,7 +136,7 @@ def verify_with_input_ids(
   It compares only one outputs from the original and the reauthored model.
 
   Args:
-    original_model (torch.nn.Module): The original model.
+    original_model (ModelWrapper): The original model.
     reauthored_model (torch.nn.Module): The model reauthored with ai_edge_torch
       Generative API.
     input_ids (torch.Tensor): The input token IDs to forward.
@@ -119,7 +168,7 @@ def verify_with_input_ids(
 
 
 def verify_model_with_prompts(
-    original_model: torch.nn.Module,
+    original_model: ModelWrapper,
     reauthored_model: torch.nn.Module,
     tokenizer: torch.nn.Module,
     prompts: str,
@@ -130,7 +179,7 @@ def verify_model_with_prompts(
   original and the reauthored model.
 
   Args:
-    original_model (torch.nn.Module): The original model.
+    original_model (ModelWrapper): The original model.
     reauthored_model (torch.nn.Module): The model reauthored with ai_edge_torch
       Generative API.
     tokenizer (torch.nn.Module): The tokenizer.
@@ -156,7 +205,7 @@ def verify_model_with_prompts(
 
 
 def verify_reauthored_model(
-    original_model: torch.nn.Module,
+    original_model: ModelWrapper,
     reauthored_model: torch.nn.Module,
     tokenizer: torch.nn.Module,
     prompts: List[str],
@@ -174,7 +223,7 @@ def verify_reauthored_model(
   It prints out "PASS" or "FAILED" to the console.
 
   Args:
-    original_model (torch.nn.Module): The original model.
+    original_model (ModelWrapper): The original model.
     reauthored_model (torch.nn.Module): The model reauthored with ai_edge_torch
       Generative API.
     tokenizer (torch.nn.Module): The tokenizer.

ai_edge_torch/model.py

@@ -27,6 +27,8 @@ from typing import Callable
 import numpy.typing as npt
 import tensorflow as tf
 
+from ai_edge_litert import interpreter as tfl_interpreter  # pylint: disable=g-direct-tensorflow-import
+
 DEFAULT_SIGNATURE_NAME = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
 
 
@@ -65,7 +67,7 @@ class TfLiteModel(Model):
       tflite_model: A TFlite serialized object.
     """
     self._tflite_model = tflite_model
-    self._interpreter_builder = lambda: tf.lite.Interpreter(
+    self._interpreter_builder = lambda: tfl_interpreter.Interpreter(
         model_content=self._tflite_model,
         experimental_default_delegate_latest_features=True,
     )
@@ -75,12 +77,13 @@ class TfLiteModel(Model):
     return self._tflite_model
 
   def set_interpreter_builder(
-      self, builder: Callable[[], tf.lite.Interpreter]
+      self, builder: Callable[[], tfl_interpreter.Interpreter]
   ) -> None:
     """Sets a custom interpreter builder.
 
     Args:
-      builder: A function that returns a `tf.lite.Interpreter` or its subclass.
+      builder: A function that returns a `tfl_interpreter.Interpreter` or its
+        subclass.
     """
     self._interpreter_builder = builder
 
@@ -166,7 +169,7 @@ class TfLiteModel(Model):
 
     # Check if this is indeed a tflite model:
     try:
-      interpreter = tf.lite.Interpreter(model_content=model_content)
+      interpreter = tfl_interpreter.Interpreter(model_content=model_content)
       interpreter.get_signature_list()
     except:
       return None

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.3.0.dev20240919"
+__version__ = "0.3.0.dev20240920"

{ai_edge_torch_nightly-0.3.0.dev20240919.dist-info → ai_edge_torch_nightly-0.3.0.dev20240920.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20240919
+Version: 0.3.0.dev20240920
 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.3.0.dev20240919.dist-info → ai_edge_torch_nightly-0.3.0.dev20240920.dist-info}/RECORD

@@ -2,8 +2,8 @@ ai_edge_torch/__init__.py,sha256=48qP37uHT90YPs4eIUQxCiWVwqGEX3idCUs6mQKvX1U,116
 ai_edge_torch/config.py,sha256=FMWeCH2b7HYILBvaI1iZNnYCO4WAhDOwBZBmIE-xrF0,909
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/fx_pass_base.py,sha256=D86Gw3pIRcpnTebUPKlnPbPGJae1S6Fw4DZZ3ZkD0zw,3730
-ai_edge_torch/model.py,sha256=NYV6Mkaje_ditIEI_s_7nLP_-8i4kbGM8nRzieVkbUI,5397
-ai_edge_torch/version.py,sha256=N5hYc9s2RU44J1_oe0UfJhTFo0d4JvMlKvxNlYtK0GI,706
+ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
+ai_edge_torch/version.py,sha256=-oH0R07HZpydzqltOWclHB1dbcc4VycTlZcnDYtS89g,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=5uPwHhmc6kwiIz-CqaiHDejf2SOWMHrb-rYEHm69wKc,3801
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -25,7 +25,7 @@ ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitio
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py,sha256=L_x8BrF7UDah-SYl-pG11I6CIckdU9kBTUHcmwW4cts,2420
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=mzfL9cf0qBnpmxM_OlMQFvQsEZV2B_Mia9yEJV4J7rI,7135
 ai_edge_torch/_convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/test/test_convert.py,sha256=FSufFZEeTLBpUnzE1Iy-LvNN0mhDynWMNg7Mei8RpLQ,14973
+ai_edge_torch/_convert/test/test_convert.py,sha256=40QRxQFNeSRr4dLXJkzG-wKUlvJtsfv62cdvRrmBv5w,15097
 ai_edge_torch/_convert/test/test_convert_composites.py,sha256=BCIODgxMI_3MxMLfNWYMGjcz-al-J3z5eDHCiZJXNwY,7992
 ai_edge_torch/_convert/test/test_convert_multisig.py,sha256=6_C2R9--KyNR7_oezZIAfyTSR97tOeEWy4XGcbSxBDE,5778
 ai_edge_torch/_convert/test/test_to_channel_last_io.py,sha256=1o-gUiwzIuO67FNAJ8DeyKv8fVUeZVNNNwofNVDjYeU,3024
@@ -39,22 +39,22 @@ ai_edge_torch/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrK
 ai_edge_torch/generative/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=09VbyWErOMP9BXGwZpwvqzN5RaOqRigsELfxNRVeWns,2024
-ai_edge_torch/generative/examples/gemma/convert_to_tflite.py,sha256=qJKQu6lKuSVhn8JR7KUeInq0u6yqgxEi7hfKCrZrIqY,2019
+ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=RZDs6oY-NLYrPNtfuJDweIHzGUL2kzpIc3AW_1p8gGg,2186
+ai_edge_torch/generative/examples/gemma/convert_to_tflite.py,sha256=t8Qg10obnEzeoMeyHnZhyNBN7G85SGy-au8Y8nehq8E,2181
 ai_edge_torch/generative/examples/gemma/gemma.py,sha256=hjpSPzEjPHuxwRJ-vHHtCCf2PSTnm30Mp0ajYYtDivo,7489
 ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=gCLOti-4xHunjphNBbx9St6faRteSakm8Oex6R1Xek0,10272
 ai_edge_torch/generative/examples/openelm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=HnqP3te1Qvy4SKaaqPrsG05eojiKDJShp4H3jPC9tYg,2023
+ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=85FVEt6cKFP2UzCLC78tAkbwGlGhAArtG7Wa75NxJik,2185
 ai_edge_torch/generative/examples/openelm/openelm.py,sha256=gGkHELNrt4xqnu11fCh3sJbZ7OsPyvoiF1J1aKCs5r8,7532
-ai_edge_torch/generative/examples/openelm/verify.py,sha256=2qFdyLfcefdA3s1KQ-ZGWo4XReMXkEQAvpUEyJE5iqM,2057
+ai_edge_torch/generative/examples/openelm/verify.py,sha256=BvK4c8jodQBy2l3NnvCjlBB0qaA7EYwPNKklvFR4k_o,2103
 ai_edge_torch/generative/examples/phi/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=viIkbAgknE3zxavTZtib87cMIG2_-jJXtxJPcmB2pGQ,2007
+ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=3go690yX6PFeXMdpY7y4JZorAwxX0HT_b_pKZieauvk,2169
 ai_edge_torch/generative/examples/phi/phi2.py,sha256=YwAszA53aOjvaMJ5wua2-5rP79N21Un_Y5yBCfFSYNU,6189
-ai_edge_torch/generative/examples/phi/verify.py,sha256=R9BjOArnn-3svoIApmP1NwO47n8KIFikOF0_MEgTOa4,1770
+ai_edge_torch/generative/examples/phi/verify.py,sha256=5bKONolW8JIsQAzMHIvh_OSytoJVVJqDZEcxjhciFnI,2136
 ai_edge_torch/generative/examples/smollm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/smollm/convert_to_tflite.py,sha256=86hvBleyFXWmwy3Ke5J7x7WcCtG20D2kiBNrodE0R4w,2017
+ai_edge_torch/generative/examples/smollm/convert_to_tflite.py,sha256=zPrDTDeRVWFi9DS32uNi-RLpzOStFOk5MhNla4ixeew,2179
 ai_edge_torch/generative/examples/smollm/smollm.py,sha256=hyhMk-b5762Q2xmjdD47g85dcbBSNJXNPIsifm1DRto,3239
-ai_edge_torch/generative/examples/smollm/verify.py,sha256=JzidfVMMFDXzDdwn7ToDPuMo6eaoENNZGpEzX3f61Jk,1976
+ai_edge_torch/generative/examples/smollm/verify.py,sha256=wsoy3CaHZhrdJjkJJYir7xxxwgCvLprMnh8QxT0hEkc,2013
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=kDWG6MlIGa89zC5KSRcJlw2c4ITuw8KcchtfmF55f4g,3545
 ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=tL6w2dr6VP66IXjSKo9StDNP-wl0RO3fh6dIliiYlFA,4656
@@ -78,9 +78,9 @@ ai_edge_torch/generative/examples/test_models/__init__.py,sha256=hHLluseD2R0Hh4W
 ai_edge_torch/generative/examples/test_models/toy_model.py,sha256=QyLeCqDnk71WvvFH68g9UeF-HytonSk1ItGF9dc7Zj8,5854
 ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=e_Kqm5dStSrNE9_aIYC-vYJRsqLn-hJVkmR4QjYqZI0,5913
 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=Yg5G1LePoryeTib35lqICqaDW6foLUzSRgwJ2FlklIw,2040
+ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=ekxd8efjMgEvauUu3PidWOC-DszPHn5sqU753F7sJIM,2201
 ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=tlWpa7Aun3u3w5b-9EBtW7olhmSf8W-tn5bKUIwC-ys,6044
-ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=jld5PlGOQXMIWc1WoDYL_1nnsoVzRfrg-WgnsxRgaEU,2041
+ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=27oBf706_AKX7amfp2THF9J0G3AUEEecGaXv025idKA,2086
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=jrzCB3ZyY_t5jJM1e2Czdt3DjAIL43R0_a-T-I7wOzw,1155
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=hhxSQvkDMv0isZJhmuLiod66ZODaJ8uSPSVTJVHBabQ,1931
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -107,8 +107,8 @@ ai_edge_torch/generative/quantize/supported_schemes.py,sha256=FjdycEOvxRgBmQdZVu
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/test/test_kv_cache.py,sha256=W6Bh0gYDzmwb0j9HdD5_D7Z7FPToP2HSyFrmwIXuFqo,3793
 ai_edge_torch/generative/test/test_loader.py,sha256=8y74ChO3CZCfEi1eCf3-w47kRgAI4qPYCXpi8rTQXMA,3378
-ai_edge_torch/generative/test/test_model_conversion.py,sha256=DBlqxW2IT-dZYzEfOMAp86Wtqiu6kgSWZ9BKZR1Clrw,5467
-ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=dUYFarOldejqbMpa0j0vIDvXlWPAancuI8di3XkGxm8,4498
+ai_edge_torch/generative/test/test_model_conversion.py,sha256=s-EVLOQGjIeVtgNI8Ggs37pkRdErAliT6NhrrFigPOE,5459
+ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=PtePuBqVMLjxq2cDIIXXqaz7zsn3R19oilFyIVJRFi8,4490
 ai_edge_torch/generative/test/test_quantize.py,sha256=8geJhKwYBU20m0mdGPD1BUFwQ0lZKNtCB04SOLO18y4,5980
 ai_edge_torch/generative/test/utils.py,sha256=YvEhO2HIj1LkBs5du1UxY-cGRW9HMyAYsOUhgsTrTpA,1796
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
@@ -116,7 +116,7 @@ ai_edge_torch/generative/utilities/converter.py,sha256=MQUg2ZLmfk_2csWmQWKD_II0b
 ai_edge_torch/generative/utilities/loader.py,sha256=b9iotIhVDX-Zc9XjIDUaLxnV395AyBnkQe3dV5YA7Co,13297
 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/verifier.py,sha256=QAv1uJdI5o1yfphr_DpzxhZswKa4VG3JZUpqbCCWKMk,7114
+ai_edge_torch/generative/utilities/verifier.py,sha256=7DoYtkilz4wjWnXfdydIGNgTG1udZIydFxdbpIcKbMQ,8625
 ai_edge_torch/hlfb/__init__.py,sha256=sH4um75na-O8tzxN6chFyp6Y4xnexsE7kUQpZySv6dE,735
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=cjTprggj_cuktSCm7-A25e7Shop3k63ylp7sdZmtZ8o,4790
 ai_edge_torch/hlfb/mark_pattern/passes.py,sha256=pjkKcI1nHECPluAt87cFBrt1DP0f3ge7rHq1NhCkBIE,1936
@@ -163,8 +163,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 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.3.0.dev20240919.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20240919.dist-info/METADATA,sha256=NkHYIOMz-5DNKJuSQ8wE-3Nz1R6a9YZ59M-Nq8sAnJg,1859
-ai_edge_torch_nightly-0.3.0.dev20240919.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20240919.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20240919.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20240920.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20240920.dist-info/METADATA,sha256=m60oD-H8W2EMVolDGw02tMYcKDrotaTaLtsZwzr_Kyk,1859
+ai_edge_torch_nightly-0.3.0.dev20240920.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20240920.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20240920.dist-info/RECORD,,