ai-edge-torch-nightly 0.3.0.dev20250107__py3-none-any.whl → 0.3.0.dev20250109__py3-none-any.whl

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/test/test_model_conversion_large.py

@@ -150,6 +150,16 @@ class TestModelConversion(googletest.TestCase):
       ai_edge_torch.config.in_oss,
       reason="tests with custom ops are not supported in oss",
   )
+
+  def test_smollm2(self):
+    config = smollm.get_fake_model_config_v2()
+    pytorch_model = smollm.SmolLM2(config).eval()
+    self._test_model(config, pytorch_model, "prefill", atol=1e-4, rtol=1e-5)
+  @googletest.skipIf(
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
+  )
+
   def test_openelm(self):
     config = openelm.get_fake_model_config()
     pytorch_model = openelm.OpenELM(config).eval()

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,12 +22,15 @@ 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",
@@ -85,13 +88,6 @@ 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,
     )
@@ -103,6 +99,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()
@@ -113,13 +110,23 @@ class DecoderOnlyModel(nn.Module):
 
     # token embeddings of shape (b, t, n_embd)
     input_embeds = self.tok_embedding(tokens)
-    cos, sin = self.rope_cache
-    rope = (cos.index_select(0, input_pos), sin.index_select(0, input_pos))
+
+    # 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 = self.config.build_rope(
+        input_pos=input_pos,
+        n_elem=n_elem,
+        base=attn_config.rotary_base,
+        head_dim=attn_config.head_dim,
+        # input_pos=input_pos, n_elem=n_elem, base=attn_config.rotary_base
+    )
+
     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
