ai-edge-torch-nightly 0.3.0.dev20240829__py3-none-any.whl → 0.3.0.dev20240901__py3-none-any.whl

ai_edge_torch/generative/examples/experimental/gemma/gemma.py

@@ -159,15 +159,44 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   return config
 
 
-def get_fake_model_config_2b_for_test(**kwargs) -> cfg.ModelConfig:
-  config = get_model_config_2b(**kwargs)
-  config.num_layers = 2
+def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
+  attn_config = cfg.AttentionConfig(
+      num_heads=8,
+      head_dim=256,
+      num_query_groups=1,
+      rotary_percentage=1.0,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_TANH),
+      intermediate_size=128,
+  )
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM,
+      epsilon=1e-6,
+      zero_centered=True,
+  )
+  config = cfg.ModelConfig(
+      vocab_size=128,
+      num_layers=2,
+      max_seq_len=2 * kv_cache_max_len,
+      embedding_dim=2048,
+      kv_cache_max_len=kv_cache_max_len,
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+      final_norm_config=norm_config,
+      parallel_residual=False,
+      lm_head_use_bias=False,
+      enable_hlfb=True,
+  )
   return config
 
 
 def build_2b_model(checkpoint_path, test_model=False, **kwargs) -> nn.Module:
   config = (
-      get_fake_model_config_2b_for_test(**kwargs)
+      get_fake_model_config(**kwargs)
       if test_model
       else get_model_config_2b(**kwargs)
   )

ai_edge_torch/generative/examples/gemma/gemma.py

@@ -147,9 +147,39 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   return config
 
 
-def get_fake_model_config_2b_for_test() -> cfg.ModelConfig:
-  config = get_model_config_2b()
-  config.num_layers = 2
+# TODO(b/363021962): Clean up this part to streamline fake model config generation.
+def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
+  attn_config = cfg.AttentionConfig(
+      num_heads=8,
+      head_dim=256,
+      num_query_groups=1,
+      rotary_percentage=1.0,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_TANH),
+      intermediate_size=128,
+  )
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM,
+      epsilon=1e-6,
+      zero_centered=True,
+  )
+  config = cfg.ModelConfig(
+      vocab_size=128,
+      num_layers=2,
+      max_seq_len=2 * kv_cache_max_len,
+      embedding_dim=2048,
+      kv_cache_max_len=kv_cache_max_len,
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+      final_norm_config=norm_config,
+      parallel_residual=False,
+      lm_head_use_bias=False,
+      enable_hlfb=True,
+  )
   return config
 
 

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -209,9 +209,47 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   return config
 
 
-def get_fake_model_config_2b_for_test() -> cfg.ModelConfig:
-  config = get_model_config_2b()
-  config.num_layers = 2
+def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
+  attn_config = cfg.AttentionConfig(
+      num_heads=4,
+      head_dim=64,
+      num_query_groups=4,
+      rotary_percentage=1.0,
+      qkv_transpose_before_split=True,
+      logit_softcap=50.0,
+      sliding_window_size=64,
+      attn_types=[cfg.AttentionType.GLOBAL, cfg.AttentionType.LOCAL_SLIDING]
+      * 13,
+  )
+
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM,
+      epsilon=1e-6,
+      zero_centered=True,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_TANH),
+      intermediate_size=128,
+      pre_ff_norm_config=norm_config,
+      post_ff_norm_config=norm_config,
+  )
+  config = cfg.ModelConfig(
+      vocab_size=128,
+      num_layers=2,
+      max_seq_len=2 * kv_cache_max_len,
+      embedding_dim=128,
+      kv_cache_max_len=kv_cache_max_len,
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+      final_norm_config=norm_config,
+      parallel_residual=False,
+      lm_head_use_bias=False,
+      enable_hlfb=True,
+      final_logit_softcap=30.0,
+  )
   return config
 
 

ai_edge_torch/generative/examples/phi2/phi2.py

