ai-edge-torch-nightly 0.6.0.dev20250620__py3-none-any.whl → 0.6.0.dev20250816__py3-none-any.whl

ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py

@@ -16,6 +16,7 @@
 from typing import Any, Callable
 from ai_edge_torch import fx_infra
 from ai_edge_torch import lowertools
+from ai_edge_torch.odml_torch import optimization_barrier as optimization_barrier_lib
 import torch
 import torch.utils._pytree as pytree
 
@@ -276,6 +277,7 @@ def _aten_embedding(gm: torch.fx.GraphModule, node: torch.fx.Node):
     # Explicitly reshape back to the original shape. This places the ReshapeOp
     # outside of the HLFB.
     output = torch.reshape(output, (*(original_idx_shape), embedding_dim))
+    output, _ = optimization_barrier_lib.optimization_barrier(output, idx)
     return output
 
   node.target = embedding

ai_edge_torch/_convert/test/test_convert.py

@@ -576,6 +576,24 @@ class TestConvert(googletest.TestCase):
       self.fail(f"Conversion failed with bloat16 inputs: {err}")
     # pylint: enable=broad-except
 
+  def test_convert_model_with_torch_linspace_operation(self):
+    """Test converting a simple model with torch.linspace operation."""
+
+    class SampleModel(nn.Module):
+
+      def forward(self, x: torch.Tensor):
+        return torch.linspace(0.5, 10.5, steps=x.shape[0], dtype=torch.float64)
+
+    model = SampleModel().eval()
+    args = (torch.randint(0, 100, (10, 10), dtype=torch.int64),)
+
+    try:
+      # Expect this to fix the error during conversion
+      ai_edge_torch.convert(model, args)
+    except Exception as err:
+      self.fail(f"Conversion failed with int64 inputs: {err}")
+    # pylint: enable=broad-except
+
   def test_compile_model(self):
     """Tests AOT compilation of a simple Add module."""
 

ai_edge_torch/generative/examples/gemma/gemma1.py

@@ -23,7 +23,7 @@ import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
 from torch import nn
 
-TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+TENSOR_NAMES_FUSED_QKV = loading_utils.ModelLoader.TensorNames(
     ff_up_proj="model.layers.{}.mlp.up_proj",
     ff_down_proj="model.layers.{}.mlp.down_proj",
     ff_gate_proj="model.layers.{}.mlp.gate_proj",
@@ -36,6 +36,24 @@ TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
     lm_head=None,
 )
 
+TENSOR_NAMES_SEP_QKV = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj="model.layers.{}.mlp.up_proj",
+    ff_down_proj="model.layers.{}.mlp.down_proj",
+    ff_gate_proj="model.layers.{}.mlp.gate_proj",
+    attn_query_proj="model.layers.{}.self_attn.q_proj",
+    attn_key_proj="model.layers.{}.self_attn.k_proj",
+    attn_value_proj="model.layers.{}.self_attn.v_proj",
+    attn_output_proj="model.layers.{}.self_attn.o_proj",
+    pre_attn_norm="model.layers.{}.input_layernorm",
+    post_attn_norm="model.layers.{}.post_attention_layernorm",
+    embedding="model.embed_tokens",
+    final_norm="model.norm",
+)
+
+TENSOR_NAMES_DICT = {
+    "safetensors": TENSOR_NAMES_SEP_QKV,
+    "kaggle": TENSOR_NAMES_FUSED_QKV,
+}
 
 class Gemma1(model_builder.DecoderOnlyModel):
   """A Gemma1 model built from the Edge Generative API layers."""
@@ -94,11 +112,28 @@ def build_2b_model(
     custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
     mask_cache_size: int = 0,
 ) -> nn.Module:
-  return model_builder.build_decoder_only_model(
-      checkpoint_path=checkpoint_path,
-      config=get_model_config_2b(),
-      tensor_names=TENSOR_NAMES,
-      model_class=Gemma1,
-      custom_loader=custom_loader,
-      mask_cache_size=mask_cache_size,
+
+  # A list to store the reasons for each failure
+  key_errors = []
+
+  for tensor_names in TENSOR_NAMES_DICT.values():
+    try:
+      return model_builder.build_decoder_only_model(
+          checkpoint_path=checkpoint_path,
+          config=get_model_config_2b(),
+          tensor_names=tensor_names,
+          model_class=Gemma1,
+          custom_loader=custom_loader,
+          mask_cache_size=mask_cache_size,
+      )
+    except KeyError as ke:
+      # Store the specific key that was missing for later
+      key_errors.append(f"Missing key: {ke}")
+      continue
+
+  # If the loop finishes, raise an error with all the collected details
+  error_details = "\n".join(key_errors)
+  raise RuntimeError(
+      "Failed to build model after trying all configurations. "
+      f"Encountered the following errors:\n{error_details}"
   )

ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py

@@ -42,6 +42,14 @@ def main(_):
         ),
         mask_cache_size=converter.get_mask_cache_size_from_flags(),
     )
