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

ai_edge_torch/generative/examples/tiny_llama/verify.py

@@ -0,0 +1,62 @@
+# 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.
+# ==============================================================================
+
+"""Verifies the reauthored TinyLlama-1.1B model."""
+
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "Show me the program to add 2 and 3.",
+    "The input prompts to generate answers.",
+)
+
+
+def main(_):
+  checkpoint = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+  verifier.log_msg("Loading the original model from", checkpoint)
+  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
+  )
+  reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  verifier.log_msg("Building the reauthored model from", reauthored_checkpoint)
+  reauthored_model = tiny_llama.build_model(reauthored_checkpoint)
+
+  verifier.log_msg("Loading the tokenizer from", checkpoint)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=wrapper_model,
+      reauthored_model=reauthored_model,
+      tokenizer=tokenizer,
+      prompts=_PROMPTS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/layers/builder.py

@@ -75,7 +75,9 @@ def build_norm(dim: int, config: cfg.NormalizationConfig):
         zero_centered_gamma=config.zero_centered,
     )
   elif config.type == cfg.NormalizationType.LAYER_NORM:
-    return normalization.LayerNorm(dim, config.epsilon, config.enable_hlfb)
+    return normalization.LayerNorm(
+        dim, config.epsilon, config.enable_hlfb, config.use_input_shape
+    )
   elif config.type == cfg.NormalizationType.GROUP_NORM:
     return normalization.GroupNorm(
         config.group_num, dim, config.epsilon, config.enable_hlfb

ai_edge_torch/generative/layers/model_config.py

@@ -69,6 +69,9 @@ class NormalizationConfig:
   zero_centered: bool = False
   # Number of groups used in group normalization.
   group_num: Optional[float] = None
+  # Whether to use the input shape to determine the dimension of normalization
+  # when type is LAYER_NORM.
+  use_input_shape: bool = True
 
 
 @dataclass

ai_edge_torch/generative/layers/normalization.py

@@ -78,7 +78,7 @@ class GroupNorm(torch.nn.Module):
       group_num (int): Number of groups to separate the channels into.
       dim (int): Dimension of the input tensor.
       eps (float): A small float value to ensure numerical stability (default:
-        1e-6).
+        1e-5).
       enable_hlfb (bool): Whether to convert this normalization into a single
         op.
     """
@@ -112,7 +112,13 @@ class GroupNorm(torch.nn.Module):
 
 class LayerNorm(torch.nn.Module):
 
-  def __init__(self, dim: int, eps: float = 1e-5, enable_hlfb: bool = False):
+  def __init__(
+      self,
+      dim: int,
+      eps: float = 1e-5,
+      enable_hlfb: bool = False,
+      use_input_shape: bool = True,
+  ):
     """Initialize the LayerNorm layer.
 
     Args:
@@ -121,9 +127,12 @@ class LayerNorm(torch.nn.Module):
         1e-6).
       enable_hlfb (bool): Whether to convert this normalization into a single
         op.
+      use_input_shape (bool): Whether to use the input shape to determine the
+        dimension of normalization (default: True).
     """
     super().__init__()
     self.enable_hlfb = enable_hlfb
+    self.use_input_shape = use_input_shape
     self.eps = eps
     self.weight = torch.nn.Parameter(torch.ones(dim))
     self.bias = torch.nn.Parameter(torch.ones(dim))
@@ -139,19 +148,18 @@ class LayerNorm(torch.nn.Module):
     """
     if self.enable_hlfb:
       return layer_norm_with_hlfb(
-          x,
-          self.weight,
-          self.bias,
-          self.eps,
+          x, self.weight, self.bias, self.eps, self.use_input_shape
       )
+
+    if self.use_input_shape:
+      normalized_shape = x.shape
+      weight = self.weight.broadcast_to(x.shape)
+      bias = self.bias.broadcast_to(x.shape)
     else:
-      return F.layer_norm(
-          x,
-          x.shape,
-          self.weight.broadcast_to(x.shape),
-          self.bias.broadcast_to(x.shape),
-          self.eps,
-      )
+      normalized_shape = self.weight.shape
+      weight = self.weight
+      bias = self.bias
+    return F.layer_norm(x, normalized_shape, weight, bias, self.eps)
 
 
 def group_norm_with_hlfb(
@@ -193,6 +201,7 @@ def layer_norm_with_hlfb(
     w: torch.Tensor,
     b: torch.Tensor,
     eps: float,
+    use_input_shape: bool,
 ):
   """Layer Normalization with high-level function boundary enabled.
 
@@ -201,18 +210,20 @@ def layer_norm_with_hlfb(
     w (torch.Tensor): The weight tensor for the normalization.
     b (torch.Tensor): The bias tensor for the normalization.
     eps (float): A small float value to ensure numerical stability.
+    use_input_shape (bool): Whether to use the input shape to determine the
+      dimension of normalization.
 
   Returns:
     The output tensor of Layer Normalization.
   """
   builder = StableHLOCompositeBuilder(name="odml.layer_norm", attr={"eps": eps})
   x, w, b = builder.mark_inputs(x, w, b)
-  y = F.layer_norm(
-      x,
-      x.shape,
-      weight=w.broadcast_to(x.shape),
-      bias=b.broadcast_to(x.shape),
-      eps=eps,
-  )
+  if use_input_shape:
+    normalized_shape = x.shape
+    w = w.broadcast_to(x.shape)
+    b = b.broadcast_to(x.shape)
+  else:
+    normalized_shape = w.shape
+  y = F.layer_norm(x, normalized_shape, w, b, eps=eps)
   y = builder.mark_outputs(y)
   return y

ai_edge_torch/generative/layers/scaled_dot_product_attention.py

@@ -119,15 +119,25 @@ def scaled_dot_product_attention_with_hlfb(
     # Handle the GQA case, where q.shape[1] % k.shape[1] == 0.
     k = k.repeat_interleave(q.shape[1] // k.shape[1], dim=1)
     v = v.repeat_interleave(q.shape[1] // v.shape[1], dim=1)
-  y = F.scaled_dot_product_attention(
-      q,
-      k,
-      v,
-      attn_mask=mask,
-      dropout_p=0.0,
-      is_causal=mask is None,
-      scale=scale,
-  )
+  if softcap is None:
+    y = F.scaled_dot_product_attention(
+        q,
+        k,
+        v,
+        attn_mask=mask,
+        dropout_p=0.0,
+        is_causal=mask is None,
+        scale=scale,
+    )
+  else:
+    q.mul_(scale)
+    scores = q @ k.transpose(-1, -2)
+    scores = scores / softcap
+    scores = torch.tanh(scores)
+    scores = scores * softcap
+    scores = scores + mask
+    out = F.softmax(scores.float(), dim=-1).type_as(q)
+    y = torch.matmul(out, v)
 
   result = y.transpose(1, 2)
   result = builder.mark_outputs(result)

ai_edge_torch/generative/layers/unet/blocks_2d.py

@@ -41,22 +41,22 @@ class ResidualBlock2D(nn.Module):
     )
     self.conv_1 = nn.Conv2d(
         config.in_channels,
-        config.out_channels,
+        config.hidden_channels,
         kernel_size=3,
         stride=1,
         padding=1,
     )
     if config.time_embedding_channels is not None:
       self.time_emb_proj = nn.Linear(
-          config.time_embedding_channels, config.out_channels
+          config.time_embedding_channels, config.hidden_channels
       )
     else:
       self.time_emb_proj = None
     self.norm_2 = layers_builder.build_norm(
-        config.out_channels, config.normalization_config
+        config.hidden_channels, config.normalization_config
     )
     self.conv_2 = nn.Conv2d(
-        config.out_channels,
+        config.hidden_channels,
         config.out_channels,
         kernel_size=3,
         stride=1,
@@ -391,6 +391,7 @@ class DownEncoderBlock2D(nn.Module):
           ResidualBlock2D(
               unet_cfg.ResidualBlock2DConfig(
                   in_channels=input_channels,
+                  hidden_channels=config.out_channels,
                   out_channels=config.out_channels,
                   time_embedding_channels=config.time_embedding_channels,
                   normalization_config=config.normalization_config,
@@ -492,6 +493,7 @@ class UpDecoderBlock2D(nn.Module):
           ResidualBlock2D(
               unet_cfg.ResidualBlock2DConfig(
                   in_channels=input_channels,
+                  hidden_channels=config.out_channels,
                   out_channels=config.out_channels,
                   time_embedding_channels=config.time_embedding_channels,
                   normalization_config=config.normalization_config,
@@ -602,6 +604,7 @@ class SkipUpDecoderBlock2D(nn.Module):
           ResidualBlock2D(
               unet_cfg.ResidualBlock2DConfig(
                   in_channels=resnet_in_channels + res_skip_channels,
+                  hidden_channels=config.out_channels,
                   out_channels=config.out_channels,
                   time_embedding_channels=config.time_embedding_channels,
                   normalization_config=config.normalization_config,
@@ -706,6 +709,7 @@ class MidBlock2D(nn.Module):
         ResidualBlock2D(
             unet_cfg.ResidualBlock2DConfig(
                 in_channels=config.in_channels,
+                hidden_channels=config.in_channels,
                 out_channels=config.in_channels,
                 time_embedding_channels=config.time_embedding_channels,
                 normalization_config=config.normalization_config,
@@ -724,6 +728,7 @@ class MidBlock2D(nn.Module):
           ResidualBlock2D(
               unet_cfg.ResidualBlock2DConfig(
                   in_channels=config.in_channels,
+                  hidden_channels=config.in_channels,
                   out_channels=config.in_channels,
                   time_embedding_channels=config.time_embedding_channels,
                   normalization_config=config.normalization_config,

ai_edge_torch/generative/layers/unet/model_config.py

@@ -48,6 +48,7 @@ class DownSamplingConfig:
 @dataclasses.dataclass
 class ResidualBlock2DConfig:
   in_channels: int
+  hidden_channels: int
   out_channels: int
   normalization_config: layers_cfg.NormalizationConfig
   activation_config: layers_cfg.ActivationConfig

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):
@@ -96,7 +96,7 @@ class TestModelConversion(googletest.TestCase):
   def test_gemma2(self):
     config = gemma2.get_fake_model_config()
     pytorch_model = gemma2.Gemma2(config).eval()
-    self._test_model(config, pytorch_model, "prefill", atol=1e-1, rtol=1e-3)
+    self._test_model(config, pytorch_model, "prefill", atol=1e-4, rtol=1e-5)
 
   @googletest.skipIf(
       ai_edge_config.Config.use_torch_xla,

ai_edge_torch/generative/utilities/stable_diffusion_loader.py

@@ -412,6 +412,7 @@ class BaseLoader(loader.ModelLoader):
   ):
     residual_block_config = unet_config.ResidualBlock2DConfig(
         in_channels=config.in_channels,
+        hidden_channels=config.in_channels,
         out_channels=config.in_channels,
         time_embedding_channels=config.time_embedding_channels,
         normalization_config=config.normalization_config,
@@ -466,6 +467,7 @@ class BaseLoader(loader.ModelLoader):
           f"{converted_state_param_prefix}.resnets.{i}",
           unet_config.ResidualBlock2DConfig(
               in_channels=input_channels,
+              hidden_channels=config.out_channels,
               out_channels=config.out_channels,
               time_embedding_channels=config.time_embedding_channels,
               normalization_config=config.normalization_config,
@@ -508,6 +510,7 @@ class BaseLoader(loader.ModelLoader):
           f"{converted_state_param_prefix}.resnets.{i}",
           unet_config.ResidualBlock2DConfig(
               in_channels=input_channels,
+              hidden_channels=config.out_channels,
               out_channels=config.out_channels,
               time_embedding_channels=config.time_embedding_channels,
               normalization_config=config.normalization_config,
@@ -554,6 +557,7 @@ class BaseLoader(loader.ModelLoader):
           f"{converted_state_param_prefix}.resnets.{i}",
           unet_config.ResidualBlock2DConfig(
               in_channels=resnet_in_channels + res_skip_channels,
+              hidden_channels=config.out_channels,
               out_channels=config.out_channels,
               time_embedding_channels=config.time_embedding_channels,
               normalization_config=config.normalization_config,

ai_edge_torch/generative/utilities/verifier.py

@@ -0,0 +1,249 @@
+# 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.
+# ==============================================================================
+
+"""Common utility functions to verify the reauthored models."""
+
+import datetime
+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,
+    kv_cache: kv_utils.KVCache,
+) -> tuple[torch.Tensor, kv_utils.KVCache]:
+  """Forwards the model reauthored with ai_edge_torch Generative API.
+
+  Args:
+    model (torch.nn.Module): The model to forward. It should be a model built
+      with ai_edge_torch Generative API.
+    tokens (torch.Tensor): The input tokens to forward.
+    kv_cache (KVCache): The KV cache to forward.
+
+  Returns:
+    The output logits and the updated KV cache.
+  """
+  input_pos = torch.arange(0, tokens.shape[1], dtype=torch.int)
+  output = model.forward(tokens, input_pos, kv_cache)
+  return output["logits"], output["kv_cache"]
+
+
+def generate(
+    model: torch.nn.Module, prompts: torch.Tensor, response_len: int
+) -> torch.Tensor:
+  """Generates the response to the prompts.
+
+  It appends tokens output by the model to the prompts and feeds them back to
+  the model up to decode_len.
+
+  Args:
+    model (torch.nn.Module): The model to generate. It should be a model built
+      with ai_edge_torch Generative API.
+    prompts (torch.Tensor): The prompts to generate.
+    response_len (int): The number of tokens to generate.
+
+  Returns:
+    The generated tokens.
+  """
+  input_ids = prompts[0].int().tolist()
+  kv_cache = kv_utils.KVCache.from_model_config(model.config)
+  for _ in range(response_len - len(input_ids)):
+    logits, kv_cache = forward(model, torch.tensor([input_ids]), kv_cache)
+    generated_token = logits[0][-1].argmax().item()
+    input_ids.append(generated_token)
+  return torch.tensor([input_ids])
+
+
+def verify_with_input_ids(
+    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,
+    rtol: float = 1e-05,
+    atol: float = 1e-05,
+) -> bool:
+  """Verifies if the model reauthored generates the same output of the oringal.
+
+  It compares only one outputs from the original and the reauthored model.
+
+  Args:
+    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.
+    kv_cache_max_len (int): The maximum sequence length of the KV cache.
+    rtol (float): The relative tolerance for the comparison.
+    atol (float): The absolute tolerance for the comparison.
+
+  Returns:
+    True if the model reauthored generates the same output of the original.
+  """
+  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
+  input_ids_len = input_ids.shape[1]
+  tokens[0, :input_ids_len] = input_ids
+
+  log_msg("Forwarding the original model...")
+  outputs_original = original_model.forward(tokens)
+  logits_original = outputs_original.logits[0, input_ids_len - 1, :]
+  log_msg("logits_original: ", logits_original)
+
+  log_msg("Forwarding the reauthored model...")
+  kv_cache = kv_utils.KVCache.from_model_config(reauthored_model.config)
+  outputs_reauthored = forward(reauthored_model, tokens, kv_cache)
+  logits_reauthored = outputs_reauthored[0][0, input_ids_len - 1, :]
+  log_msg("logits_reauthored:", logits_reauthored)
+
+  return torch.allclose(
+      logits_original, logits_reauthored, rtol=rtol, atol=atol
+  )
+
+
+def verify_model_with_prompts(
+    original_model: ModelWrapper,
+    reauthored_model: torch.nn.Module,
+    tokenizer: torch.nn.Module,
+    prompts: str,
+) -> bool:
+  """Verifies if the model reauthored generates the same answer of the oringal.
+
+  It compares an answer, i.e. multiple continuous outputs generated by the
+  original and the reauthored model.
+
+  Args:
+    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.
+    prompts (str): The input prompts to generate answers.
+
+  Returns:
+    True if the model reauthored generates the same answer of the original.
+  """
+  prompt_tokens = tokenizer.encode(prompts, return_tensors="pt")
+
+  log_msg("Generating answer with the original model...")
+  outputs_original = original_model.generate(prompt_tokens)
+  response_original = tokenizer.decode(outputs_original[0])
+  log_msg("outputs_from_original_model: [[", response_original, "]]")
+
+  log_msg("Generating answer with the reauthored model...")
+  generate_len = len(outputs_original[0])
+  outputs_reauthored = generate(reauthored_model, prompt_tokens, generate_len)
+  response_reauthored = tokenizer.decode(outputs_reauthored[0])
+  log_msg("outputs from reauthored model: [[", response_reauthored, "]]")
+
+  return response_original == response_reauthored
+
+
+def verify_reauthored_model(
+    original_model: ModelWrapper,
+    reauthored_model: torch.nn.Module,
+    tokenizer: torch.nn.Module,
+    prompts: List[str],
+    rtol: float = 1e-05,
+    atol: float = 1e-05,
+):
+  """Verifies the reauthored model against the original model.
+
+  It verifies the reauthored model with two methods:
+  1. It compares the output of the original and the reauthored model with an
+     arbitrary input.
+  2. It compares the answer generated by the original and the reauthored model
+     with a prompt.
+
+  It prints out "PASS" or "FAILED" to the console.
+
+  Args:
+    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.
+    prompts (List[str]): List of the input prompts to generate answers.
+    rtol (float): The relative tolerance for the comparison.
+    atol (float): The absolute tolerance for the comparison.
+  """
+  log_msg("Verifying the reauthored model with an arbitrary input...")
+  if verify_with_input_ids(
+      original_model, reauthored_model, rtol=rtol, atol=atol
+  ):
+    log_msg("PASS")
+  else:
+    log_msg("FAILED")
+
+  for p in prompts:
+    log_msg("Verifying the reauthored model with prompts:", p)
+    if verify_model_with_prompts(
+        original_model, reauthored_model, tokenizer, p
+    ):
+      log_msg("PASS")
+    else:
+      log_msg("FAILED")

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.dev20240918"
+__version__ = "0.3.0.dev20240920"

{ai_edge_torch_nightly-0.3.0.dev20240918.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.dev20240918
+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