@@ -139,9 +139,36 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   return config
 
 
-def get_fake_model_config_for_test() -> cfg.ModelConfig:
-  config = get_model_config()
-  config.num_layers = 2
+def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
+  attn_config = cfg.AttentionConfig(
+      num_heads=16,
+      head_dim=80,
+      num_query_groups=4,
+      rotary_percentage=0.4,
+      qkv_use_bias=True,
+      output_proj_use_bias=True,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.SEQUENTIAL,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_TANH),
+      intermediate_size=128,
+      use_bias=True,
+  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.LAYER_NORM)
+  config = cfg.ModelConfig(
+      vocab_size=128,
+      num_layers=2,
+      max_seq_len=2 * kv_cache_max_len,
+      kv_cache_max_len=kv_cache_max_len,
+      embedding_dim=128,
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      final_norm_config=norm_config,
+      parallel_residual=True,
+      lm_head_use_bias=True,
+      enable_hlfb=True,
+  )
   return config
 
 

ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py

@@ -137,11 +137,11 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   return config
 
 
-def get_fake_model_config_for_test() -> cfg.ModelConfig:
+def get_fake_model_config() -> cfg.ModelConfig:
   config = get_model_config()
   config.vocab_size = 128
   config.num_layers = 2
-  config.ff_config.intermediate_size = 256
+  config.ff_config.intermediate_size = 64
   return config
 
 

ai_edge_torch/generative/quantize/example.py

@@ -22,7 +22,7 @@ import torch
 
 def main():
   # Build a PyTorch model as usual
-  config = gemma.get_fake_model_config_2b_for_test()
+  config = gemma.get_fake_model_config()
   model = gemma.Gemma(config)
   idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
   tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")

ai_edge_torch/generative/test/test_model_conversion.py

@@ -16,6 +16,7 @@
 import copy
 
 import ai_edge_torch
+from ai_edge_torch import config as ai_edge_config
 from ai_edge_torch.generative.examples.gemma import gemma, gemma2
 from ai_edge_torch.generative.examples.phi2 import phi2
 from ai_edge_torch.generative.examples.test_models import toy_model_with_kv_cache  # NOQA
@@ -25,11 +26,27 @@ import numpy as np
 import torch
 
 from absl.testing import absltest as googletest
+from tensorflow.lite.python import interpreter
 
 
 class TestModelConversion(googletest.TestCase):
   """Unit tests that check for model conversion and correctness."""
 
+  def setUp(self):
+    super().setUp()
+    # Builder function for an Interpreter that supports custom ops.
+    self._interpreter_builder = (
+        lambda tflite_model: lambda: interpreter.InterpreterWithCustomOps(
+            custom_op_registerers=["GenAIOpsRegisterer"],
+            model_content=tflite_model,
+            experimental_default_delegate_latest_features=True,
+        )
+    )
+
+  @googletest.skipIf(
+      ai_edge_config.Config.use_torch_xla,
+      reason="tests with custom ops are not supported on oss",
+  )
   def test_toy_model_with_kv_cache(self):
     config = toy_model_with_kv_cache.get_model_config()
     pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config).eval()
@@ -38,22 +55,27 @@ class TestModelConversion(googletest.TestCase):
     )
 
     edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
+    edge_model.set_interpreter_builder(
+        self._interpreter_builder(edge_model.tflite_model())
+    )
 
-    # TODO: b/338288901 - re-enable test to check output tensors.
-    skip_output_check = True
-    if not skip_output_check:
-      self.assertTrue(
-          model_coverage.compare_tflite_torch(
-              edge_model,
-              pytorch_model,
-              (idx, input_pos),
-              num_valid_inputs=1,
-              atol=1e-5,
-              rtol=1e-5,
-          )
-      )
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model,
+            pytorch_model,
+            (idx, input_pos),
+            num_valid_inputs=1,
+            atol=1e-5,
+            rtol=1e-5,
+        )
+    )
 
+  @googletest.skipIf(
+      ai_edge_config.Config.use_torch_xla,
+      reason="tests with custom ops are not supported on oss",
+  )
   def test_toy_model_with_multi_batches(self):
