ai-edge-torch-nightly 0.3.0.dev20250105__py3-none-any.whl → 0.3.0.dev20250108__py3-none-any.whl

ai_edge_torch/generative/layers/rotary_position_embedding.py

@@ -32,64 +32,57 @@ def apply_rope(
   """
   x = x.transpose(1, 2)
   head_size = x.size(-1)
-  x1, x2 = torch.split(x, head_size // 2, dim=-1)
-  left = x1 * cos - x2 * sin
-  right = x2 * cos + x1 * sin
-  roped = torch.cat([left, right], dim=-1)
+  x1 = x[..., : head_size // 2]  # (B, nh, T, hs/2)
+  x2 = x[..., head_size // 2 :]  # (B, nh, T, hs/2)
+  rotated = torch.cat((-x2, x1), dim=-1)  # (B, nh, T, hs)
+  roped = (x * cos) + (rotated * sin)
   return roped.transpose(1, 2).type_as(x)
 
 
-def build_rope(
+def apply_rope_inline(
+    q: torch.Tensor,
+    k: torch.Tensor,
     input_pos: torch.Tensor,
     n_elem: int,
-    head_dim: int,
     base: int = 10_000,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-  """Computes rotary positional embedding cosine and sine tensors.
+  """Computes rotary positional embedding inline for a query and key.
 
   Args:
+    q: the query tensor.
+    k: the key tensor.
     input_pos: the sequence indices for the query and key
     n_elem: number of elements of the head dimension for RoPE computation
-    base: the base of the exponentiated value for RoPE.
 
   Returns:
-    cos, sin tensors
+    output the RoPE'd query and key.
   """
 
   if n_elem <= 0:
-    return None, None
+    return q, k
 
   theta = 1.0 / (base ** (torch.arange(0, n_elem, 2).float() / n_elem))
   freq_exponents = (2.0 / n_elem) * torch.arange(
-      head_dim // 2, dtype=torch.float32
+      q.shape[-1] // 2, dtype=torch.float32
   )
   timescale = float(base) ** freq_exponents
   radians = input_pos.clone().unsqueeze(0).unsqueeze(-1) / timescale.unsqueeze(
       0
   ).unsqueeze(0)
-  cos = torch.cos(radians)
-  sin = torch.sin(radians)
-  return cos, sin
-
+  cos = torch.cos(radians).type_as(q)
+  sin = torch.sin(radians).type_as(q)
 
-def apply_rope_inline(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-  """Computes rotary positional embedding inline for a query and key.
-
-  Args:
-    q: the query tensor.
-    k: the key tensor.
-    cos: the cosine tensor.
-    sin: the sine tensor.
-
-  Returns:
-    output the RoPE'd query and key.
-  """
+  def apply(x, sin, cos):
+    x = x.transpose(1, 2)
+    b, h, s, d = x.shape
+    ans = torch.split(x, d // 2, dim=-1)
+    x1, x2 = ans
+    left = x1 * cos - x2 * sin
+    right = x2 * cos + x1 * sin
+    res = torch.cat([left, right], dim=-1)
+    res = res.transpose(1, 2)
+    return res
 
-  q_roped = apply_rope(q, cos, sin)
-  k_roped = apply_rope(k, cos, sin)
+  q_roped = apply(q, sin, cos)
+  k_roped = apply(k, sin, cos)
   return q_roped, k_roped

ai_edge_torch/generative/test/test_lora.py

@@ -0,0 +1,147 @@
+# Copyright 2025 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.
+# ==============================================================================
+
+"""A suite of tests to validate LoRA utilities."""
+
+from ai_edge_torch.generative.layers import lora as lora_utils
+import ai_edge_torch.generative.layers.model_config as cfg
+import torch
+from absl.testing import absltest as googletest
+from tensorflow.python.platform import resource_loader  # pylint: disable=g-direct-tensorflow-import
+
+
+class TestLora(googletest.TestCase):
+  """Tests for LoRA utilities."""
+
+  def test_safetensors_builder(self):
+    """Converts a safetensors file to a LoRA module."""
+
+    tensor_names = lora_utils.LoRATensorNames(
+        attn_query_w_a=(
+            "base_model.model.model.layers.{}.self_attn.q_proj.lora_A.weight"
+        ),
+        attn_query_w_b=(
+            "base_model.model.model.layers.{}.self_attn.q_proj.lora_B.weight"
+        ),
+        attn_key_w_a=(
+            "base_model.model.model.layers.{}.self_attn.k_proj.lora_A.weight"
+        ),
+        attn_key_w_b=(
+            "base_model.model.model.layers.{}.self_attn.k_proj.lora_B.weight"
+        ),
+        attn_value_w_a=(
+            "base_model.model.model.layers.{}.self_attn.v_proj.lora_A.weight"
+        ),
+        attn_value_w_b=(
+            "base_model.model.model.layers.{}.self_attn.v_proj.lora_B.weight"
+        ),
+        attn_output_w_a=(
+            "base_model.model.model.layers.{}.self_attn.o_proj.lora_A.weight"
+        ),
+        attn_output_w_b=(
+            "base_model.model.model.layers.{}.self_attn.o_proj.lora_B.weight"
+        ),
+    )
+
+    safetensors_file = resource_loader.get_path_to_datafile(
+        "fixtures/test_lora_rank16.safetensors"
+    )
+    config = self._get_test_config(
+        num_layers=1,
+        head_dim=8,
+        num_query_groups=1,
+        kv_cache_max_len=16,
+    )
+    lora = lora_utils.LoRA.from_safetensors(
+        safetensors_file,
+        scale=1.0,
+        lora_tensor_names=tensor_names,
+        config=config,
+    )
+    self.assertEqual(lora.get_rank(), 16)
+
+  def test_torch_export(self):
+    """Tests the export of the LoRA module."""
+
+    class TestModel(torch.nn.Module):
+
+      def forward(self, x: torch.Tensor, lora: lora_utils.LoRA) -> torch.Tensor:
+        x += lora_utils.apply_lora(x, lora.adapters[0].attention.query)
+        return x
+
+    n = 1
+    head_dim = 2
+    num_query_groups = 1
+    key_length = 4
+    config = self._get_test_config(
+        num_layers=n,
+        head_dim=head_dim,
+        num_query_groups=num_query_groups,
+        kv_cache_max_len=key_length,
+    )
+    inputs = torch.zeros((n, 1, head_dim))
+    lora = lora_utils.LoRA.zeros(rank=16, config=config)
+    model = TestModel()
+    exported_program = torch.export.export(model, (inputs, lora))
+    input_specs = exported_program.graph_signature.input_specs
+    # 9 inputs: 1 for x, 2 for query lora, 2 for key lora, 2 for value lora,
+    # 2 for output lora.
+    self.assertLen(input_specs, 9)
+    self.assertEqual(input_specs[0].arg.name, "x")
+    self.assertEqual(input_specs[1].arg.name, "lora_atten_q_a_prime_weight_0")
+    self.assertEqual(input_specs[2].arg.name, "lora_atten_q_b_prime_weight_0")
+    self.assertEqual(input_specs[3].arg.name, "lora_atten_k_a_prime_weight_0")
+    self.assertEqual(input_specs[4].arg.name, "lora_atten_k_b_prime_weight_0")
+    self.assertEqual(input_specs[5].arg.name, "lora_atten_v_a_prime_weight_0")
+    self.assertEqual(input_specs[6].arg.name, "lora_atten_v_b_prime_weight_0")
+    self.assertEqual(input_specs[7].arg.name, "lora_atten_o_a_prime_weight_0")
+    self.assertEqual(input_specs[8].arg.name, "lora_atten_o_b_prime_weight_0")
+
+  def test_lora_tflite_serialization(self):
+    """Tests the serialization of the LoRA module."""
+    config = self._get_test_config(
+        num_layers=2,
+        head_dim=8,
+        num_query_groups=1,
+        kv_cache_max_len=16,
+    )
+    lora = lora_utils.LoRA.random(rank=16, config=config)
+    flatbuffer_model = lora.to_tflite()
+    recovered_lora = lora_utils.LoRA.from_flatbuffers(flatbuffer_model)
+    self.assertEqual(lora, recovered_lora)
+
+  def _get_test_config(
+      self, num_layers, head_dim, num_query_groups, kv_cache_max_len
+  ):
+    """Returns a test model config."""
+    attn_config = cfg.AttentionConfig(
+        num_heads=1, head_dim=head_dim, num_query_groups=num_query_groups
+    )
+    block_config = cfg.TransformerBlockConfig(
+        attn_config=attn_config, ff_config=None
+    )
+    config = cfg.ModelConfig(
+        kv_cache_max_len=kv_cache_max_len,
+        embedding_dim=head_dim,
+        block_configs=block_config,
+        num_layers=num_layers,
+        max_seq_len=None,
+        vocab_size=None,
+    )
+    return config
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_torch/generative/utilities/converter.py

@@ -15,16 +15,15 @@
 
 """Common utility functions for model conversion."""
 
-from functools import partial
-from typing import Any, Union
-
+import os
+from typing import Optional, Union
 from ai_edge_torch._convert import converter as converter_utils
+from ai_edge_torch.generative.layers import lora as lora_utils
 import ai_edge_torch.generative.layers.kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.quantize import quant_recipes
 from ai_edge_torch.generative.utilities.model_builder import ExportConfig
 import torch
-import torch.nn as nn
 
 
 class ExportableModule(torch.nn.Module):
@@ -41,11 +40,13 @@ class ExportableModule(torch.nn.Module):
 
 def convert_to_tflite(
     pytorch_model: torch.nn.Module,
-    tflite_path: str,
+    output_path: str,
+    output_name_prefix: str,
     prefill_seq_len: Union[int, list[int]],
     pixel_values_size: torch.Size = None,
     quantize: bool = True,
     config: cfg.ModelConfig = None,
+    lora_ranks: Optional[list[int]] = None,
     export_config: ExportConfig = None,
 ):
   """Converts a nn.Module model to multi-signature tflite model.
@@ -79,21 +80,65 @@ def convert_to_tflite(
 
   Args:
       pytorch_model (torch.nn.Module): PyTorch model to convert to tflite.
-      tflite_path (str): The tflite file path to export.
-      prefill_seq_len (Union[int, list[int]]): A list of prefill lengths to
-        export.
+      output_path (str): The path to export the tflite model.
+      output_name_prefix (str): The prefix of the tflite model name.
+      prefill_seq_len (Union[int, list[int]]): The prefill sequence length to
+        use. If a list, the model will have multiple prefill signatures.
       pixel_values_size (torch.Size, optional): The size of pixel values to pass
         to the model. If None, the model is not expected to take pixel values.
       quantize (bool, optional): Whether the model should be quanized. Defaults
         to True.
       config (cfg.ModelConfig, optional): The model config used to configure KV
         cache. If None, it uses the config of the pytorch_model.
+      lora_ranks (list[int], optional): The ranks of the LORA layers. If None,
+        no LoRA signatures will be added.
   """
+  # pylint: disable=protected-access
+  torch._dynamo.config.cache_size_limit = 64
+
+  config = config if config else pytorch_model.config
   prefill_seq_lens = (
       [prefill_seq_len] if isinstance(prefill_seq_len, int) else prefill_seq_len
   )
+  loras = [None]
+  if lora_ranks is not None:
+    for rank in lora_ranks:
+      lora = lora_utils.LoRA.zeros(rank, config)
+      loras.append(lora)
+
+  quant_suffix = 'q8' if quantize else 'f32'
+  kv_size = config.kv_cache_max_len
+  lora_suffix = (
+      '' if not lora_ranks else f'_lora{",".join(map(str, lora_ranks))}'
+  )
+  output_filename = (
+      f'{output_name_prefix}_{quant_suffix}_ekv{kv_size}{lora_suffix}.tflite'
+  )
+  output_file = os.path.join(output_path, output_filename)
+
+  _export_helper(
+      pytorch_model,
+      output_file,
+      prefill_seq_lens,
+      pixel_values_size,
+      quantize,
+      config,
+      loras,
+      export_config,
+  )
 
-  # Tensors used to trace the model graph during conversion.
+
+def _export_helper(
+    pytorch_model: torch.nn.Module,
+    output_file: str,
+    prefill_seq_lens: list[int],
+    pixel_values_size: torch.Size,
+    quantize: bool,
+    config: cfg.ModelConfig,
+    loras: list[None | lora_utils.LoRA],
+    export_config: ExportConfig,
+):
+  """Helper function to export a model to tflite."""
   prefill_tokens_list = []
   prefill_input_pos_list = []
   for seq_len in prefill_seq_lens:
@@ -108,9 +153,7 @@ def convert_to_tflite(
 
   decode_token = torch.tensor([[0]], dtype=torch.int)
   decode_input_pos = torch.tensor([0], dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(
-      config if config else pytorch_model.config
-  )
+  kv = kv_utils.KVCache.from_model_config(config)
 
   quant_config = quant_recipes.full_int8_dynamic_recipe() if quantize else None
 
@@ -119,44 +162,54 @@ def convert_to_tflite(
   mod = ExportableModule(pytorch_model, export_config=export_config)
 
   converter = converter_utils.Converter()
-  for i in range(len(prefill_seq_lens)):
-    prefill_seq_len = prefill_seq_lens[i]
-    prefill_tokens = prefill_tokens_list[i]
-    prefill_input_pos = prefill_input_pos_list[i]
-    if i == 0 and len(prefill_seq_lens) == 1:
-      prefill_signature_name = 'prefill'
-    else:
-      prefill_signature_name = f'prefill_{prefill_seq_len}'
-    converter.add_signature(
-        prefill_signature_name,
-        mod,
-        sample_kwargs={
-            'tokens': prefill_tokens,
-            'input_pos': prefill_input_pos,
-            'kv_cache': kv,
-        },
-    )
-    if prefill_pixel_values is not None:
+  for lora in loras:
+    for i in range(len(prefill_seq_lens)):
+      prefill_seq_len = prefill_seq_lens[i]
+      prefill_tokens = prefill_tokens_list[i]
+      prefill_input_pos = prefill_input_pos_list[i]
+      if i == 0 and len(prefill_seq_lens) == 1:
+        prefill_signature_name = 'prefill'
+      else:
+        prefill_signature_name = f'prefill_{prefill_seq_len}'
+
+      sample_kwargs = {
+          'tokens': prefill_tokens,
+          'input_pos': prefill_input_pos,
+          'kv_cache': kv,
+      }
+      if lora is not None:
+        prefill_signature_name += f'_lora_r{lora.get_rank()}'
+        sample_kwargs['lora'] = lora
+
       converter.add_signature(
-          prefill_signature_name + '_pixel',
+          prefill_signature_name,
           mod,
-          sample_kwargs={
-              'tokens': prefill_tokens,
-              'input_pos': prefill_input_pos,
-              'kv_cache': kv,
-              'pixel_values': prefill_pixel_values,
-          },
+          sample_kwargs=sample_kwargs,
       )
 
-  converter.add_signature(
-      'decode',
-      mod,
-      sample_kwargs={
-          'tokens': decode_token,
-          'input_pos': decode_input_pos,
-          'kv_cache': kv,
-      },
-  )
+      if prefill_pixel_values is not None:
+        converter.add_signature(
+            prefill_signature_name + '_pixel',
+            mod,
+            sample_kwargs={
+                **sample_kwargs,
+                'pixel_values': prefill_pixel_values,
+            },
+        )
+
+    sample_kwargs = {
+        'tokens': decode_token,
+        'input_pos': decode_input_pos,
+        'kv_cache': kv,
+    }
+    if lora is not None:
+      sample_kwargs['lora'] = lora
+
+    converter.add_signature(
+        'decode' if lora is None else f'decode_lora_r{lora.get_rank()}',
+        mod,
+        sample_kwargs=sample_kwargs,
+    )
 
   edge_model = converter.convert(quant_config=quant_config)
-  edge_model.export(tflite_path)
+  edge_model.export(output_file)

ai_edge_torch/generative/utilities/model_builder.py

@@ -22,13 +22,14 @@ from typing import Optional, Tuple
 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
+from ai_edge_torch.generative.layers import lora as lora_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.layers.rotary_position_embedding as rotary_pos_emb
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
 from torch import nn
 
+
 TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
     ff_up_proj="model.layers.{}.mlp.up_proj",
     ff_down_proj="model.layers.{}.mlp.down_proj",
@@ -86,6 +87,13 @@ class DecoderOnlyModel(nn.Module):
         config.embedding_dim,
         config.final_norm_config,
     )
+    # ROPE parameters for all attn_configs are the same. Take the first one.
+    attn_config = config.block_config(0).attn_config
+    self.rope_cache = attn_utils.build_rope_cache(
+        size=config.kv_cache_max,
+        dim=int(attn_config.rotary_percentage * attn_config.head_dim),
+        base=attn_config.rotary_base,
+    )
     self.mask_cache = attn_utils.build_causal_mask_cache(
         size=config.kv_cache_max,
     )
@@ -97,6 +105,7 @@ class DecoderOnlyModel(nn.Module):
       tokens: torch.Tensor,
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
+      lora: Optional[lora_utils.LoRA] = None,
       export_config: Optional[ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     _, seq_len = tokens.size()
@@ -107,28 +116,23 @@ class DecoderOnlyModel(nn.Module):
 
     # token embeddings of shape (b, t, n_embd)
     input_embeds = self.tok_embedding(tokens)
-
-    # ROPE parameters for all attn_configs are the same. Take the first one.
-    attn_config = self.config.block_config(0).attn_config
-    n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
-    rope = rotary_pos_emb.build_rope(
-        input_pos, n_elem, attn_config.head_dim, attn_config.rotary_base
-    )
-
+    cos, sin = self.rope_cache
+    rope = (cos.index_select(0, input_pos), sin.index_select(0, input_pos))
     mask = self.mask_cache.index_select(2, input_pos)
     mask = mask[:, :, :, : self.config.kv_cache_max]
 
-    return self._forward_with_embeds(
-        input_embeds, rope, mask, input_pos, kv_cache, export_config
+    return self.forward_with_embeds(
+        input_embeds, rope, mask, input_pos, kv_cache, lora, export_config
     )
 
-  def _forward_with_embeds(
+  def forward_with_embeds(
       self,
       input_embeds: torch.Tensor,
       rope: Tuple[torch.Tensor, torch.Tensor],
       mask: torch.Tensor,
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
+      lora: Optional[lora_utils.LoRA] = None,
       export_config: Optional[ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     """Forwards the model with input embeddings."""
@@ -141,13 +145,14 @@ class DecoderOnlyModel(nn.Module):
     if self.config.embedding_scale is not None:
       x = x * self.config.embedding_scale
 
-    updated_kv_entries = []
+    updated_kv_entires = []
     for i, block in enumerate(self.transformer_blocks):
       kv_entry = kv_cache.caches[i] if kv_cache else None
-      x, kv_entry = block(x, rope, mask, input_pos, kv_entry)
+      lora_adapter = lora.adapters[i] if lora else None
+      x, kv_entry = block(x, rope, mask, input_pos, kv_entry, lora_adapter)
       if kv_entry:
-        updated_kv_entries.append(kv_entry)
-    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entries))
+        updated_kv_entires.append(kv_entry)
+    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entires))
 
     if export_config is not None:
       if (

ai_edge_torch/generative/utilities/verifier.py

@@ -16,7 +16,7 @@
 """Common utility functions to verify the reauthored models."""
 
 import logging
-from typing import Any,List
+from typing import Any, List, Optional
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.generative.utilities.model_builder import ExportConfig
@@ -134,7 +134,7 @@ class ReauthoredModelWrapper(ModelWrapper):
       prompts: torch.Tensor,
       max_new_tokens: int,
       pixel_values: torch.Tensor = None,
-      eos_token_id: int = 1,
+      eos_token_id: Optional[int] = None,
   ) -> torch.IntTensor:
     input_ids = prompts[0].int().tolist()
     tokens = torch.tensor([input_ids])
@@ -146,7 +146,7 @@ class ReauthoredModelWrapper(ModelWrapper):
       )
       generated_token = logits[0][-1].argmax().item()
       input_ids.append(generated_token)
-      if generated_token == eos_token_id:
+      if eos_token_id is not None and generated_token == eos_token_id:
         break
       tokens = torch.tensor([[generated_token]])
       input_pos = torch.tensor([len(input_ids) - 1])
@@ -253,7 +253,7 @@ def verify_model_with_prompts(
   outputs_reauthored = reauthored_model.generate(
       prompt_tokens,
       max_new_tokens,
-      eos_token_id=tokenizer.tokenizer.eos_token_id,
+      eos_token_id=getattr(tokenizer.tokenizer, "eos_token_id", None),
   )
   response_reauthored = tokenizer.decode(outputs_reauthored[0])
   logging.info("outputs from reauthored model: [[%s]]", response_reauthored)

ai_edge_torch/odml_torch/_torch_future.py

@@ -73,3 +73,16 @@ def safe_run_decompositions(exported_program, decomp_table=None):
       node.target = lambda self, size: torch.reshape(self.contiguous(), size)
 
   return exported_program.run_decompositions(decomp_table)
+
+
+def dummy_decomp_table():
+  """Build dummy decomp table for run_decompositions without any decompositions.
+
+  Compatible for torch<=2.5.
+
+  Returns:
+    Decomp table for ExportedProgram.run_decompositions.
+  """
+  return {
+      torch._ops.OperatorBase(): lambda: None,
+  }

ai_edge_torch/odml_torch/export.py

@@ -238,6 +238,9 @@ def _convert_i64_to_i32(exported_program: torch.export.ExportedProgram):
   def in_i32(x: int):
     return -2147483648 <= x <= 2147483647
 
+  def to_int32(x: torch.Tensor):
+    return torch.ops.aten._to_copy.default(x, dtype=torch.int32)
+
   def rewrite_arange(node: torch.fx.Node):
     tensor_meta = node.meta.get("tensor_meta", None)
     if not tensor_meta:
@@ -249,7 +252,7 @@ def _convert_i64_to_i32(exported_program: torch.export.ExportedProgram):
     if not (in_i32(start) and in_i32(end)):
       return
     op = node.target
-    node.target = lambda *args, **kwargs: op(*args, **kwargs).type(torch.int32)
+    node.target = lambda *args, **kwargs: to_int32(op(*args, **kwargs))
 
   graph_module = exported_program.graph_module
   for node in graph_module.graph.nodes:
@@ -305,8 +308,9 @@ def exported_program_to_mlir(
 
   _convert_i64_to_i32(exported_program)
 
+  # No decompositions but just retracing/cananicalization.
   exported_program = _torch_future.safe_run_decompositions(
-      exported_program, lowerings.decompositions()
+      exported_program, _torch_future.dummy_decomp_table()
   )
 
   # Passes below mutate the exported program to a state not executable by torch.

ai_edge_torch/odml_torch/lowerings/decomp.py

@@ -55,6 +55,10 @@ def decompositions():
       ],
   )
 
+  # Override noop aten op decompositions for faster run_decompositions.
+  decompositions[torch.ops.aten.alias.default] = lambda x: x
+  decompositions[torch.ops.aten.detach.default] = lambda x: x
+
   # Override _safe_softmax decompositions with regular softmax.
   # _safe_softmax introduces additional check-select ops to guard extreme
   # input values to softmax, which could make the converted model inefficient

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20250105.dist-info → ai_edge_torch_nightly-0.3.0.dev20250108.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20250105
+Version: 0.3.0.dev20250108
 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