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

ai_edge_torch/generative/examples/openelm/openelm.py

@@ -15,16 +15,12 @@
 
 """Example of building an OpenELM model."""
 
-import os
-import pathlib
-
 from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
-import numpy as np
 import torch
 from torch import nn
 
@@ -210,28 +206,3 @@ def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
   loader.load(model, strict=False)
   model.eval()
   return model
-
-
-def define_and_run(checkpoint_path: str) -> None:
-  """Instantiates and runs an OpenELM model."""
-
-  current_dir = pathlib.Path(__file__).parent.resolve()
-  openelm_goldens = torch.load(current_dir / "openelm_lm_logits.pt")
-  kv_cache_max_len = 1024
-  model = build_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
-  idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
-  tokens[0, :4] = idx
-  input_pos = torch.arange(0, kv_cache_max_len, dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(model.config)
-  output = model.forward(tokens, input_pos, kv)
-  assert torch.allclose(
-      openelm_goldens, output["logits"][0, idx.shape[1] - 1, :], atol=1e-05
-  )
-
-
-if __name__ == "__main__":
-  input_checkpoint_path = os.path.join(
-      pathlib.Path.home(), "Downloads/llm_data/openelm"
-  )
-  define_and_run(input_checkpoint_path)

ai_edge_torch/generative/examples/openelm/verify.py

@@ -0,0 +1,61 @@
+# 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 OpenELM-3B model."""
+
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.openelm import openelm
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "What is the meaning of life?",
+    "The input prompts to generate answers.",
+)
+
+
+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
+  )
+
+  # 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 = openelm.build_model(reauthored_checkpoint)
+
+  tokenizer_checkpoint = "meta-llama/Llama-2-7b-hf"
+  verifier.log_msg("Loading the tokenizer from", tokenizer_checkpoint)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=original_model,
+      reauthored_model=reauthored_model,
+      tokenizer=tokenizer,
+      prompts=_PROMPTS.value,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/phi/phi2.py

@@ -15,16 +15,12 @@
 
 """Example of building a Phi-2 model."""
 
-import os
-import pathlib
-
 from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
-import numpy as np
 import torch
 from torch import nn
 
@@ -143,7 +139,10 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=10240,
       use_bias=True,
   )
-  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.LAYER_NORM)
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.LAYER_NORM,
+      use_input_shape=False,  # Phi-2 does layer-norm with the weight shape.
+  )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -182,29 +181,3 @@ def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
   loader.load(model)
   model.eval()
   return model
-
-
-def define_and_run(checkpoint_path: str) -> None:
-  """Instantiates and runs a Phi-2 model."""
-
-  current_dir = pathlib.Path(__file__).parent.resolve()
-  phi2_goldens = torch.load(current_dir / "phi2_lm_logits.pt")
-  kv_cache_max_len = 1024
-  model = build_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
-  idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
-  tokens[0, :4] = idx
-  input_pos = torch.arange(0, kv_cache_max_len, dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(model.config)
-  output = model.forward(tokens, input_pos, kv)
-  print("comparing with goldens..")
-  assert torch.allclose(
-      phi2_goldens, output["logits"][0, idx.shape[1] - 1, :], atol=1e-02
-  )
-
-
-if __name__ == "__main__":
-  input_checkpoint_path = os.path.join(
-      pathlib.Path.home(), "Downloads/llm_data/phi2"
-  )
-  define_and_run(input_checkpoint_path)

ai_edge_torch/generative/examples/phi/verify.py

@@ -0,0 +1,53 @@
+# 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 Phi-2 model."""
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.phi import phi2
+from ai_edge_torch.generative.utilities import verifier
+import kagglehub
+import transformers
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "What is the meaning of life?",
+    "The input prompts to generate answers.",
+)
+
+
+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)
+
+  verifier.log_msg("Building the reauthored model from", checkpoint)
+  reauthored_model = phi2.build_model(checkpoint)
+
+  verifier.log_msg("Loading the tokenizer from", checkpoint)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=original_model,
+      reauthored_model=reauthored_model,
+      tokenizer=tokenizer,
+      prompts=_PROMPTS.value,
+      atol=1e-03,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/smollm/smollm.py

@@ -16,15 +16,10 @@
 """Example of building a SmolLM model."""
 
 import copy
-import os
-import pathlib
 
 from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
-from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
-import numpy as np
-import torch
 from torch import nn
 
 TENSOR_NAMES = copy.copy(tiny_llama.TENSOR_NAMES)
@@ -104,28 +99,3 @@ def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
   loader.load(model, strict=False)
   model.eval()
   return model