+    self.skipTest("b/362842043")
     config = toy_model_with_kv_cache.get_model_config()
     config.batch_size = 2
     pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config).eval()
@@ -62,21 +84,25 @@ class TestModelConversion(googletest.TestCase):
     )
 
     edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
+    edge_model.set_interpreter_builder(
+        self._interpreter_builder(edge_model.tflite_model())
+    )
 
-    # TODO: b/338288901 - re-enable test to check output tensors.
-    skip_output_check = True
-    if not skip_output_check:
-      self.assertTrue(
-          model_coverage.compare_tflite_torch(
-              edge_model,
-              pytorch_model,
-              (idx, input_pos),
-              num_valid_inputs=1,
-              atol=1e-5,
-              rtol=1e-5,
-          )
-      )
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model,
+            pytorch_model,
+            (idx, input_pos),
+            num_valid_inputs=1,
+            atol=1e-5,
+            rtol=1e-5,
+        )
+    )
 
+  @googletest.skipIf(
+      ai_edge_config.Config.use_torch_xla,
+      reason="tests with custom ops are not supported on oss",
+  )
   def test_toy_model_with_kv_cache_with_hlfb(self):
     config = toy_model_with_kv_cache.get_model_config()
     config.enable_hlfb = True
@@ -86,49 +112,27 @@ class TestModelConversion(googletest.TestCase):
     )
 
     edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
+    edge_model.set_interpreter_builder(
+        self._interpreter_builder(edge_model.tflite_model())
+    )
 
-    # TODO: b/338288901 - re-enable test to check output tensors.
-    skip_output_check = True
-    if not skip_output_check:
-      self.assertTrue(
-          model_coverage.compare_tflite_torch(
-              edge_model,
-              pytorch_model,
-              (idx, input_pos),
-              num_valid_inputs=1,
-              atol=1e-5,
-              rtol=1e-5,
-          )
-      )
-
-  def test_tiny_llama(self):
-    self.skipTest("b/338288901")
-    config = tiny_llama.get_fake_model_config_for_test()
-    pytorch_model = tiny_llama.TinyLLamma(config).eval()
-
-    idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-    tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
-    tokens[0, :4] = idx
-    input_pos = torch.arange(0, 10)
-
-    edge_model = ai_edge_torch.convert(pytorch_model, (tokens, input_pos))
-
-    # TODO: b/338288901 - re-enable test to check output tensors.
-    skip_output_check = True
-    if not skip_output_check:
-      self.assertTrue(
-          model_coverage.compare_tflite_torch(
-              edge_model,
-              pytorch_model,
-              (tokens, input_pos),
-              num_valid_inputs=1,
-              atol=1e-5,
-              rtol=1e-5,
-          )
-      )
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model,
+            pytorch_model,
+            (idx, input_pos),
+            num_valid_inputs=1,
+            atol=1e-5,
+            rtol=1e-5,
+        )
+    )
 
+  @googletest.skipIf(
+      ai_edge_config.Config.use_torch_xla,
+      reason="tests with custom ops are not supported on oss",
+  )
   def test_tiny_llama_multisig(self):
-    config = tiny_llama.get_fake_model_config_for_test()
+    config = tiny_llama.get_fake_model_config()
     pytorch_model = tiny_llama.TinyLLamma(config).eval()
 
     # prefill
@@ -149,22 +153,25 @@ class TestModelConversion(googletest.TestCase):
         .signature("decode", pytorch_model, (decode_token, decode_input_pos))
         .convert()
     )
+    edge_model.set_interpreter_builder(
+        self._interpreter_builder(edge_model.tflite_model())
+    )
 
-    # TODO: b/338288901 - re-enable test to check output tensors.
-    skip_output_check = True
-    if not skip_output_check:
-      copied_model = copy.deepcopy(pytorch_model)
+    copied_model = copy.deepcopy(pytorch_model)
 