+  elif _MODEL_SIZE.value == '270m':
+    pytorch_model = gemma3.build_model_270m(
+        checkpoint_path,
+        custom_loader=loader.maybe_get_custom_loader(
+            checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+        ),
+        mask_cache_size=converter.get_mask_cache_size_from_flags(),
+    )
   else:
     raise ValueError(f'Unsupported model size: {_MODEL_SIZE.value}')
 

ai_edge_torch/generative/examples/gemma3/decoder.py

@@ -391,6 +391,60 @@ def get_decoder_config_1b() -> cfg.ModelConfig:
   return config
 
 
+def get_decoder_config_270m() -> cfg.ModelConfig:
+  """Returns the model config for a Gemma3 270M model."""
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_TANH),
+      intermediate_size=2048,
+      pre_ff_norm_config=norm_config,
+      post_ff_norm_config=norm_config,
+  )
+
+  def get_block_config(idx: int) -> cfg.TransformerBlockConfig:
+    attn_config = cfg.AttentionConfig(
+        num_heads=4,
+        head_dim=256,
+        num_query_groups=1,
+        rotary_base=1_000_000 if (idx + 1) % 6 == 0 else 10_000,
+        rotary_percentage=1.0,
+        qkv_transpose_before_split=True,
+        query_norm_config=norm_config,
+        key_norm_config=norm_config,
+        logit_softcap=None,
+        sliding_window_size=512,
+        attn_type=(
+            cfg.AttentionType.GLOBAL
+            if (idx + 1) % 6 == 0
+            else cfg.AttentionType.LOCAL_SLIDING
+        ),
+    )
+    return cfg.TransformerBlockConfig(
+        attn_config=attn_config,
+        ff_config=ff_config,
+        pre_attention_norm_config=norm_config,
+        post_attention_norm_config=norm_config,
+    )
+
+  num_layers = 18
+  embedding_dim = 640
+  config = cfg.ModelConfig(
+      vocab_size=262_144,
+      num_layers=num_layers,
+      max_seq_len=32_768,
+      embedding_dim=embedding_dim,
+      embedding_scale=embedding_dim**0.5,
+      block_configs=[get_block_config(i) for i in range(num_layers)],
+      final_norm_config=norm_config,
+      lm_head_use_bias=False,
+      final_logit_softcap=None,
+  )
+  return config
+
+
 def get_fake_decoder_config_1b() -> cfg.ModelConfig:
   """Returns a fake model config for a Gemma3 1B model."""
   config = get_decoder_config_1b()
@@ -427,3 +481,25 @@ def build_model_1b(
       )
     except KeyError as ke:
       continue
+
+
+def build_model_270m(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    mask_cache_size: int = 0,
+) -> nn.Module:
+  """Builds a Gemma3 270M model."""
+  # TODO(b/403644647): Better error handling for loading checkpoints with
+  # different tensor names.
+  for tensor_names in TENSOR_NAMES_DICT.values():
+    try:
+      return model_builder.build_decoder_only_model(
+          checkpoint_path=checkpoint_path,
+          config=get_decoder_config_270m(),
+          tensor_names=tensor_names,
+          model_class=Decoder,
+          custom_loader=custom_loader,
+          mask_cache_size=mask_cache_size,
+      )
+    except KeyError as _:
+      continue

ai_edge_torch/generative/examples/gemma3/gemma3.py

@@ -179,3 +179,22 @@ def build_model_1b(
   # TODO: Load the parameters of decoder from checkpoint.
   model.eval()
   return model
+
+
+def build_model_270m(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    mask_cache_size: int = 0,
+) -> decoder.Decoder:
+  """Builds a Gemma3 270M model."""
+  if checkpoint_path:
+    model = decoder.build_model_270m(
+        checkpoint_path, custom_loader, mask_cache_size
+    )
+  else:
+    config = decoder.get_decoder_config_270m()
+    model = decoder.Decoder(config, mask_cache_size)
+  # TODO: Load the parameters of decoder from checkpoint.
+  model.eval()
+  return model
+

ai_edge_torch/generative/layers/einsum.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 # Einsum layer implementation.
 
-from typing import Sequence
+from typing import Callable, Sequence
 import torch
 from torch import nn
 
@@ -22,7 +22,12 @@ from torch import nn
 class Einsum(nn.Module):
   """Einsum layer wrapping over torch.einsum."""
 
-  def __init__(self, shape: Sequence[int], einsum_str: str):
+  def __init__(
+      self,
+      shape: Sequence[int],
+      einsum_str: str,
+      init_fn: Callable[..., torch.Tensor] = lambda *args, **kwargs: None,
+  ):
     super().__init__()
     self.shape = shape
     self.einsum_str = einsum_str
@@ -30,6 +35,7 @@ class Einsum(nn.Module):
         torch.empty(shape, dtype=torch.float32),
         requires_grad=False,
     )
+    init_fn(self.w)
     self.einsum_fn = lambda x: torch.einsum(einsum_str, x, self.w)
 
   def forward(self, x: torch.Tensor) -> torch.Tensor:

ai_edge_torch/generative/quantize/quant_recipes.py

@@ -76,7 +76,6 @@ def all_supported_int4_dynamic_block_recipe(
           default=quant_recipe_utils.create_layer_quant_int4_dynamic_block(
               block_size
           ),
-          embedding=quant_recipe_utils.create_layer_quant_int8_dynamic(),
           _model_config=mcfg,
       )
   )

ai_edge_torch/lowertools/translate_recipe.py

@@ -25,13 +25,13 @@ _OpQuantConfig = quantizer.qtyping.OpQuantizationConfig
 
 _DEFAULT_REGEX_STR = '.*'
 _SINGULAR_TRANSFORMER_BLOCK_REGEX_STR = 'transformer_block'
-_IDX_TRANSFORMER_BLOCKS_REGEX_STR = 'transformer_blocks\[{}\]'
+_IDX_TRANSFORMER_BLOCKS_REGEX_STR = r'transformer_blocks\[{}\]'
 _ATTENTION_REGEX_STR = 'ai_edge_torch.generative.layers.attention'
 _FEEDFORWARD_REGEX_STR = 'ai_edge_torch.generative.layers.feed_forward'
 _EMBEDDING_REGEX_STR = 'Embedding_tok_embedding'
 # TODO: b/415833584 - Improve the regex for pre-softmax layer.
 _DECODE_LOGITS_REGEX_STR = 'StatefulPartitionedCall'
-_ANY_TWO_DIGITS_REGEX_STR = '\d{1,2}'
+_ANY_TWO_DIGITS_REGEX_STR = r'\d{1,2}'
 
 
 def _get_nbits_from_dtype(dtype: quant_attrs.Dtype) -> int:

ai_edge_torch/odml_torch/export.py

@@ -21,7 +21,6 @@ import operator
 from typing import Any, Callable, Optional
 
 from ai_edge_torch import fx_infra
-from jax.lib import xla_extension
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import func
 from jax._src.lib.mlir.dialects import hlo as stablehlo
@@ -35,6 +34,8 @@ from . import lowerings
 
 LoweringContext = lowerings.context.LoweringContext
 
+from jaxlib._jax.mlir import serialize_portable_artifact
+
 
 def _build_flat_inputs(exported_program: torch.export.ExportedProgram):
   """Build flattened inputs and metadata from exported program's signature."""
@@ -233,7 +234,7 @@ class MlirLowered:
       target_version = stablehlo.get_version_from_compatibility_requirement(
           stablehlo.StablehloCompatibilityRequirement.WEEK_12
       )