+        input_embeds, rope, mask, input_pos, kv_cache, lora, export_config
     )
 
   def forward_with_embeds(
@@ -129,6 +136,7 @@ class DecoderOnlyModel(nn.Module):
       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 +149,14 @@ class DecoderOnlyModel(nn.Module):
     if self.config.embedding_scale is not None:
       x = x * self.config.embedding_scale
 
-    updated_kv_entires = []
+    updated_kv_entries = []
     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_entires.append(kv_entry)
-    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entires))
+        updated_kv_entries.append(kv_entry)
+    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entries))
 
     if export_config is not None:
       if (

ai_edge_torch/hlfb/mark_pattern/__init__.py

@@ -17,7 +17,7 @@ from typing import Any
 import uuid
 
 from ai_edge_torch import lowertools
-from ai_edge_torch.hlfb.mark_pattern import passes
+from ai_edge_torch.hlfb.mark_pattern import fx_utils
 from ai_edge_torch.hlfb.mark_pattern import pattern as pattern_module
 import torch
 
@@ -87,7 +87,7 @@ def mark_pattern(
     m.meta["ORIGINAL_NODE"] = n
 
   # Sanitize graph_module to match in the same way as pattern's graph_module.
-  graph_module_to_match = passes.remove_clone_ops(graph_module_to_match)
+  graph_module_to_match = fx_utils.remove_clone_ops(graph_module_to_match)
 
   match_with_attrs = pattern.match(graph_module_to_match)
 
@@ -111,13 +111,25 @@ def mark_pattern(
           is_input=True,
       )
 
-      # Only replace input by the marker node for those nodes used in the pattern.
+      # Only replace input by the marker node for those nodes used in the
+      # pattern.
       in_pattern_nodes = set(match.nodes_map.values())
       for user in input_node.users.keys():
-        if user in in_pattern_nodes:
-          user.meta["ORIGINAL_NODE"].replace_input_with(
-              input_node.meta["ORIGINAL_NODE"], new_input_node
-          )
+        if user not in in_pattern_nodes:
+          continue
+
+        user.meta["ORIGINAL_NODE"].replace_input_with(
+            input_node.meta["ORIGINAL_NODE"], new_input_node
+        )
+        # Pattern matching graph sanitization may remove clone ops, which means
+        # the user's input in the original graph may be a clone op. When
+        # replacing the input with the marker node, we need to further try
+        # replacing the input of the clone op that connects to the user.
+        for original_user_input in user.meta["ORIGINAL_NODE"].all_input_nodes:
+          if fx_utils.is_clone_op(original_user_input):
+            original_user_input.replace_input_with(
+                input_node.meta["ORIGINAL_NODE"], new_input_node
+            )
 
     for i, pattern_output_node in enumerate(pattern.output_nodes):
       output_node = match.nodes_map[pattern_output_node]

ai_edge_torch/hlfb/mark_pattern/{passes.py → fx_utils.py}

@@ -12,11 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Passes to clean up the model graph for pattern matching."""
+"""FX graph utilities for pattern matching clean ups."""
 
 import torch
 
 
+def is_clone_op(node: torch.fx.Node) -> bool:
+  """Checks if the node is a clone op."""
+  return (
+      node.op == "call_function" and node.target == torch.ops.aten.clone.default
+  )
+
+
 def remove_clone_ops(gm: torch.fx.GraphModule):
   """Removes clone ops from the graph.
 
@@ -32,7 +39,7 @@ def remove_clone_ops(gm: torch.fx.GraphModule):
     The graph module with clone ops removed.
   """
   for node in gm.graph.nodes:
-    if node.op == "call_function" and node.name.startswith("clone"):
+    if is_clone_op(node):
       node.replace_all_uses_with(node.args[0])
       gm.graph.erase_node(node)
 

ai_edge_torch/hlfb/mark_pattern/pattern.py

@@ -18,13 +18,14 @@ import dataclasses
 from typing import Any, Callable, Optional, Union
 
 from ai_edge_torch import fx_pass_base
-from ai_edge_torch.hlfb.mark_pattern import passes
+from ai_edge_torch.hlfb.mark_pattern import fx_utils
 import torch
-from torch.export.graph_signature import TensorArgument
-from torch.fx import Graph
-from torch.fx import GraphModule
-from torch.fx.passes.utils.matcher_utils import InternalMatch
-from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
+
+Graph = torch.fx.Graph
+GraphModule = torch.fx.GraphModule
+TensorArgument = torch.export.graph_signature.TensorArgument
+InternalMatch = torch.fx.passes.utils.matcher_utils.InternalMatch
+SubgraphMatcher = torch.fx.passes.utils.matcher_utils.SubgraphMatcher
 
 
 def _are_equal(x: Any, y: Any) -> bool:
@@ -219,8 +220,8 @@ class Pattern:
     # Sanitize graph_module for more precise pattern matching.
     # The graph_module to match against this pattern should apply equivalent
     # sanitization.
-    self.graph_module = passes.remove_clone_ops(self.graph_module)
-    self.graph_module = passes.remove_dangling_args(self.graph_module)
+    self.graph_module = fx_utils.remove_clone_ops(self.graph_module)
+    self.graph_module = fx_utils.remove_dangling_args(self.graph_module)
 
     # Builds list of ordered input and output nodes.
     self.graph_nodes_map = {}

ai_edge_torch/hlfb/test/test_mark_pattern.py

@@ -58,6 +58,32 @@ class TestMarkPattern(googletest.TestCase):
         {"stablehlo.custom_call @mark_tensor": 6},
     )
 
+  def test_mark_pattern_with_clone_inputs(self):
+
+    class TestModel(torch.nn.Module):
+
+      def forward(self, x):
+        return torch.ops.aten.clone.default(x * x) + x
+
+    pattern = pattern_module.Pattern(
+        "test.add",
+        lambda a, b: a + b,
+        export_args=(torch.rand(2, 2), torch.rand(2, 2)),
+    )
+
+    model = TestModel().eval()
+    args = (torch.rand(20, 20),)
+    exported_program = torch.export.export(model, args)
+    mark_pattern.mark_pattern(exported_program.graph_module, pattern)
+    mlir = _export_stablehlo_mlir(exported_program)
+
+    lowertools.assert_string_count(
+        self,
+        mlir,
+        {'stablehlo.composite "test.add"': 1},
+        {"stablehlo.custom_call @mark_tensor": 3},
+    )
+
   def test_mark_pattern_with_attr_builder(self):
     class TestModel(torch.nn.Module):
 

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.dev20250107"
+__version__ = "0.3.0.dev20250109"