-      self.assertTrue(
-          model_coverage.compare_tflite_torch(
-              edge_model,
-              pytorch_model,
-              (prefill_tokens, prefill_input_pos),
-              signature_name="prefill",
-              num_valid_inputs=1,
-          )
-      )
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model,
+            pytorch_model,
+            (prefill_tokens, prefill_input_pos),
+            signature_name="prefill",
+            num_valid_inputs=1,
+        )
+    )
 
+    # TODO(b/362840003): figure why this decode output has big numerical diff.
+    skip_output_check = True
+    if not skip_output_check:
       self.assertTrue(
           model_coverage.compare_tflite_torch(
               edge_model,
@@ -175,87 +182,6 @@ class TestModelConversion(googletest.TestCase):
           )
       )
 
-  def test_gemma(self):
-    self.skipTest("b/338288901")
-    config = gemma.get_fake_model_config_2b_for_test()
-    model = gemma.Gemma(config)
-
-    idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-    tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
-    tokens[0, :4] = idx
-    input_pos = torch.arange(0, 10)
-
-    edge_model = ai_edge_torch.convert(model, (tokens, input_pos))
-
-    # TODO: b/338288901 - re-enable test to check output tensors.
-    skip_output_check = True
-    if not skip_output_check:
-      # TODO(talumbau, haoliang): debug numerical diff.
-      self.assertTrue(
-          model_coverage.compare_tflite_torch(
-              edge_model,
-              model,
-              (tokens, input_pos),
-              num_valid_inputs=1,
-              atol=1e-2,
-              rtol=1e-5,
-          )
-      )
-
-  def test_gemma2(self):
-    self.skipTest("b/338288901")
-    config = gemma2.get_fake_model_config_2b_for_test()
-    model = gemma2.Gemma2(config)
-    model.eval()
-
-    idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-    tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
-    tokens[0, :4] = idx
-    input_pos = torch.arange(0, 10)
-
-    edge_model = ai_edge_torch.convert(model, (tokens, input_pos))
-
-    # TODO: b/338288901 - re-enable test to check output tensors.
-    skip_output_check = True
-    if not skip_output_check:
-      # TODO(talumbau, haoliang): debug numerical diff.
-      self.assertTrue(
-          model_coverage.compare_tflite_torch(
-              edge_model,
-              model,
-              (tokens, input_pos),
-              num_valid_inputs=1,
-              atol=1e-2,
-              rtol=1e-5,
-          )
-      )
-
-  def test_phi2(self):
-    self.skipTest("b/338288901")
-    config = phi2.get_fake_model_config_for_test()
-    pytorch_model = phi2.Phi2(config).eval()
-
-    idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-    tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
-    tokens[0, :4] = idx
-    input_pos = torch.arange(0, 10)
-
-    edge_model = ai_edge_torch.convert(pytorch_model, (tokens, input_pos))
-
-    # TODO: b/338288901 - re-enable test to check output tensors.
-    skip_output_check = True
-    if not skip_output_check:
-      self.assertTrue(
-          model_coverage.compare_tflite_torch(
-              edge_model,
-              pytorch_model,
-              (tokens, input_pos),
-              num_valid_inputs=1,
-              atol=1e-5,
-              rtol=1e-5,
-          )
-      )
-
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_torch/generative/test/test_model_conversion_large.py

