ai-edge-torch-nightly 0.4.0.dev20250225__py3-none-any.whl → 0.4.0.dev20250227__py3-none-any.whl

ai_edge_torch/_convert/conversion.py

@@ -125,6 +125,10 @@ def convert_signatures(
     else:
       exported_program = torch.export.export(**kwargs, strict=True)
 
+    exported_program = fx_infra.graph_utils.reset_from_node_meta(
+        exported_program
+    )
+
     exported_program = fx_infra.safe_run_decompositions(
         exported_program,
         fx_infra.decomp.pre_convert_decomp(),

ai_edge_torch/fx_infra/graph_utils.py

@@ -13,7 +13,10 @@
 # limitations under the License.
 # ==============================================================================
 """FX graph utilities."""
+
+from packaging import version
 import torch
+from torch.fx import traceback
 
 
 def remove_dangling_args(graph_module: torch.fx.GraphModule):
@@ -40,3 +43,32 @@ def remove_assert_tensor_metadata_nodes(graph_module: torch.fx.GraphModule):
   graph_module.graph.lint()
   graph_module.recompile()
   return graph_module
+
+
+def is_torch_version_under(torch_version: str) -> bool:
+  """Checks if the current torch version is under the given version."""
+  if not torch_version:
+    raise ValueError("torch_version cannot be empty.")
+  current_version = version.parse(torch.__version__)
+  compared_version = version.parse(torch_version)
+  return current_version < compared_version
+
+
+def reset_from_node_meta(ep: torch.export.ExportedProgram):
+  """Resets the "from_node" meta field to fx node name only for the exported program."""
+
+  for node in ep.graph.nodes:
+    if not hasattr(node, "meta") or "from_node" not in node.meta:
+      continue
+    if is_torch_version_under("2.6.0.dev0"):
+      # For torch version under 2.6.0, the history stack is a list of tuple. We
+      # will only keep the current node's name in the history stack.
+      history = [(node.name,)]
+    else:
+      # Clean up the history stack by keeping only the current node info (fx
+      # node name and graph id) in a list of size 1. Clear the "from_node" field
+      # to prevent redundant additions to the history stack.
+      history = [traceback.NodeSource(node)]
+      history[0].from_node = []
+    node.meta["from_node"] = history
+  return ep

ai_edge_torch/generative/examples/smollm/convert_to_tflite.py

@@ -22,7 +22,7 @@ from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.smollm import smollm
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import model_builder
 
 _CHECKPOINT_PATH = flags.DEFINE_string(
     'checkpoint_path',
@@ -59,6 +59,11 @@ _LORA_RANKS = flags.DEFINE_multi_integer(
     None,
     'If set, the model will be converted with the provided list of LoRA ranks.',
 )
+_DECODE_BATCH_SIZE = flags.DEFINE_integer(
+    'decode_batch_size',
+    1,
+    'The batch size for the decode signature.',
+)
 
 
 def main(_):
@@ -72,7 +77,9 @@ def main(_):
       prefill_seq_len=_PREFILL_SEQ_LENS.value,
       quantize=_QUANTIZE.value,
       lora_ranks=_LORA_RANKS.value,
-      export_config=ExportConfig(),
+      export_config=model_builder.ExportConfig(
+          decode_batch_size=_DECODE_BATCH_SIZE.value
+      ),
   )
 
 

ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py

@@ -22,17 +22,22 @@ from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.smollm import smollm
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import model_builder
 
 _CHECKPOINT_PATH = flags.DEFINE_string(
     'checkpoint_path',
     os.path.join(pathlib.Path.home(), 'Downloads/llm_data/smollm2'),
     'The path to the model checkpoint, or directory holding the checkpoint.',
 )
-_TFLITE_PATH = flags.DEFINE_string(
-    'tflite_path',
+_OUTPUT_PATH = flags.DEFINE_string(
+    'output_path',
     '/tmp/',
-    'The tflite file path to export.',
+    'The path to export the tflite model.',
+)
+_OUTPUT_NAME_PREFIX = flags.DEFINE_string(
+    'output_name_prefix',
+    'smollm2',
+    'The prefix of the output tflite model name.',
 )
 _PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(
     'prefill_seq_lens',
@@ -49,6 +54,16 @@ _QUANTIZE = flags.DEFINE_bool(
     True,
     'Whether the model should be quantized.',
 )
+_LORA_RANKS = flags.DEFINE_multi_integer(
+    'lora_ranks',
+    None,
+    'If set, the model will be converted with the provided list of LoRA ranks.',
+)
+_DECODE_BATCH_SIZE = flags.DEFINE_integer(
+    'decode_batch_size',
+    1,
+    'The batch size for the decode signature.',
+)
 
 
 def main(_):
@@ -56,14 +71,16 @@ def main(_):
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
 
-  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
-  output_filename = f'smollm2_{quant_suffix}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
-      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
+      output_path=_OUTPUT_PATH.value,
+      output_name_prefix=_OUTPUT_NAME_PREFIX.value,
       prefill_seq_len=_PREFILL_SEQ_LENS.value,
       quantize=_QUANTIZE.value,
-      export_config=ExportConfig(),
+      lora_ranks=_LORA_RANKS.value,
+      export_config=model_builder.ExportConfig(
+          decode_batch_size=_DECODE_BATCH_SIZE.value
+      ),
   )
 
 

ai_edge_torch/generative/layers/attention.py

@@ -48,7 +48,6 @@ class TransformerBlock(nn.Module):
         config.pre_attention_norm_config,
     )
     self.atten_func = CausalSelfAttention(
-        model_config.batch_size,
         model_config.embedding_dim,
         config.attn_config,
         model_config.enable_hlfb,
@@ -115,7 +114,6 @@ class CausalSelfAttention(nn.Module):
 
   def __init__(
       self,
-      batch_size: int,
       dim: int,
       config: cfg.AttentionConfig,
       enable_hlfb: bool,
@@ -123,14 +121,12 @@ class CausalSelfAttention(nn.Module):
     """Initialize an instance of CausalSelfAttention.
 
     Args:
-      batch_size (int): batch size of the input tensor.
       dim (int): causal attention's input/output dimmension.
       config (cfg.AttentionConfig): attention specific configurations.
       enable_hlfb (bool): whether hlfb is enabled or not.
     """
     super().__init__()
     self.kv_cache = None
-    self.batch_size = batch_size
     qkv_shape = (
         config.num_heads + 2 * config.num_query_groups
     ) * config.head_dim
@@ -179,11 +175,6 @@ class CausalSelfAttention(nn.Module):
     """
     # Batch size, sequence length, embedding dimensionality.
     B, T, E = x.size()
-    assert B == self.batch_size, (
-        "batch size of input tensor must match with the batch size specified in"
-        " the model configuration."
-    )
-
     qkv = self.qkv_projection(x)
 
     # Assemble into a number of query groups to support MHA, MQA and GQA.
@@ -290,7 +281,6 @@ class CrossAttention(nn.Module):
 
   def __init__(
       self,
-      batch_size: int,
       query_dim: int,
       cross_dim: int,
       hidden_dim: int,
@@ -301,7 +291,6 @@ class CrossAttention(nn.Module):
     """Initialize an instance of CrossAttention.
 
     Args:
-      batch_size (int): batch size of the input tensor.
       query_dim (int): query tensor's dimension.
       cross_dim (int): cross attention's dimensions, for key and value tensors.
       hidden_dim (int): hidden dimension that q, k, v tensors project to.
@@ -376,7 +365,6 @@ class CrossAttention(nn.Module):
 
     if rope is not None:
       # Compute rotary positional embedding for query and key.
-      n_elem = int(self.config.rotary_percentage * self.config.head_dim)
       cos, sin = rope
       q, k = rotary_pos_emb.apply_rope_inline(q, k, cos, sin)
 

ai_edge_torch/generative/layers/kv_cache.py

@@ -18,14 +18,11 @@
 import dataclasses
 from typing import List, Tuple
 
-from ai_edge_torch import hlfb
 from ai_edge_torch.generative.layers import model_config
 from ai_edge_torch.generative.utilities.dynamic_update_slice import dynamic_update_slice
 import torch
 import torch.utils._pytree as pytree
 
-BATCH_SIZE = 1
-
 
 @dataclasses.dataclass
 class KVCacheEntry:
@@ -45,9 +42,10 @@ class KVCacheEntry:
       config: model_config.AttentionConfig,
       dtype: torch.dtype = torch.float32,
       device: torch.device = None,
+      batch_size: int = 1,
   ) -> "KVCacheEntry":
     """Build an instance of the class based on model config."""
-    shape = (BATCH_SIZE, kv_cache_max, config.num_query_groups, config.head_dim)
+    shape = (batch_size, kv_cache_max, config.num_query_groups, config.head_dim)
     k = torch.zeros(shape, dtype=dtype, device=device)
     v = torch.zeros(shape, dtype=dtype, device=device)
     obj = cls(k_cache=k, v_cache=v)
@@ -66,6 +64,7 @@ class KVCache:
       config: model_config.ModelConfig,
       dtype: torch.dtype = torch.float32,
       device: torch.device = None,
+      batch_size: int = 1,
   ) -> "KVCache":
     """Build an instance of the class based on model config.
 
@@ -75,17 +74,21 @@ class KVCache:
           Defaults to torch.float32.
         device (torch.device, optional): The device placement of the cache
           tensors. Defaults to None.
+        batch_size (int, optional): The batch size of the cache tensors.
+          Defaults to 1.
 
     Returns:
         KVCache: The created cache object.
     """
     caches = [
         KVCacheEntry.from_model_config(
-            config.kv_cache_max if not config.block_config(idx).kv_cache_max_len 
+            config.kv_cache_max
+            if not config.block_config(idx).kv_cache_max_len
             else config.block_config(idx).kv_cache_max_len,
             config.block_config(idx).attn_config,
             dtype,
             device,
+            batch_size,
         )
         for idx in range(config.num_layers)
     ]

ai_edge_torch/generative/layers/model_config.py

@@ -220,9 +220,6 @@ class ModelConfig:
   # The maximum sequence length of the KV cache. Should not exceed max_seq_len.
   kv_cache_max_len: int = 0
 
-  # Default batch size of the exported model. Default value is 1.
-  batch_size: int = 1
-
   # Softcap on the model output logits.
   final_logit_softcap: Optional[float] = None
 

ai_edge_torch/generative/utilities/converter.py

@@ -110,6 +110,11 @@ def convert_to_tflite(
   lora_suffix = (
       '' if not lora_ranks else f'_lora{",".join(map(str, lora_ranks))}'
   )
+
+  if export_config is not None:
+    if export_config.decode_batch_size > 1:
+      output_name_prefix += f'_dbs{export_config.decode_batch_size}'
+
   output_filename = (
       f'{output_name_prefix}_{quant_suffix}_ekv{kv_size}{lora_suffix}.tflite'
   )
@@ -162,9 +167,14 @@ def _export_helper(
   if prefill_masks:
     assert len(prefill_masks) == len(prefill_seq_lens)
 
-  decode_token = torch.tensor([[0]], dtype=torch.int)
+  decode_token = torch.tensor(
+      [[0] for _ in range(export_config.decode_batch_size)], dtype=torch.int
+  )
   decode_input_pos = torch.tensor([0], dtype=torch.int)
-  kv = export_config.kvcache_cls.from_model_config(config)
+  prefill_kv = export_config.kvcache_cls.from_model_config(config)
+  decode_kv = export_config.kvcache_cls.from_model_config(
+      config, batch_size=export_config.decode_batch_size
+  )
 
   quant_config = quant_recipes.full_int8_dynamic_recipe() if quantize else None
 
@@ -183,7 +193,7 @@ def _export_helper(
       sample_kwargs = {
           'tokens': prefill_tokens,
           'input_pos': prefill_input_pos,
-          'kv_cache': kv,
+          'kv_cache': prefill_kv,
       }
       if prefill_masks is not None:
         sample_kwargs['mask'] = prefill_masks[i]
@@ -211,7 +221,7 @@ def _export_helper(
     sample_kwargs = {
         'tokens': decode_token,
         'input_pos': decode_input_pos,
-        'kv_cache': kv,
+        'kv_cache': decode_kv,
     }
     if export_config.decode_mask is not None:
       sample_kwargs['mask'] = export_config.decode_mask

ai_edge_torch/generative/utilities/model_builder.py

@@ -60,6 +60,8 @@ class ExportConfig:
   decode_mask: Optional[torch.Tensor | List[torch.Tensor]] = None
   # The KV Cache class for K and V buffers in attention.
   kvcache_cls: type = kv_utils.KVCache
+  # The batch size of the decode signature.
+  decode_batch_size: int = 1
 
 
 class DecoderOnlyModel(nn.Module):

ai_edge_torch/odml_torch/debuginfo/_build.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 import re
+from ai_edge_torch.fx_infra import graph_utils
 import torch
 
 
@@ -48,8 +49,6 @@ def _get_canonical_filename(filename: str):
     The canonicalized filename.
   """
 
-  # TODO(yijieyang): We should add a config option to provide a regex to strip
-  # from the debug info. Currently absolute path is used.
   return filename
 
 
@@ -57,9 +56,13 @@ def _get_canoical_nodename(node: torch.fx.Node) -> str:
   """Get the canonical node name from the node's history."""
 
   history = node.meta.get("from_node", [])
+  if not history:
+    return None
 
-  if len(history) > 1:  # Compatible with torch version under 2.6.0
-    return history[1][0]
+  # Compatible with torch version under 2.6.0. The history stack is a list of
+  # tuple. The first element of the first tuple is the node name.
+  if graph_utils.is_torch_version_under("2.6.0.dev0"):
+    return history[0][0]
 
   if not hasattr(history[0], "name"):
     return None
@@ -68,12 +71,10 @@ def _get_canoical_nodename(node: torch.fx.Node) -> str:
     names.append(history[0].name)
     history = history[0].from_node
 
-  # Based on the experiment, the third to last name in the history stack
-  # can be mapped to the original torch node name. The history stack is
-  # generated by tracing the node's transformation history during lowering.
-  if len(names) > 2:
-    return names[-3]
-  return None
+  # The history stack is generated by tracing the node's transformation history
+  # during lowering. The last name in the history stack is used to map to the
+  # original torch fx node name.
+  return names[-1]
 
 
 def build_mlir_file_debuginfo(node: torch.fx.Node):

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.4.0.dev20250225"
+__version__ = "0.4.0.dev20250227"

{ai_edge_torch_nightly-0.4.0.dev20250225.dist-info → ai_edge_torch_nightly-0.4.0.dev20250227.dist-info}/METADATA

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

@@ -2,9 +2,9 @@ ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,120
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=IlCclQHtv864xDiY226T7A6oRSPGetrRq1_B_aIyGAg,706
+ai_edge_torch/version.py,sha256=K2jtDrBNGi74j_uQYVUT6MJ2-aQFKkKy5ZYur9iWdVU,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=mckvxznKLXdF2HuJg_IxQaT5Ty-iWl_iXElHEugH3VI,5452
+ai_edge_torch/_convert/conversion.py,sha256=gpXQnifODU-mWxkUZw_3ov1lEYBw1SPVIcqj5k7pTGo,5550
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
 ai_edge_torch/_convert/converter.py,sha256=075F8LRewk_033Ebsnft7FJr3KgtIbtZ_-8udIPy6ho,9980
 ai_edge_torch/_convert/signature.py,sha256=-YKJdLk-eNEHfhdPCtcQVtZf915SoVePEFxKXPPf16c,2572
@@ -41,7 +41,7 @@ ai_edge_torch/fx_infra/__init__.py,sha256=APjkSqEfwDxcnI8k53rGi3Ef-G2L-M8fdaPGpx
 ai_edge_torch/fx_infra/_canonicalize_pass.py,sha256=GDRoDdPVQw--QQFTT5J_C3TVuphL31m6K6F1-67SE4s,1097
 ai_edge_torch/fx_infra/_safe_run_decompositions.py,sha256=ZbWheeZ8ydsxCk2aVGUgUynrkEkBOMjBCzPhS5uq4sU,2595
 ai_edge_torch/fx_infra/decomp.py,sha256=S58SCgwMHYVFl_hJwlJxvu2wcI-AGNn82gel3qmTPrU,2500
-ai_edge_torch/fx_infra/graph_utils.py,sha256=3UZAOHWOUh2LCj1E2_AKQn3gRDILi9JCdqSScjyOd4M,1535
+ai_edge_torch/fx_infra/graph_utils.py,sha256=nqGe-xIJ77RamSUh0UYyI2XHOsZqFDWax-vpRAtVR_E,2796
 ai_edge_torch/fx_infra/pass_base.py,sha256=Ic2AlhSoRFscz6l7gJKvWVNMDLQFfAw5kRf84-ZR9qM,2904
 ai_edge_torch/generative/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -102,8 +102,8 @@ ai_edge_torch/generative/examples/qwen_vl/verify.py,sha256=JUwHoC_zvcC3RC3wZ3e3e
 ai_edge_torch/generative/examples/qwen_vl/verify_decoder.py,sha256=xPWoOBLh2eK12KEhELLYymfL7xvc0chmYC98c6x37oo,2602
 ai_edge_torch/generative/examples/qwen_vl/verify_image_encoder.py,sha256=PZ392nDoJG2OmHZ_7Jet3Zu1JkN6QErxKcDc7a-PPds,3126
 ai_edge_torch/generative/examples/smollm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/smollm/convert_to_tflite.py,sha256=megskv1oiPhwHSnguoG7zV-esXp1Ns_FPeMLAYKhDb0,2522
-ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py,sha256=CjY1i0iCYxFSjhCpQZwxkmVxILgeo0zu1m0oBrHqyDU,2311
+ai_edge_torch/generative/examples/smollm/convert_to_tflite.py,sha256=hWko-RJB8eXNUfi4EzQ2yjW30YE4UB4zAz7rd2Q5qpg,2708
+ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py,sha256=sJ-o385eqQsciv0TEQRkixvS0DD6dKruAuK0zlEsDoY,2715
 ai_edge_torch/generative/examples/smollm/smollm.py,sha256=3uUltb6D3Q1aHpndcYTJrsWM_RBwLAraKDniH8ZZous,3779
 ai_edge_torch/generative/examples/smollm/verify.py,sha256=KpYxVz_lv61YWy6HLfwT68n0owZMvty5Rr3W7ZNWWSw,2702
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -136,13 +136,13 @@ ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=LRu6PSw7Lqu6HGbv1t
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=4rFrppMRKlTwwZeX1ON_cdp4yUqoTOES161IZQkJF6c,1143
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=myGjal5A8yIBoqgArd2k40rZmCgD1Ya369KR7182bhI,2129
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/attention.py,sha256=Pm8FLKh-NnOvUjqQC9oX5oghPbdivZvlPVkgOVTShoU,13703
+ai_edge_torch/generative/layers/attention.py,sha256=wLZ1jgUlcODBWgK3hnnhclHuuQDqYuGOZdYAI9EooOM,13247
 ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
 ai_edge_torch/generative/layers/builder.py,sha256=LXGuSHIx6QZAzLFm7aJvlzoMPgQwbXLFchGEKYwOOUA,5090
 ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
-ai_edge_torch/generative/layers/kv_cache.py,sha256=sGGAZD0mWYuO4FukZfDbHXoxpBOBE9lTYICvZzDj5F8,6400
+ai_edge_torch/generative/layers/kv_cache.py,sha256=jwbt0-2fd_CNWS2fp4nf0zvh6kk5citINGlFC_RtEUU,6540
 ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
-ai_edge_torch/generative/layers/model_config.py,sha256=EA1Ey5-c1IOLRNANSUnZ7gtNTA0o6OJxrz_I_mp8cjw,8244
+ai_edge_torch/generative/layers/model_config.py,sha256=wNsZDzZQoimOKdZ9FWMCktPj-pQ_0D7084hgzMT5XYo,8155
 ai_edge_torch/generative/layers/normalization.py,sha256=MbwH-n80Fob5YvjBzdqDjBizMHLzSJGYRDdbD-rL5C0,6174
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIrILJ7blceAcxrNqX1ZCN0ECKG1gz-bV8,2655
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=vp8dVx6tOe99neJhpbrtIt5fvN5NFw19JVH1v0yi5Mg,4154
@@ -173,10 +173,10 @@ ai_edge_torch/generative/test/test_quantize.py,sha256=bEJMhpQ9bIDUZVBXTW888728Fc
 ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv9og0DDLHU,2656
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
 ai_edge_torch/generative/utilities/bmm_4d.py,sha256=2BMOYiFVUsl-bjxmLkrX4N7kpO0CnhB7eDYxm_iBCr8,2533
-ai_edge_torch/generative/utilities/converter.py,sha256=_PO9lYCdNNYPVsAqh8QQVMG_8TUBshKwmaR1cdT6Ang,8065
+ai_edge_torch/generative/utilities/converter.py,sha256=VtG42CVz657XbvTj-FZJiCFW0Hm11OVKKC_mr2tjxhc,8413
 ai_edge_torch/generative/utilities/dynamic_update_slice.py,sha256=e2mhx-Vp8sUK4EXoPtpZLSx3TViqLAKs67EhKcXBjAQ,2121
 ai_edge_torch/generative/utilities/loader.py,sha256=A3SOjPXp--AsvoP1hqj5QKWE4sgxoFc3H5EBUz_Eogc,13531
-ai_edge_torch/generative/utilities/model_builder.py,sha256=5WqcxpeTdt51nVoUwt9g5kKB5wQKj2eYbiaz7k6Ofxg,6815
+ai_edge_torch/generative/utilities/model_builder.py,sha256=eY3qAcBhupIn955YnWuzUi9hoWYvl4ntRWA6PBudzMo,6888
 ai_edge_torch/generative/utilities/moonshine_loader.py,sha256=_RpFabSqtGH5PHiP3_1f6QfO14qMADUxr_HGRlVDFB0,4891
 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
@@ -205,7 +205,7 @@ ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xK
 ai_edge_torch/odml_torch/composite/mark_tensor.py,sha256=U--rwl-XkWKgkdXCXDn6yySug8FR66o1YFUAIoSaWW4,3523
 ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py,sha256=2Y52E_gLeoXpMcPpV-svXsgN3JbEIjnPVjm0xkpTUdQ,3319
 ai_edge_torch/odml_torch/debuginfo/__init__.py,sha256=3A_lMyj-B-DOhLJG6WmjKvZK5te2rXje8FrfqOhZsN0,959
-ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=k9Kas790kpMS5OrVcLzIr48ejAzcc2smrroKAHHM7TQ,3311
+ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=6Ns2rlfOilLJEk5cUxlkRwm2uxOgEF2-0S2DMcOqr6A,3319
 ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py,sha256=IvOBQyROI9WHS3umHRxsDW-1YElU9BPWzKtJA2eKWOI,1739
 ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNitEeg-IoBUGNfUxsDSA,798
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=LqwZ1vCJTSOzgzvH8LUAN-sAkF-l_pGj1AMEIzAqHCA,6638
@@ -230,8 +230,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.4.0.dev20250225.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.4.0.dev20250225.dist-info/METADATA,sha256=R_qRD9bYEGBzEuxpKkHA0pjZyNWiQ5CG8hT43wNYDeM,1966
-ai_edge_torch_nightly-0.4.0.dev20250225.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.4.0.dev20250225.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.4.0.dev20250225.dist-info/RECORD,,
+ai_edge_torch_nightly-0.4.0.dev20250227.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.4.0.dev20250227.dist-info/METADATA,sha256=cHcz3adq1WwVddazAJ06h7SKITJm70eMpFVjoNa2Jw4,1966
+ai_edge_torch_nightly-0.4.0.dev20250227.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.4.0.dev20250227.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.4.0.dev20250227.dist-info/RECORD,,