-
-
-def define_and_run(checkpoint_path: str) -> None:
-  """Instantiates and runs a SmolLM model."""
-
-  current_dir = pathlib.Path(__file__).parent.resolve()
-  smollm_goldens = torch.load(current_dir / "smollm_lm_logits.pt")
-  kv_cache_max_len = 1024
-  model = build_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
-  idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
-  tokens[0, :4] = idx
-  input_pos = torch.arange(0, kv_cache_max_len, dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(model.config)
-  output = model.forward(tokens, input_pos, kv)
-  assert torch.allclose(
-      smollm_goldens, output["logits"][0, idx.shape[1] - 1, :], atol=1e-05
-  )
-
-
-if __name__ == "__main__":
-  input_checkpoint_path = os.path.join(
-      pathlib.Path.home(), "Downloads/llm_data/smollm"
-  )
-  define_and_run(input_checkpoint_path)

ai_edge_torch/generative/examples/smollm/verify.py

@@ -0,0 +1,59 @@
+# 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 SmolLM-135M model."""
+
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.smollm import smollm
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "What is the meaning of life?",
+    "The input prompts to generate answers.",
+)
+
+
+def main(_):
+  checkpoint = "HuggingFaceTB/SmolLM-135M"
+  verifier.log_msg("Loading the original model from", checkpoint)
+  original_model = transformers.AutoModelForCausalLM.from_pretrained(checkpoint)
+
+  # 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 = smollm.build_model(reauthored_checkpoint)
+
+  verifier.log_msg("Loading the tokenizer from", checkpoint)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=original_model,
+      reauthored_model=reauthored_model,
+      tokenizer=tokenizer,
+      prompts=_PROMPTS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py

@@ -15,16 +15,12 @@
 
 """Example of building a TinyLlama model."""
 
-import os
-import pathlib
-
 from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
-import numpy as np
 import torch
 from torch import nn
 
@@ -179,28 +175,3 @@ def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
   loader.load(model)
   model.eval()
   return model
-
-
-def define_and_run(checkpoint_path: str) -> None:
-  """Instantiates and runs a TinyLlama model."""
-
-  current_dir = pathlib.Path(__file__).parent.resolve()
-  tiny_llama_goldens = torch.load(current_dir / "tiny_llama_lm_logits.pt")
-  kv_cache_max_len = 1024
-  model = build_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
-  idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
-  tokens[0, :4] = idx
-  input_pos = torch.arange(0, kv_cache_max_len, dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(model.config)
-  output = model.forward(tokens, input_pos, kv)
-  assert torch.allclose(
-      tiny_llama_goldens, output["logits"][0, idx.shape[1] - 1, :], atol=1e-02
-  )
-
-
-if __name__ == "__main__":
-  input_checkpoint_path = os.path.join(
-      pathlib.Path.home(), "Downloads/llm_data/tiny_llama"
-  )
-  define_and_run(input_checkpoint_path)

ai_edge_torch/generative/examples/tiny_llama/verify.py

@@ -0,0 +1,61 @@
+# 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)
+  original_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=original_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_large.py

@@ -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,200 @@
+# 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
+
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+import numpy as np
+import torch
+
+
+def log_msg(*args):
+  print("[%s]" % datetime.datetime.now(), *args)
+
+
+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: torch.nn.Module,
+    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 (torch.nn.Module): 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: torch.nn.Module,
+    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 (torch.nn.Module): 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: torch.nn.Module,
+    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 (torch.nn.Module): 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/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20240918.dist-info → ai_edge_torch_nightly-0.3.0.dev20240919.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.dev20240919
 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.dev20240918.dist-info → ai_edge_torch_nightly-0.3.0.dev20240919.dist-info}/RECORD

@@ -3,7 +3,7 @@ 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=jWg5qA8V0XqgFoqjk0SCsNWPRBeTmfrir9u0bucHYOU,706
+ai_edge_torch/version.py,sha256=N5hYc9s2RU44J1_oe0UfJhTFo0d4JvMlKvxNlYtK0GI,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
@@ -45,13 +45,16 @@ ai_edge_torch/generative/examples/gemma/gemma.py,sha256=hjpSPzEjPHuxwRJ-vHHtCCf2
 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/openelm.py,sha256=kQTJlCDz_DHLRLlVWE0JEpbOjIGAKtxH1fTSc-jn1nU,8498