@@ -0,0 +1,139 @@
+# 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.
+# ==============================================================================
+# Testing model conversion for a few gen-ai models.
+import copy
+
+import ai_edge_torch
+from ai_edge_torch import config as ai_edge_config
+from ai_edge_torch.generative.examples.gemma import gemma, gemma2
+from ai_edge_torch.generative.examples.phi2 import phi2
+from ai_edge_torch.generative.examples.test_models import toy_model_with_kv_cache  # NOQA
+from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
+from ai_edge_torch.testing import model_coverage
+import numpy as np
+import torch
+
+from absl.testing import absltest as googletest
+from tensorflow.lite.python import interpreter
+
+
+class TestModelConversion(googletest.TestCase):
+  """Unit tests that check for model conversion and correctness."""
+
+  def setUp(self):
+    super().setUp()
+    # Builder function for an Interpreter that supports custom ops.
+    self._interpreter_builder = (
+        lambda tflite_model: lambda: interpreter.InterpreterWithCustomOps(
+            custom_op_registerers=["GenAIOpsRegisterer"],
+            model_content=tflite_model,
+            experimental_default_delegate_latest_features=True,
+        )
+    )
+
+  @googletest.skipIf(
+      ai_edge_config.Config.use_torch_xla,
+      reason="tests with custom ops are not supported on oss",
+  )
+  def test_gemma(self):
+    config = gemma.get_fake_model_config()
+    model = gemma.Gemma(config)
+
+    idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
+    tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
+    tokens[0, :4] = idx
+    input_pos = torch.arange(0, 10)
+
+    edge_model = ai_edge_torch.convert(model, (tokens, input_pos))
+    edge_model.set_interpreter_builder(
+        self._interpreter_builder(edge_model.tflite_model())
+    )
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model,
+            model,
+            (tokens, input_pos),
+            num_valid_inputs=1,
+            atol=1e-2,
+            rtol=1e-5,
+        )
+    )
+
+  @googletest.skipIf(
+      ai_edge_config.Config.use_torch_xla,
+      reason="tests with custom ops are not supported on oss",
+  )
+  def test_gemma2(self):
+    config = gemma2.get_fake_model_config()
+    model = gemma2.Gemma2(config)
+    model.eval()
+
+    idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
+    tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
+    tokens[0, :4] = idx
+    input_pos = torch.arange(0, 10)
+
+    edge_model = ai_edge_torch.convert(model, (tokens, input_pos))
+    edge_model.set_interpreter_builder(
+        self._interpreter_builder(edge_model.tflite_model())
+    )
+
+    # TODO(b/362840003): debug numerical diff.
+    skip_output_check = True
+    if not skip_output_check:
+      self.assertTrue(
+          model_coverage.compare_tflite_torch(
+              edge_model,
+              model,
+              (tokens, input_pos),
+              num_valid_inputs=1,
+              atol=1e-2,
+              rtol=1e-5,
+          )
+      )
+
+  @googletest.skipIf(
+      ai_edge_config.Config.use_torch_xla,
+      reason="tests with custom ops are not supported on oss",
+  )
+  def test_phi2(self):
+    config = phi2.get_fake_model_config()
+    pytorch_model = phi2.Phi2(config).eval()
+
+    idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
+    tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
+    tokens[0, :4] = idx
+    input_pos = torch.arange(0, 10)
+
+    edge_model = ai_edge_torch.convert(pytorch_model, (tokens, input_pos))
+    edge_model.set_interpreter_builder(
+        self._interpreter_builder(edge_model.tflite_model())
+    )
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model,
+            pytorch_model,
+            (tokens, input_pos),
+            num_valid_inputs=1,
+            atol=1e-3,
+            rtol=1e-3,
+        )
+    )
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_torch/model.py

@@ -22,6 +22,7 @@ from __future__ import annotations
 
 import abc
 import re
+from typing import Callable
 
 import numpy.typing as npt
 import tensorflow as tf
@@ -64,6 +65,24 @@ class TfLiteModel(Model):
       tflite_model: A TFlite serialized object.
     """
     self._tflite_model = tflite_model
+    self._interpreter_builder = lambda: tf.lite.Interpreter(
+        model_content=self._tflite_model,
+        experimental_default_delegate_latest_features=True,
+    )
+
+  def tflite_model(self) -> bytes:
+    """Returns the wrapped tflite model."""
+    return self._tflite_model
+
+  def set_interpreter_builder(
+      self, builder: Callable[[], tf.lite.Interpreter]
+  ) -> None:
+    """Sets a custom interpreter builder.
+
+    Args:
+      builder: A function that returns a `tf.lite.Interpreter` or its subclass.
+    """
+    self._interpreter_builder = builder
 
   def __call__(
       self,
@@ -80,10 +99,7 @@ class TfLiteModel(Model):
       signature_name: The name of the signature to be used for inference. The
         default signature is used if not provided.
     """