-    module_bytecode = xla_extension.mlir.serialize_portable_artifact(
+    module_bytecode = serialize_portable_artifact(
         self.module_bytecode, target_version
     )
     return module_bytecode

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -331,6 +331,18 @@ def _aten_sym_size_int(lctx, x: ir.Value, dim: int):
   return stablehlo.get_dimension_size(x, dim)
 
 
+# Lowering for the addition operator (`+`).
+# Handles cases where one operand is an integer (scalar) and the other is a
+# tensor, broadcasting the scalar to the tensor's shape before addition.
+@lower(operator.add)
+def _operator_add(lctx, self: int | ir.Value, other: int | ir.Value):
+  if isinstance(self, int) and isinstance(other, ir.Value):
+    self = utils.splat(self, other.type.element_type, other.type.shape)
+  if isinstance(other, int) and isinstance(self, ir.Value):
+    other = utils.splat(other, self.type.element_type, self.type.shape)
+  return stablehlo.add(self, other)
+
+
 # Lowering for the subtraction operator (`-`).
 # Handles cases where one operand is an integer (scalar) and the other is a
 # tensor, broadcasting the scalar to the tensor's shape before subtraction.

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -78,8 +78,6 @@ lower_by_torch_xla2(torch.ops.aten._unsafe_index)
 lower_by_torch_xla2(torch.ops.aten._unsafe_view)
 lower_by_torch_xla2(torch.ops.aten.acos)
 lower_by_torch_xla2(torch.ops.aten.acosh)
-lower_by_torch_xla2(torch.ops.aten.add.Scalar)
-lower_by_torch_xla2(torch.ops.aten.add.Tensor)
 lower_by_torch_xla2(torch.ops.aten.addbmm.default)
 lower_by_torch_xla2(torch.ops.aten.addmm)
 lower_by_torch_xla2(torch.ops.aten.addmv)
@@ -159,7 +157,6 @@ lower_by_torch_xla2(torch.ops.aten.logical_and)
 lower_by_torch_xla2(torch.ops.aten.logical_not)
 lower_by_torch_xla2(torch.ops.aten.logical_or)
 lower_by_torch_xla2(torch.ops.aten.logical_xor)
-lower_by_torch_xla2(torch.ops.aten.lt)
 lower_by_torch_xla2(torch.ops.aten.max)
 lower_by_torch_xla2(torch.ops.aten.max_pool2d_with_indices)
 lower_by_torch_xla2(torch.ops.aten.max_pool2d_with_indices_backward)
@@ -170,8 +167,6 @@ lower_by_torch_xla2(torch.ops.aten.mean)
 lower_by_torch_xla2(torch.ops.aten.min)
 lower_by_torch_xla2(torch.ops.aten.minimum)
 lower_by_torch_xla2(torch.ops.aten.mm)
-lower_by_torch_xla2(torch.ops.aten.mul.Scalar)
-lower_by_torch_xla2(torch.ops.aten.mul.Tensor)
 lower_by_torch_xla2(torch.ops.aten.native_batch_norm)
 lower_by_torch_xla2(torch.ops.aten.native_layer_norm_backward)
 lower_by_torch_xla2(torch.ops.aten.ne)
@@ -215,8 +210,6 @@ lower_by_torch_xla2(torch.ops.aten.sqrt)
 lower_by_torch_xla2(torch.ops.aten.squeeze)
 lower_by_torch_xla2(torch.ops.aten.squeeze_copy)
 lower_by_torch_xla2(torch.ops.aten.stack)
-lower_by_torch_xla2(torch.ops.aten.sub.Scalar)
-lower_by_torch_xla2(torch.ops.aten.sub.Tensor)
 lower_by_torch_xla2(torch.ops.aten.sum)
 lower_by_torch_xla2(torch.ops.aten.t)
 lower_by_torch_xla2(torch.ops.aten.tan)
@@ -244,7 +237,6 @@ lower_by_torch_xla2(torch.ops.aten.view_as_real)
 lower_by_torch_xla2(torch.ops.aten.view_copy)
 lower_by_torch_xla2(torch.ops.aten.where.ScalarOther)
 lower_by_torch_xla2(torch.ops.aten.where.ScalarSelf)
-lower_by_torch_xla2(torch.ops.aten.where.self)
 lower_by_torch_xla2(torch.ops.prims.broadcast_in_dim)
 lower_by_torch_xla2(torch.ops.prims.var)
 
@@ -259,6 +251,149 @@ def _aten_copy(self, src, **kwargs):
   return _TORCH_XLA2_IMPLS[torch.ops.aten.copy](self, src)
 
 
+@registry.lower(torch.ops.aten.add.Scalar)
+def _aten_add_scalar(lctx: LoweringContext, self, other):
+  _log_usage(torch.ops.aten.add.Scalar)
+
+  @jax_bridge.wrap
+  def jax_lowering(self, other):
+    other_dtype = jnp.result_type(other)
+    promoted_type = jnp.promote_types(self.dtype, other_dtype)
+    if promoted_type == jnp.float64:
+      promoted_type = jnp.float32
+    return jnp.add(
+        self.astype(promoted_type), jnp.array(other, dtype=promoted_type)
+    )
+
+  return jax_lowering(lctx, self, other)
+
+
+@registry.lower(torch.ops.aten.add.Tensor)
+def _aten_add_tensor(lctx: LoweringContext, self, other):
+  _log_usage(torch.ops.aten.add.Tensor)
+
+  @jax_bridge.wrap
+  def jax_lowering(self, other):
+    promoted_type = jnp.promote_types(self.dtype, other.dtype)
+    if promoted_type == jnp.float64:
+      promoted_type = jnp.float32
+    return jnp.add(self.astype(promoted_type), other.astype(promoted_type))
+
+  return jax_lowering(lctx, self, other)
+
+
+@registry.lower(torch.ops.aten.sub.Scalar)
+def _aten_sub_scalar(lctx: LoweringContext, self, other):
+  _log_usage(torch.ops.aten.sub.Scalar)
+
+  @jax_bridge.wrap
+  def jax_lowering(self, other):
+    other_dtype = jnp.result_type(other)
+    promoted_type = jnp.promote_types(self.dtype, other_dtype)
+    if promoted_type == jnp.float64:
+      promoted_type = jnp.float32
+    return jnp.subtract(
+        self.astype(promoted_type), jnp.array(other, dtype=promoted_type)
+    )
+
+  return jax_lowering(lctx, self, other)
+
+
+@registry.lower(torch.ops.aten.sub.Tensor)
+def _aten_sub_tensor(lctx: LoweringContext, self, other):
+  _log_usage(torch.ops.aten.sub.Tensor)
+
+  @jax_bridge.wrap
+  def jax_lowering(self, other):
+    promoted_type = jnp.promote_types(self.dtype, other.dtype)
+    if promoted_type == jnp.float64:
+      promoted_type = jnp.float32
+    return jnp.subtract(self.astype(promoted_type), other.astype(promoted_type))
+
+  return jax_lowering(lctx, self, other)
+
+
+@registry.lower(torch.ops.aten.lt.Scalar)
+def _aten_lt_scalar(lctx: LoweringContext, self, other):
+  _log_usage(torch.ops.aten.lt.Scalar)
+
+  @jax_bridge.wrap
+  def jax_lowering(self, other):
+    other_dtype = jnp.result_type(other)
+    promoted_type = jnp.promote_types(self.dtype, other_dtype)
+    if promoted_type == jnp.float64:
+      promoted_type = jnp.float32
+    return jnp.less(
+        self.astype(promoted_type), jnp.array(other, dtype=promoted_type)
+    )
+
+  return jax_lowering(lctx, self, other)
+
+
+@registry.lower(torch.ops.aten.lt.Tensor)
+def _aten_lt_tensor(lctx: LoweringContext, self, other):
+  _log_usage(torch.ops.aten.lt.Tensor)
+
+  @jax_bridge.wrap
+  def jax_lowering(self, other):
+    promoted_type = jnp.promote_types(self.dtype, other.dtype)
+    return jnp.less(self.astype(promoted_type), other.astype(promoted_type))
+
+  return jax_lowering(lctx, self, other)
+
+
+@registry.lower(torch.ops.aten.mul.Scalar)
+def _aten_mul_scalar(lctx: LoweringContext, self, other):
+  _log_usage(torch.ops.aten.mul.Scalar)
+
+  @jax_bridge.wrap
+  def jax_lowering(self, other):
+    other_dtype = jnp.result_type(other)
+    promoted_type = jnp.promote_types(self.dtype, other_dtype)
+    if promoted_type == jnp.float64:
+      promoted_type = jnp.float32
+    return jnp.multiply(
+        self.astype(promoted_type), jnp.array(other, dtype=promoted_type)
+    )
+
+  return jax_lowering(lctx, self, other)
+
+
+@registry.lower(torch.ops.aten.mul.Tensor)
+def _aten_mul_tensor(lctx: LoweringContext, self, other):
+  _log_usage(torch.ops.aten.mul.Tensor)
+
+  @jax_bridge.wrap
+  def jax_lowering(self, other):
+    other_dtype = jnp.result_type(other)
+    promoted_type = jnp.promote_types(self.dtype, other_dtype)
+    if promoted_type == jnp.float64:
+      promoted_type = jnp.float32
+    return jnp.multiply(
+        self.astype(promoted_type), jnp.array(other, dtype=promoted_type)
+    )
+
+  return jax_lowering(lctx, self, other)
+
+
+@registry.lower(torch.ops.aten.where.self)
+def _aten_where_self(lctx: LoweringContext, condition, self, other):
+  _log_usage(torch.ops.aten.where.self)
+
+  @jax_bridge.wrap
+  def jax_lowering(condition, self, other):
+    promoted_type = jnp.promote_types(self.dtype, other.dtype)
+    if promoted_type == jnp.float64:
+      promoted_type = jnp.float32
+    return jnp.where(
+        condition,
+        self.astype(promoted_type),
+        other.astype(promoted_type),
+    )
+
+  return jax_lowering(lctx, condition, self, other)
+
+
 # Schema:
 #   - aten::einsum(str equation, Tensor[] tensors, *, int[]? path=None)
 #       -> Tensor

ai_edge_torch/odml_torch/optimization_barrier.py

@@ -0,0 +1,71 @@
+# 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.
+# ==============================================================================
+"""Optimization barrier op definition and lowering."""
+
+from ai_edge_torch.odml_torch import _torch_library
+from ai_edge_torch.odml_torch.lowerings import registry
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import torch
+import torch.utils._pytree as pytree
+
+_torch_library.ODML_TORCH_LIB.define(
+    "optimization_barrier(Tensor[] inputs) -> Tensor[]"
+)
+
+optimization_barrier_op = torch.ops.odml_torch.optimization_barrier.default
+
+
+def optimization_barrier(*inputs: pytree.PyTree):
+  """Apply optimization barrier to the tensors nested within arbitrary pytrees.
+
+  Args:
+    *inputs: A list of tensors or tensor pytrees.
+
+  Returns:
+    The tensors after optimization barrier in the same pytrees structures.
+  """
+  if len(inputs) == 1:
+    inputs = inputs[0]
+  tensors, spec = pytree.tree_flatten(inputs)
+  tensors = optimization_barrier_op(tuple(tensors))
+  outputs = pytree.tree_unflatten(tensors, spec)
+  return outputs
+
+
+@torch.library.impl(
+    _torch_library.ODML_TORCH_LIB,
+    "optimization_barrier",
+    "CompositeExplicitAutograd",
+)
+def _optimization_barrier_impl(inputs: tuple[torch.Tensor, ...]):
+  return tuple(inputs)
+
+
+@torch.library.impl(
+    _torch_library.ODML_TORCH_LIB,
+    "optimization_barrier",
+    "Meta",
+)
+def _optimization_barrier_fake(inputs: tuple[torch.Tensor, ...]):
+  return tuple([torch.empty_like(x) for x in inputs])
+
+
+@registry.lower(torch.ops.odml_torch.optimization_barrier.default)
+def _optimization_barrier_lowering(
+    lctx, inputs: tuple[ir.Value, ...]
+) -> ir.Value:
+  del lctx
+  return stablehlo.optimization_barrier(inputs)

ai_edge_torch/version.py

@@ -15,4 +15,4 @@
 
 # The next version of ai-edge-torch.
 # The minor version code should be bumped after every release.
-__version__ = "0.6.0.dev20250620"
+__version__ = "0.6.0.dev20250816"

{ai_edge_torch_nightly-0.6.0.dev20250620.dist-info → ai_edge_torch_nightly-0.6.0.dev20250816.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.6.0.dev20250620
+Version: 0.6.0.dev20250816
 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
@@ -31,7 +31,7 @@ Requires-Dist: transformers
 Requires-Dist: kagglehub
 Requires-Dist: tabulate
 Requires-Dist: torch>=2.4.0
-Requires-Dist: tf-nightly>=2.19.0.dev20250101
+Requires-Dist: tf-nightly<=2.20.0.dev20250615
 Requires-Dist: ai-edge-litert-nightly
 Requires-Dist: ai-edge-quantizer-nightly
 Requires-Dist: jax

{ai_edge_torch_nightly-0.6.0.dev20250620.dist-info → ai_edge_torch_nightly-0.6.0.dev20250816.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=lemyLCNoGYRnJsmDuGZu7qOqLbLqG6CGDFtu3ue1syU,129
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=A7loFu8jE9CsXsfMmHYZ-KDFJiaD8Kkqwm_9d3IVzk0,5638
-ai_edge_torch/version.py,sha256=L-kCkN9vMaBYCuarY-Y8kGgttAEHZdyPWaupqKVMLiA,806
+ai_edge_torch/version.py,sha256=l7eZe2YNZxV10qw-t4qOqrQ97K8OtzpEfzORpOBAi88,806
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=iQk3R-pLq4c1nfLqPB4xTRj78gghxPGzJCJtILLdg5o,6123
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -10,7 +10,7 @@ ai_edge_torch/_convert/converter.py,sha256=6MLKELzAwFoiXv-b7KRYi7gc7Z57XOeowcz9A
 ai_edge_torch/_convert/signature.py,sha256=-YKJdLk-eNEHfhdPCtcQVtZf915SoVePEFxKXPPf16c,2572
 ai_edge_torch/_convert/to_channel_last_io.py,sha256=_31phf7TYgZY2ftpNbrdlB1RhDium1lz_BXEQ6IsMFc,2893
 ai_edge_torch/_convert/fx_passes/__init__.py,sha256=uHek7k9KIW3kaEM_lcygbukJ69JLjm-xnYUWzAEIZII,1345
-ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py,sha256=dgUO-lI9Id9hIOHP5XmegVlu5Fl79GR4_b-lDUehzoo,11428
+ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py,sha256=D_7RQp3F7XzJZ3d0Sgay7hf_oz3IhKZCZ-F-sDH5LEY,11589
 ai_edge_torch/_convert/fx_passes/cast_inputs_bf16_to_f32_pass.py,sha256=90YxLVAAkiA3qKr4Um__JmPeC1bTeA2PxBCj0GETq1Q,1748
 ai_edge_torch/_convert/fx_passes/eliminate_dead_code_pass.py,sha256=jMl9YHIbx08KQHbp9UgDnxviUUWiN-FSsiUgR2HCT5s,1576
 ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=Z6E3U7SYZvMl3Ivpqa3burVOLKFndEZuNmWKNxjq2mM,2386
@@ -27,7 +27,7 @@ ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitio
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py,sha256=L_x8BrF7UDah-SYl-pG11I6CIckdU9kBTUHcmwW4cts,2420
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=oXbr9G5Jc21xd1dr2CDrp774I4crs0_kkN490K5fNn0,7312
 ai_edge_torch/_convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/test/test_convert.py,sha256=yQC0WZk_gzReguTOfgWWodK71jnfMjYoRF29_Kafnuw,18692
+ai_edge_torch/_convert/test/test_convert.py,sha256=ioeEpeS07NAs14-nHiPI-6lLTtnALxl8uNtTKvoHdgE,19316
 ai_edge_torch/_convert/test/test_convert_composites.py,sha256=BCIODgxMI_3MxMLfNWYMGjcz-al-J3z5eDHCiZJXNwY,7992
 ai_edge_torch/_convert/test/test_convert_multisig.py,sha256=6_C2R9--KyNR7_oezZIAfyTSR97tOeEWy4XGcbSxBDE,5778
 ai_edge_torch/_convert/test/test_to_channel_last_io.py,sha256=1o-gUiwzIuO67FNAJ8DeyKv8fVUeZVNNNwofNVDjYeU,3024
@@ -65,15 +65,15 @@ ai_edge_torch/generative/examples/deepseek/verify_util.py,sha256=xf2-jhv-ORwr_-R
 ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=m5N3MxNX4Yu8-0vXRszkMgfVuFKN6gskVoUIGqonJFk,1814
 ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=fR4869w1UZIuUVGLDdPof0IDEMDhqn2ej5kl7QW9vNA,1882
-ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=TH9XQAp5p4S829XbaWbJQZBwB18WizDRIQMsUkKqj38,3377
+ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=8oLB0PaSrHMy0tDT--qPIEjTm0qNT16kqRBCHqAALSc,4674
 ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=E6jotWYYIx6SXUWqurKWjiZpbfj_M2jJrBc2rQ90z1s,11782
 ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=USyist332qZkhCBof2tJwQSqtnKjTQsKAK_jCE_CO2U,1853
 ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=eAM7EVVMW-QCqjeZEss7TOkVKArgUs1La51LAC-5a9A,1962
 ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=b12naCj4zZxOjkIKrd08qovtajYuX-Ba3fbrv6kkDZs,8410
 ai_edge_torch/generative/examples/gemma3/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
-ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=UEDNN3JmI31WfE2pvacxeJpqumKK86L2dEus3yTURaY,2114
-ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=1UVv9SFFg5degX3wf-Fefx7nor1AzJj2NWBVuo8bRnM,15540
-ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=fFMyIS8si3GdwW8EsdhYk1OKyg_27xDv1HTQ2Gv4N8E,6616
+ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=b2l8R3-3xhurmXn3LRZ-C12J11C8tx77NXKze6jFYbA,2425
+ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=DEkOnfHZJPaxbjwegarRjGpBa7gr6IuWMf8xgZP2mtU,17951
+ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=PAtuqvBkq6TeNFImoJfCaLX5pZCBOcaQjFvwur_XQJ4,7141
 ai_edge_torch/generative/examples/gemma3/image_encoder.py,sha256=OCMIAQfNmPR4uQUAtlYL6j4xkG0dw2Ays4-lnThcWqQ,5110
 ai_edge_torch/generative/examples/gemma3/verify_gemma3.py,sha256=v8oNXFICmVOtQxfO7IhZ8GnbvotEkDi9lzYHjoQyOso,2464
 ai_edge_torch/generative/examples/gemma3/verify_util.py,sha256=5OmUwz38kVHYLA-v8U8evvDN9da2WioZtGo-XK6yq1o,10067
@@ -174,7 +174,7 @@ ai_edge_torch/generative/layers/attention_test.py,sha256=9v8v96TLyFPdqxEylU1JOAe
 ai_edge_torch/generative/layers/attention_utils.py,sha256=2qfg7Tzk9ikKph5w3geOHC1I6EyOCdDsWXMr7F7IOZM,7630
 ai_edge_torch/generative/layers/attention_utils_test.py,sha256=22gQ1gcRPkwqFG3_p82GZfRKVE3udEssSy58wNOqv0w,2431
 ai_edge_torch/generative/layers/builder.py,sha256=2bUgkyowDkDznkF8XaHyZs4nowHr1QEHYLM7pMaFmIk,4921
-ai_edge_torch/generative/layers/einsum.py,sha256=EsZSWNVWUs0-1plp4TBnhP4ZhaRDBa2VlDO6hWpUAqU,1288
+ai_edge_torch/generative/layers/einsum.py,sha256=LH4CNHr-pFfLUuCpwbYL3GpoAMgHJ4nLju3XCqA4VwM,1416
 ai_edge_torch/generative/layers/einsum_test.py,sha256=ltIE773bvvNLv_9aLQxFwe1MgQ762sez0c5E2tejxuA,1079
 ai_edge_torch/generative/layers/feed_forward.py,sha256=_GmtHxwL068l9gh_F_WFcFk7La-Tl5SfoQ9v2hMabZM,5541
 ai_edge_torch/generative/layers/feed_forward_test.py,sha256=Y5l1eC9NgfYixHcfIfE1W4FGh7oC-9UGGyHdKS9tQKc,1880
@@ -196,7 +196,7 @@ ai_edge_torch/generative/quantize/example.py,sha256=1lfVNUd2cEyRUnoZ7BLbRJ9IN-FT
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=plMsd7JBi98r2NHsAdMdvS6TPTXAoRFLCwOXu8H3-24,2004
 ai_edge_torch/generative/quantize/quant_recipe.py,sha256=CEW-ewHxwb59x_GISx4jr7WMihvn-jKWVcBonllzDS4,5724
 ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=h3k_na6rbR08Ip79-2JbkeH8RDk_rrnEGiytuzFDhqc,2678
-ai_edge_torch/generative/quantize/quant_recipes.py,sha256=45DJfcQXZ1FA1qI4LgYoYE4UD4yvfIYoY9LgYTeKFVw,2845
+ai_edge_torch/generative/quantize/quant_recipes.py,sha256=mmPsgQ2vloMWflqJ6ALmD1lANacTkWotEhurb5foK30,2771
 ai_edge_torch/generative/quantize/supported_schemes.py,sha256=TwR2FpQuBEORy6FshEyHNBMKARWlA2MVtTfX9tXV5aE,1488
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/test/test_custom_dus.py,sha256=MjIhTvkTko872M35XMciobvICcDWTcIDJ3rociko-wM,3267
@@ -231,12 +231,13 @@ ai_edge_torch/lowertools/common_utils.py,sha256=4HQtquPZ6oiId8vR_1ykW_uK4ELnyo5z
 ai_edge_torch/lowertools/odml_torch_utils.py,sha256=QRuS7S5lULRWEh3J1sWIsnKh-rbX7rd9tt6JJHbMPfo,8317
 ai_edge_torch/lowertools/test_utils.py,sha256=mdxTlhqHABZEQ_GEmPFCL8LIAWtqRtYZUGdSY1ieZjw,1949
 ai_edge_torch/lowertools/torch_xla_utils.py,sha256=1EytIw2R6dthhLhf69wN1L9BaQTeybCD0wga-PhHcMI,9518
-ai_edge_torch/lowertools/translate_recipe.py,sha256=JNsRc1Jmpj5W6PBww8KRMkbtxcv7ssl8Rr1R3x5_7to,6283
+ai_edge_torch/lowertools/translate_recipe.py,sha256=Oavs0ENKVnIeB-WidXvokTPqNlFfOuP0GMV0d5RK2Rg,6285
 ai_edge_torch/odml_torch/__init__.py,sha256=S8jOzE9nLof-6es3XDiGJRN-9H_XTxsVm9dE7lD3RWo,812
 ai_edge_torch/odml_torch/_torch_future.py,sha256=jSYHf1CMTJzMizPMbu2b39hAt0ZTR6gQLq67GMe9KTo,2336
 ai_edge_torch/odml_torch/_torch_library.py,sha256=Lw1gqL2HWNRspdTwNhIkYAHDyafHedHtkXyKKxn-Wss,805
-ai_edge_torch/odml_torch/export.py,sha256=lbLpdGa8MDE8oWNA7aSV3tOCQ9P9I2Ox95dSPEssn-g,14930
+ai_edge_torch/odml_torch/export.py,sha256=FDseiAOOcgN8HOuwv6HT9sALZgXGj7vZi6kwZPWPF14,14935
 ai_edge_torch/odml_torch/export_utils.py,sha256=QeA37Irlty6AiIBuqmHmJgn3lqahBQ5xsh6IKRoKm1g,4774
+ai_edge_torch/odml_torch/optimization_barrier.py,sha256=2lmSiu5iXWLFWpupZHvsVeNYNzG5AVGSK3K_CNhS5Sk,2290
 ai_edge_torch/odml_torch/tf_integration.py,sha256=NN29WeXmHZ0S1RPDFHUnBi2DEjMvAtwczStPYIsQ1w8,4849
 ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xKA9rknO93JS9kw9w_c,778
 ai_edge_torch/odml_torch/composite/mark_tensor.py,sha256=U--rwl-XkWKgkdXCXDn6yySug8FR66o1YFUAIoSaWW4,3523
@@ -248,11 +249,11 @@ ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNi
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=LqwZ1vCJTSOzgzvH8LUAN-sAkF-l_pGj1AMEIzAqHCA,6638
 ai_edge_torch/odml_torch/jax_bridge/utils.py,sha256=T8isGc896VrHZ6c_L5pYmLpolQ7ibcOlgWfPuVFPzIg,2264
 ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=uJ-niilt1c-D6QJzLwgvCUf62le_JsxQTlqj_iP_Ps0,1009
-ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=VWb5HEeVljnuXi1eecKp1ieOIcBrSLlu7YIZnxnrozU,12198
+ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=sC4N5-7RS9yKecs97kM9J56enGvsZj1CJo7y79cuzRg,12784
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=Q0aDzyUcZMoSzSbOU-r3LJMgPe6fble0QwdYVIOHHHk,6887
 ai_edge_torch/odml_torch/lowerings/_decomp_registry.py,sha256=ybOdoFE5HIJTkyiYcc73zpyUyUpioVnAca6k0wyJPs4,2572
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=tkaDo232HjuZvJHyua0n6tdHecifUuVzclJAGq7PPYs,11428
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=xUuQjoR0NJhuwG36GuycpKHo9jg783bDSHj9wE4F1Sg,15439
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=khJIvDVk2s332Nd2Be-5dM6-wp5DGff61HCV5lskHmQ,3011
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQDJnaAAxhNFAd7NfQm5SOVEp_nno_krA,6178
 ai_edge_torch/odml_torch/lowerings/_rand.py,sha256=g6SuqDkuC6hD35lyP1-5H7ASDIzPSmKukeNT5naZSv8,4133
@@ -268,8 +269,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 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.6.0.dev20250620.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.6.0.dev20250620.dist-info/METADATA,sha256=W4cZLDBaywmznVjn7haIlLen5cKvXF7VYVIapnx4h0E,2074
-ai_edge_torch_nightly-0.6.0.dev20250620.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.6.0.dev20250620.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.6.0.dev20250620.dist-info/RECORD,,
+ai_edge_torch_nightly-0.6.0.dev20250816.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.6.0.dev20250816.dist-info/METADATA,sha256=x3HCuvXK0BSgJhf34jwigF6tJIjfq55lFwrPZIrPjVA,2074
+ai_edge_torch_nightly-0.6.0.dev20250816.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.6.0.dev20250816.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.6.0.dev20250816.dist-info/RECORD,,