+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/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/phi2.py,sha256=mGyBI-nORoI-LhZkI4MFAonkUflIX9iimAer_K8jpck,7088
+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/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/smollm.py,sha256=_nK2DAOiSuxv5o8ip0i-gmhvvjwF5e7Dm3m5VTcsR2M,4276
+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/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
@@ -76,23 +79,24 @@ ai_edge_torch/generative/examples/test_models/toy_model.py,sha256=QyLeCqDnk71Wvv
 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/tiny_llama.py,sha256=Upo8jjqR0VKvkdczTI-Lr-1GDg0R2g4SUUGEMTUZ5uY,7023
+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/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
 ai_edge_torch/generative/layers/attention.py,sha256=Z0Y_G8IG0LmvLX2u9D8__Fkr22szB-az6wMNnZpzhkA,13233
 ai_edge_torch/generative/layers/attention_utils.py,sha256=68GXGR2HSWBFViTxX7cHifzVG-kcLS2IL2tQJPIpupg,7344
-ai_edge_torch/generative/layers/builder.py,sha256=iuAv8D7HY-azBDy7-UBILMdjuKjpe38rE2gK4H3erwE,5092
+ai_edge_torch/generative/layers/builder.py,sha256=toT9Tl1x9o5KbG-eGOEViUr4fd_4f-XLZdMQT0Ae5_8,5130
 ai_edge_torch/generative/layers/feed_forward.py,sha256=dfS1psdmomgs4EbwzkYyV_xx1xl3P1lU-3GoS8m0Avw,4221
 ai_edge_torch/generative/layers/kv_cache.py,sha256=2El7kZYnQRCRcVc63xgiAdBh9oVOksDu35p9XggvaGE,6148
-ai_edge_torch/generative/layers/model_config.py,sha256=zV3pA7giuKPrQdH81dpZz8D6LfGD-1YHuXuhIlypKc0,6784
-ai_edge_torch/generative/layers/normalization.py,sha256=iod9oNkoDS5m-yFY_Y_XMyvCU5a88ESd_s5WY34ErKA,6129
+ai_edge_torch/generative/layers/model_config.py,sha256=d0Y-EFb4Rr7iLZ4Bsdf1i92KuhY1BXRqyeUN2kuu510,6923
+ai_edge_torch/generative/layers/normalization.py,sha256=l_36uFdruJwqqyubnBTM0M-iGiJfeFafyXKPPK8KHVo,6713
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=CZqOoibLcHvUgrgaIIWAlmk3XgE2inzx340MN-npLoU,1347
-ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=VW-VP8e7FTSPCdu-6DVxpwNrIdgX0R_kq6F6MSEiyXE,3848
+ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=gXxh3papKy4FBpGEX7VyZ7rZ1Js6aHK70Q6DKrVSckY,4154
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=ZA--ohqmTfXeYQaBP1WpwFOf-TGHZmUMONocPL_hlFc,27244
+ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=c8rtlfDaeKmUfiiTKPmQhNW-U5vW9jFB2pPPcvT6qsc,27527
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
-ai_edge_torch/generative/layers/unet/model_config.py,sha256=EzF2qpuoW_qBTYO2uuThh4PN0BqF2vXQHgmfJJKVOSg,9244
+ai_edge_torch/generative/layers/unet/model_config.py,sha256=8ze9kVWMuyZVQcgK7hWYw9TM1W9lXD-2j0iMHlxoGX4,9267
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/quantize/example.py,sha256=n_YFFP3dpKjeNKYZicDGL5LqtjqwhYEIaDrC6-Ci2vE,1539
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=n1Fm8BFC8gJa_oiwwAOOghJyHtOXYZ4q-5ZRy4pHrIw,1957
@@ -104,14 +108,15 @@ ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudj
 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=TD7dELN5cVw5z9dvspFKO74Y_qIJ_VK0MYUoPdRf82Y,4498
+ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=dUYFarOldejqbMpa0j0vIDvXlWPAancuI8di3XkGxm8,4498
 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
 ai_edge_torch/generative/utilities/converter.py,sha256=MQUg2ZLmfk_2csWmQWKD_II0bXq4X3McI5i-qWraieE,2987
 ai_edge_torch/generative/utilities/loader.py,sha256=b9iotIhVDX-Zc9XjIDUaLxnV395AyBnkQe3dV5YA7Co,13297
-ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=sMMidBhGxD-0bJw5FYNVMLb7uIre3zszJ1xBAsyeDGQ,35961
+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/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
@@ -158,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.dev20240918.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20240918.dist-info/METADATA,sha256=dMaIr8Iny84IfNGQGSrtlTGkYlH_mAMmgvGWm5-pkxM,1859
-ai_edge_torch_nightly-0.3.0.dev20240918.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20240918.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20240918.dist-info/RECORD,,
+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,,