-    interpreter = tf.lite.Interpreter(
-        model_content=self._tflite_model,
-        experimental_default_delegate_latest_features=True,
-    )
+    interpreter = self._interpreter_builder()
     interpreter.allocate_tensors()
 
     signature_list = interpreter.get_signature_list()

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20240829.dist-info → ai_edge_torch_nightly-0.3.0.dev20240901.dist-info}/METADATA

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

@@ -1,8 +1,8 @@
 ai_edge_torch/__init__.py,sha256=48qP37uHT90YPs4eIUQxCiWVwqGEX3idCUs6mQKvX1U,1168
 ai_edge_torch/config.py,sha256=PCd9PVrbUNeVIUDFUCnW4goDWU4bjouK28yMYU6VOi0,877
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
-ai_edge_torch/model.py,sha256=7tox6sdFIlCYPLDYpjFcD8cPTSivURCL_VV6-Dt5Sfc,4910
-ai_edge_torch/version.py,sha256=OF9oSdUOGcmdEp2HSZmEIeCPlRhL3cpviHc_dExhcX8,706
+ai_edge_torch/model.py,sha256=NYV6Mkaje_ditIEI_s_7nLP_-8i4kbGM8nRzieVkbUI,5397
+ai_edge_torch/version.py,sha256=llxFt4Jrb5-zJ8uUalWlL75tHHddPLBl0Nyk0B7ecZU,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=kcv_QgNgeyDmrqwdzHicGNP68w6zF7GJg7YkMEIXp4Q,3759
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -42,7 +42,7 @@ ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQe
 ai_edge_torch/generative/examples/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/experimental/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/experimental/gemma/convert_to_tflite.py,sha256=lpiPFSh3SJd6WwuZ0QegSva3__iSz2tUD7L7QfkAe4I,3085
-ai_edge_torch/generative/examples/experimental/gemma/gemma.py,sha256=8313wSsddvuxZ5ZYVdaITBV2FF1k22dcCujnq0UZvKs,6699
+ai_edge_torch/generative/examples/experimental/gemma/gemma.py,sha256=EdElPCDLYxnNvkPMJkE3WKvESze1ehgShEk2NnbrXLg,7527
 ai_edge_torch/generative/examples/experimental/phi/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/experimental/phi/convert_to_tflite.py,sha256=DavrdGmqUgoThsGNRv3LXMW5tvJdYEvj66Hf1XRqkXU,3055
 ai_edge_torch/generative/examples/experimental/phi/phi2.py,sha256=u-VJX5mjzQKspXtAhNi53LCITtag-3nCaRTKdk5Z1sc,6231
@@ -52,11 +52,11 @@ ai_edge_torch/generative/examples/experimental/tiny_llama/tiny_llama.py,sha256=z
 ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=pseJExH35lSAK0ZtzSHB1sFtRtF_EuT2xcSpGU0gKVI,2524
 ai_edge_torch/generative/examples/gemma/convert_to_tflite.py,sha256=w589IJETATd6Z9_1XCIWbrlCV3E92X_5ac3VVCVFXG0,2522
-ai_edge_torch/generative/examples/gemma/gemma.py,sha256=cCki-0cKvmGxK4Md6dRNdPDWZUyhkJUI854OCTFf3h0,6262
-ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=q9Zil66EvRKrSpLVQHxKHu_8NL0HAgY2FbtThoTZVUY,8226
+ai_edge_torch/generative/examples/gemma/gemma.py,sha256=pzD9dYUYg8E6fFACh-8B8G9NHFXOVEWBjf5aDeipU2s,7202
+ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=ypd6uBb4FgDpuWm_w8JNYBAf4eFxWbYccs8vCgBhi-I,9374
 ai_edge_torch/generative/examples/phi2/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/phi2/convert_to_tflite.py,sha256=ON6zLO-nFS8eJ2yhyWzT5x2Somr-Ca-VjpjT7OGFU10,2506
-ai_edge_torch/generative/examples/phi2/phi2.py,sha256=C_kFYsPrEQ9GJCnc6h-jh8B5qQryvEpI6O6t4FBxg1I,5858
+ai_edge_torch/generative/examples/phi2/phi2.py,sha256=91mWxEtKgDtUhCAewWNwH_UOOCzy6tPdf6LNRlxZhrc,6700
 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=0WniBWQ6_NcQc5WycX3YRRX7Os9AGQSxfc1m2HKBqg8,4479
@@ -82,7 +82,7 @@ ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.p
 ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=mQkcpSe6HlRLMkIRCEHc9ZXL7jxEp9RWSGUQjjd-r2w,4841
 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=CLRqO7ycMbpy7J3_Czp1sLx6hcdwGD9zVq04yRba0e8,2550
-ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=mXXFYJfo8yegSOFOndCR0oYxFPchYb9vTJ4ThXGIFLU,5940
+ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=JmwU1sniO37vnCFc8dklbd-0ofTZK0PaBv_Ksn1Vq6M,5930
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=fmNNXawJ722M4cTUuTx289rT0NHxBEsOy_k8baqCOms,1173
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=sXis0U4u-RoIp_NyrmWJNnqFqpqRuZOrhfsJIO6rMps,2028
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -100,7 +100,7 @@ ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=4a0wh0co8Avz1wvxS3Xqsgr
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
 ai_edge_torch/generative/layers/unet/model_config.py,sha256=NvBJj09a7ZC-ChGE_ex-_kLnE_fjzrY6txbLSh1pMKA,9208
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/quantize/example.py,sha256=mqi3zFUp4w198DGnRkmZCWUZdUXTkvg1_tdTdOk9IkA,1535
+ai_edge_torch/generative/quantize/example.py,sha256=Bmc-WowIJIfDgt84CNw2LhyLRi7SFcw8BQEu4byTKJU,1523
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=n1Fm8BFC8gJa_oiwwAOOghJyHtOXYZ4q-5ZRy4pHrIw,1957
 ai_edge_torch/generative/quantize/quant_recipe.py,sha256=tKnuJq6hPD23JPCB9nPAlE1UHAwdbChkgPShiVaz4CE,5156
 ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=4fgmP_GgeiFUOkIaC9ZZXC12eO3DQZdrWDXRz5YXiwU,2270
@@ -111,7 +111,8 @@ ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/translate_recipe.py,sha
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/test/test_experimental_ekv.py,sha256=8qv_eVtJW9GPvBEf2hPQe3tpdJ33XShya6MCX1FqrZM,4355
 ai_edge_torch/generative/test/test_loader.py,sha256=_y5EHGgoNOmCuYonsB81UJScHVsTAQXUVd44czMAw6k,3379
-ai_edge_torch/generative/test/test_model_conversion.py,sha256=-RBTQSERP4szm8s8s_WRmGF3mWZA5E2w2QNtl2MqORw,8475
+ai_edge_torch/generative/test/test_model_conversion.py,sha256=wQLVjMnKHBCVCU_I-xAUZvlOFoDiwYwKQDvCZ2mjtOM,6193
+ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=o3l7HFHP-sg8aHeLNTSpMF91YovPODjp4QzYUnSJiIE,4479
 ai_edge_torch/generative/test/test_quantize.py,sha256=JEsk9SAkHK0SFm44K_quISc5yBBS6yvtBP1MDyFHdFw,5344
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
 ai_edge_torch/generative/utilities/loader.py,sha256=QFZ2lkeoYQ9MZ1CAFVxBHG4OT192SH74UtJCvbDsdeI,12727
@@ -161,8 +162,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.dev20240829.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20240829.dist-info/METADATA,sha256=LrexCNdY177vrp17WaGa53bxHH9vuZXT64O5by4HE6Y,1878
-ai_edge_torch_nightly-0.3.0.dev20240829.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20240829.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20240829.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20240901.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20240901.dist-info/METADATA,sha256=6d0enZRdCzGdIp1MfpBysngGttpwRTkPilaM5jXDN2g,1878
+ai_edge_torch_nightly-0.3.0.dev20240901.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20240901.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20240901.dist-info/RECORD,,