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

ai_edge_torch/_config.py

@@ -22,6 +22,18 @@ import os
 __all__ = ["config"]
 
 
+def _get_bool_env_var(name: str, default: bool) -> bool:
+  var = os.environ.get(name, "false")
+  var = var.lower().strip()
+  if var in ("y", "yes", "t", "true", "on", "1"):
+    return True
+  elif var in ("n", "no", "f", "false", "off", "0"):
+    return False
+  else:
+    logging.warning("Invalid %s value is ignored: %s.", name, var)
+    return default
+
+
 class _Config:
   """ai-edge-torch global configs."""
 
@@ -33,20 +45,25 @@ class _Config:
     To use torch_xla as the lowering backend, set environment variable
     `USE_TORCH_XLA` to "true".
     """
-    var = os.environ.get("USE_TORCH_XLA", "false")
-    var = var.lower().strip()
-    if var in ("y", "yes", "t", "true", "on", "1"):
-      return True
-    elif var in ("n", "no", "f", "false", "off", "0"):
-      return False
-    else:
-      logging.warning("Invalid USE_TORCH_XLA value is ignored: %s.", var)
-      return False
+    return _get_bool_env_var("USE_TORCH_XLA", default=False)
 
   @property
   def in_oss(self) -> bool:
     """True if the code is not running in google internal environment."""
     return True
 
+  @property
+  def enable_group_norm_composite(self) -> bool:
+    """True if lowering group norm in StableHLO composite.
+
+    Currently only supports NHWC group norm generated by
+    OptimizeLayoutTransposesPass.
+    """
+    return _get_bool_env_var("ENABLE_GROUP_NORM_COMPOSITE", default=False)
+
+  @enable_group_norm_composite.setter
+  def enable_group_norm_composite(self, value: bool):
+    os.environ["ENABLE_GROUP_NORM_COMPOSITE"] = "y" if value else "n"
+
 
 config = _Config()

ai_edge_torch/_convert/conversion.py

@@ -14,9 +14,9 @@
 # ==============================================================================
 
 import logging
-import os
 from typing import Any, Literal, Optional, Union
 
+import ai_edge_torch
 from ai_edge_torch import fx_pass_base
 from ai_edge_torch import lowertools
 from ai_edge_torch import model
@@ -26,8 +26,6 @@ from ai_edge_torch.generative import fx_passes as generative_fx_passes
 from ai_edge_torch.quantize import quant_config as qcfg
 import torch
 
-os.environ["EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM"] = "1"
-
 
 def _run_convert_passes(
     exported_program: torch.export.ExportedProgram,
@@ -35,21 +33,27 @@ def _run_convert_passes(
   exported_program = generative_fx_passes.run_generative_passes(
       exported_program
   )
-  exported_program = fx_pass_base.run_passes(
-      exported_program,
-      [
-          fx_passes.BuildInterpolateCompositePass(),
-          fx_passes.CanonicalizePass(),
-          fx_passes.OptimizeLayoutTransposesPass(),
-          fx_passes.CanonicalizePass(),
-          fx_passes.BuildAtenCompositePass(),
-          fx_passes.CanonicalizePass(),
-          fx_passes.RemoveNonUserOutputsPass(),
-          fx_passes.CanonicalizePass(),
-          fx_passes.InjectMlirDebuginfoPass(),
-          fx_passes.CanonicalizePass(),
-      ],
-  )
+
+  passes = [
+      fx_passes.BuildInterpolateCompositePass(),
+      fx_passes.CanonicalizePass(),
+      fx_passes.OptimizeLayoutTransposesPass(),
+      fx_passes.CanonicalizePass(),
+      fx_passes.BuildAtenCompositePass(),
+      fx_passes.CanonicalizePass(),
+      fx_passes.RemoveNonUserOutputsPass(),
+      fx_passes.CanonicalizePass(),
+  ]
+
+  # Debuginfo is not injected automatically by odml_torch. Only inject
+  # debuginfo via fx pass when using torch_xla.
+  if ai_edge_torch.config.use_torch_xla:
+    passes += [
+        fx_passes.InjectMlirDebuginfoPass(),
+        fx_passes.CanonicalizePass(),
+    ]
+
+  exported_program = fx_pass_base.run_passes(exported_program, passes)
   return exported_program
 
 

ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py

@@ -62,6 +62,7 @@ def _wrap_call_function_node_with_debuginfo_writer(node: torch.fx.GraphModule):
 
 
 class InjectMlirDebuginfoPass(fx_pass_base.PassBase):
+  """DEPRECATED: Debuginfo is injected automatically by odml_torch."""
 
   def call(self, graph_module: torch.fx.GraphModule):
     for node in graph_module.graph.nodes:

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py

@@ -17,6 +17,7 @@
 import dataclasses
 import operator
 
+import ai_edge_torch
 from ai_edge_torch import lowertools
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_rewrite
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import utils
@@ -155,6 +156,7 @@ def _qdq_layout_sensitive_inputs_getter(node: Node):
 @layout_sensitive_inputs_getters.register(
     aten._native_batch_norm_legit_no_training
 )
+@layout_sensitive_inputs_getters.register(aten.group_norm)
 @layout_sensitive_inputs_getters.register(aten.native_group_norm)
 def _first_arg_getter(node):
   return [node.args[0]]
@@ -188,6 +190,17 @@ def _aten_norm_checker(node):
   return NHWCable(can_be=len(val[0].shape) == 4, must_be=False)
 
 
+@nhwcable_node_checkers.register(aten.group_norm)
+def _aten_group_norm_checker(node):
+  val = node.meta.get("val")
+  if not hasattr(val, "shape"):
+    return NHWCable(can_be=False, must_be=False)
+
+  can_be = len(val.shape) == 4
+  must_be = can_be and ai_edge_torch.config.enable_group_norm_composite
+  return NHWCable(can_be=can_be, must_be=must_be)
+
+
 @nhwcable_node_checkers.register(aten.native_group_norm)
 def _aten_native_group_norm_checker(node):
   val = node.meta.get("val")

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py

@@ -16,6 +16,7 @@
 
 import operator
 
+import ai_edge_torch
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_mark
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import op_func_registry
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import utils
@@ -23,6 +24,7 @@ import torch
 import torch.utils._pytree as pytree
 
 aten = torch.ops.aten
+StableHLOCompositeBuilder = ai_edge_torch.hlfb.StableHLOCompositeBuilder
 
 __all__ = ["rewrite_nhwc_node", "has_nhwc_rewriter"]
 
@@ -342,6 +344,39 @@ def _aten__native_batch_norm_legit_no_training(node):
   node.target = batch_norm
 
 
+@rewriters.register(aten.group_norm.default)
+def _aten_group_norm(node):
+  def group_norm(input, num_groups: int, weight=None, bias=None, eps=1e-5):
+    is_composite_supported = (
+        ai_edge_torch.config.enable_group_norm_composite
+        and weight is not None
+        and bias is not None
+    )
+
+    builder = None
+    if is_composite_supported:
+      builder = StableHLOCompositeBuilder(
+          name="odml.group_norm",
+          attr={
+              "num_groups": num_groups,
+              "epsilon": eps,
+              "reduction_axes": [3],
+              "channel_axis": 3,
+          },
+      )
+      input, weight, bias = builder.mark_inputs(input, weight, bias)
+
+    input = utils.tensor_to_nchw(input)
+    output = aten.group_norm.default(input, num_groups, weight, bias, eps=eps)
+    output = utils.tensor_to_nhwc(output)
+
+    if builder is not None:
+      output = builder.mark_outputs(output)
+    return output
+
+  node.target = group_norm
+
+
 @rewriters.register(aten.native_group_norm.default)
 def _aten_native_group_norm(node):
 
@@ -354,6 +389,7 @@ def _aten_native_group_norm(node):
       flattened_inner_size: int,
       num_groups: int,
       eps: float,
+      **kwargs,
   ):
     input_reshaped = torch.reshape(
         input,

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -15,14 +15,13 @@
 
 """Example of building a Gemma2 model."""
 
-from typing import List, Optional, Tuple
+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
 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
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
@@ -104,12 +103,17 @@ class Gemma2(nn.Module):
         config.embedding_dim,
         config.final_norm_config,
     )
-    self.mask_cache = attn_utils.build_causal_mask_cache(
-        size=config.kv_cache_max,
-    )
     # Gemma2 has same hyper parameters for each layer except for attention
     # types. Use the first layer.
     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,
+    )
     self.sliding_window_mask_cache = attn_utils.build_sliding_window_mask_cache(
         size=config.kv_cache_max,
         window_size=attn_config.sliding_window_size,
@@ -136,48 +140,29 @@ class Gemma2(nn.Module):
         f"Cannot forward sequence of length {seq_len}, max seq length is only"
         f" {self.config.max_seq_len}"
     )
-
-    # token embeddings of shape (b, t, n_embd)
-    input_embeds = self.tok_embedding(tokens)
-    # RoPE parameters are the same for all blocks. Use the first layer.
-    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
-    )
-    mask = [self.get_attention_mask(
-        self.config.block_config(i).attn_config.attn_type, input_pos
-    ) for i in range(self.config.num_layers)]
-
-    return self._forward_with_embeds(
-        input_embeds, rope, mask, input_pos, kv_cache, export_config
-    )
-
-  def _forward_with_embeds(
-      self,
-      input_embeds: torch.Tensor,
-      rope: Tuple[torch.Tensor, torch.Tensor],
-      mask: List[torch.Tensor],
-      input_pos: torch.Tensor,
-      kv_cache: kv_utils.KVCache,
-      export_config: Optional[model_builder.ExportConfig] = None,
-  ) -> dict[torch.Tensor, kv_utils.KVCache]:
-    """Forwards the model with input embeddings."""
     assert len(self.transformer_blocks) == len(kv_cache.caches), (
         "The number of transformer blocks and the number of KV cache entries"
         " must be the same."
     )
 
-    if self.config.embedding_scale is not None:
-      input_embeds = input_embeds * self.config.embedding_scale
-    x = input_embeds
-    updated_kv_entries = []
+    cos, sin = self.rope_cache
+    cos = cos.index_select(0, input_pos)
+    sin = sin.index_select(0, input_pos)
+
+    # token embeddings of shape (b, t, n_embd)
+    x = self.tok_embedding(tokens)
+    x = x * (self.config.embedding_dim**0.5)
+
+    updated_kv_entires = []
     for i, block in enumerate(self.transformer_blocks):
+      mask = self.get_attention_mask(
+          block.config.attn_config.attn_type, input_pos
+      )
       kv_entry = kv_cache.caches[i] if kv_cache else None
-      x, kv_entry = block(x, rope, mask[i], input_pos, kv_entry)
+      x, kv_entry = block(x, (cos, sin), mask, input_pos, kv_entry)
       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 (
@@ -243,13 +228,11 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
     )
 
   num_layers = 26
-  embedding_dim = 2304
   config = cfg.ModelConfig(
       vocab_size=256000,
       num_layers=num_layers,
       max_seq_len=8192,
-      embedding_dim=embedding_dim,
-      embedding_scale=embedding_dim**0.5,
+      embedding_dim=2304,
       kv_cache_max_len=kv_cache_max_len,
       block_configs=[get_block_config(i) for i in range(num_layers)],
       final_norm_config=norm_config,
@@ -266,7 +249,6 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
   config.num_layers = 2
   config.max_seq_len = 2 * kv_cache_max_len
   config.embedding_dim = 128
-  config.embedding_scale = config.embedding_dim**0.5
   config.block_configs = config.block_configs[: config.num_layers]
   for block_config in config.block_configs:
     block_config.attn_config.num_heads = 4

ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py

@@ -72,14 +72,14 @@ class ToyModelWithKVCache(torch.nn.Module):
     mask = self.mask_cache.index_select(2, input_pos)
     mask = mask[:, :, :, : self.config.max_seq_len]
 
-    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, (cos, sin), mask, input_pos, kv_entry)
       if kv_entry:
-        updated_kv_entries.append(kv_entry)
+        updated_kv_entires.append(kv_entry)
 
-    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entries))
+    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entires))
 
     if export_config is not None:
       if (

ai_edge_torch/generative/layers/attention.py

@@ -26,6 +26,33 @@ import torch
 from torch import nn
 
 
+def _embed_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    n_elem: int,
+    rope: Tuple[torch.Tensor, torch.Tensor],
+) -> Tuple[torch.Tensor, torch.Tensor]:
+  """Embed rotary positional embedding for query and key.
+
+  Args:
+    q (torch.Tensor): query tensor.
+    k (torch.Tensor): key tensor.
+    n_elem (int): number of elements to embed rotarty positional embedding.
+    rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+  """
+  if n_elem > 0:
+    cos, sin = rope
+    q_roped = rotary_pos_emb.apply_rope(
+        q[..., :n_elem], cos.repeat(1, 2), sin.repeat(1, 2)
+    )
+    k_roped = rotary_pos_emb.apply_rope(
+        k[..., :n_elem], cos.repeat(1, 2), sin.repeat(1, 2)
+    )
+    q = torch.cat((q_roped, q[..., n_elem:]), dim=-1)
+    k = torch.cat((k_roped, k[..., n_elem:]), dim=-1)
+  return q, k
+
+
 class TransformerBlock(nn.Module):
 
   def __init__(
@@ -211,8 +238,7 @@ class CausalSelfAttention(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)
+      q, k = _embed_rope(q, k, n_elem, rope)
 
     if kv_cache is not None:
       kv_cache = kv_utils.update(kv_cache, input_pos, k, v)
@@ -348,8 +374,7 @@ 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)
+      q, k = _embed_rope(q, k, n_elem, rope)
 
     if kv_cache is not None:
       kv_cache = kv_utils.update(kv_cache, input_pos, k, v)

ai_edge_torch/generative/layers/normalization.py

@@ -80,6 +80,7 @@ class RMSNorm(torch.nn.Module):
       output = self._norm(x.float()).type_as(x)
       return output * w
 
+
 class GroupNorm(torch.nn.Module):
 
   def __init__(
@@ -115,16 +116,7 @@ class GroupNorm(torch.nn.Module):
     Returns:
       torch.Tensor: output tensor after applying GroupNorm.
     """
-    if self.enable_hlfb:
-      return group_norm_with_hlfb(
-          x,
-          self.weight,
-          self.bias,
-          self.group_num,
-          self.eps,
-      )
-    else:
-      return F.group_norm(x, self.group_num, self.weight, self.bias, self.eps)
+    return F.group_norm(x, self.group_num, self.weight, self.bias, self.eps)
 
 
 class LayerNorm(torch.nn.Module):
@@ -169,46 +161,6 @@ class LayerNorm(torch.nn.Module):
     )
 
 
-def group_norm_with_hlfb(
-    x: torch.Tensor,
-    w: torch.Tensor,
-    b: torch.Tensor,
-    num_groups: int,
-    eps: float,
-):
-  """Group Normalization with high-level function boundary enabled.
-
-  Args:
-    x (torch.Tensor): Input tensor for Group Normalization, with BCHW shape.
-    w (torch.Tensor): The weight tensor for the normalization.
-    b (torch.Tensor): The bias tensor for the normalization.
-    num_groups (int): Number of groups to separate the channels into.
-    eps (float): A small float value to ensure numerical stability.
-
-  Returns:
-    The output tensor of Group Normalization.
-  """
-  x = torch.permute(x, (0, 2, 3, 1))
-
-  builder = StableHLOCompositeBuilder(
-      name="odml.group_norm",
-      attr={
-          "num_groups": num_groups,
-          "epsilon": eps,
-          "reduction_axes": [3],
-          "channel_axis": 3,
-      },
-  )
-  x, w, b = builder.mark_inputs(x, w, b)
-  x = torch.permute(x, (0, 3, 1, 2))
-  y = F.group_norm(x, num_groups, weight=w, bias=b, eps=eps)
-  y = torch.permute(y, (0, 2, 3, 1))
-  y = builder.mark_outputs(y)
-
-  y = torch.permute(y, (0, 3, 1, 2))
-  return y
-
-
 def rms_norm_with_hlfb(
     x: torch.Tensor,
     w: torch.Tensor,

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/utilities/model_builder.py

@@ -24,7 +24,6 @@ from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.model_config as cfg
-import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
 from torch import nn
@@ -86,6 +85,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,
     )
@@ -107,22 +113,16 @@ 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(
+    return self.forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache, export_config
     )
 
-  def _forward_with_embeds(
+  def forward_with_embeds(
       self,
       input_embeds: torch.Tensor,
       rope: Tuple[torch.Tensor, torch.Tensor],
@@ -141,13 +141,13 @@ 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)
       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/lowertools/torch_xla_utils.py

@@ -27,6 +27,9 @@ if "PJRT_DEVICE" not in os.environ:
   # https://github.com/google-ai-edge/ai-edge-torch/issues/326
   os.environ["PJRT_DEVICE"] = "CPU"
 
+os.environ["EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM"] = "1"
+
+
 from ai_edge_torch import model
 from ai_edge_torch._convert import conversion_utils
 from ai_edge_torch._convert import signature as signature_module

ai_edge_torch/odml_torch/export.py

@@ -202,7 +202,7 @@ class MlirLowered:
       target_version = stablehlo.get_minimum_version()
     else:
       target_version = stablehlo.get_version_from_compatibility_requirement(
-          stablehlo.StablehloCompatibilityRequirement.WEEK_4
+          stablehlo.StablehloCompatibilityRequirement.WEEK_12
       )
     module_bytecode = xla_extension.mlir.serialize_portable_artifact(
         self.module_bytecode, target_version
@@ -222,11 +222,6 @@ class MlirLowered:
     # Lazy importing TF when execution is needed.
     return self.tf_function(*args)
 
-  def to_flatbuffer(self):
-    from . import tf_integration
-
-    return tf_integration.mlir_to_flatbuffer(self)
-
 
 # TODO(b/331481564) Make this a ai_edge_torch FX pass.
 def _convert_i64_to_i32(exported_program: torch.export.ExportedProgram):

ai_edge_torch/odml_torch/tf_integration.py

@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""APIs to convert lowered MLIR from PyTorch to TensorFlow and TFLite artifacts."""
+"""APIs to convert lowered MLIR from PyTorch to TensorFlow artifacts."""
 
 import re
-import tempfile
 
 import tensorflow as tf
 import torch
@@ -104,20 +103,26 @@ def _extract_call_args(
 def _wrap_as_tf_func(lowered, tf_state_dict):
   """Build tf.function from lowered and tf_state_dict."""
 
-  def inner(*args):
+  version = 6
+  if hasattr(tfxla, "call_module_maximum_supported_version"):
+    version = tfxla.call_module_maximum_supported_version()
+
+  def tf_func(*args):
     t_outs = [torch_dtype_to_tf(sig.dtype) for sig in lowered.output_signature]
     s_outs = [_get_shape_with_dynamic(sig) for sig in lowered.output_signature]
     call_args = _extract_call_args(lowered, args, tf_state_dict)
     return tfxla.call_module(
         tuple(call_args),
-        version=5,
+        version=version,
         Tout=t_outs,  # dtype information
-        Sout=s_outs,  # Shape information
+        Sout=s_outs,  # shape information
         function_list=[],
-        module=lowered.module_bytecode,
+        module=lowered.module_bytecode_vhlo,
+        has_token_input_output=False,
+        platforms=["CPU"],
     )
 
-  return inner
+  return tf_func
 
 
 def _make_input_signatures(
@@ -149,46 +154,3 @@ def mlir_to_tf_function(lowered: export.MlirLowered):
       _wrap_as_tf_func(lowered, tf_state_dict),
       input_signature=_make_input_signatures(lowered),
   )
-
-
-def mlir_to_flatbuffer(lowered: export.MlirLowered):
-  """Convert the MLIR lowered to a TFLite flatbuffer binary."""
-  tf_state_dict = _build_tf_state_dict(lowered)
-  signature_names = [tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
-  tf_signatures = [_make_input_signatures(lowered)]
-  tf_functions = [_wrap_as_tf_func(lowered, tf_state_dict)]
-
-  tf_module = tf.Module()
-  tf_module.f = []
-
-  for tf_sig, func in zip(tf_signatures, tf_functions):
-    tf_module.f.append(
-        tf.function(
-            func,
-            input_signature=tf_sig,
-        )
-    )
-
-    tf_module._variables = list(tf_state_dict.values())
-
-  tf_concrete_funcs = [
-      func.get_concrete_function(*tf_sig)
-      for func, tf_sig in zip(tf_module.f, tf_signatures)
-  ]
-
-  # We need to temporarily save since TFLite's from_concrete_functions does not
-  # allow providing names for each of the concrete functions.
-  with tempfile.TemporaryDirectory() as temp_dir_path:
-    tf.saved_model.save(
-        tf_module,
-        temp_dir_path,
-        signatures={
-            sig_name: tf_concrete_funcs[idx]
-            for idx, sig_name in enumerate(signature_names)
-        },
-    )
-
-    converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir_path)
-    tflite_model = converter.convert()
-
-  return tflite_model

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241226.dist-info → ai_edge_torch_nightly-0.3.0.dev20250107.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20241226
+Version: 0.3.0.dev20250107
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.3.0.dev20241226.dist-info → ai_edge_torch_nightly-0.3.0.dev20250107.dist-info}/RECORD

@@ -1,11 +1,11 @@
 ai_edge_torch/__init__.py,sha256=rq9ZtMJLG8yYNC4tNE4rpl94UAUClZW7f4GAr6HBVDQ,1208
-ai_edge_torch/_config.py,sha256=QIrerb6uHMahRvMilmhodJ_6jfiRps3qgLOBeidPnS4,1614
+ai_edge_torch/_config.py,sha256=PKtOtBOup-cM0wBdQxby6HzuhLhIC3oq-TBG8FF4znE,2161
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/fx_pass_base.py,sha256=518ziQ0TUxqum2qZXqlD8qr65pHPh8ZNLnwFC6zvK3k,4253
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=E_WDoV5Y1AG1Kq9M0_73bQYoSSnhDCJ7dxLCdKpkJJE,706
+ai_edge_torch/version.py,sha256=X0ZEB5T3xcR8MsIE8VOHDAdHnCZTzJLBQQ9j2xZ4_qA,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=SzbR16V2JEfkCjjPwRVAFUbFnzu-_1iHPKgGT9Yz7gQ,5678
+ai_edge_torch/_convert/conversion.py,sha256=_PoH0E1gbbsWhLGwDRwUtW2G_IgNzNF7pKQbn9ct6-4,5778
 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=rGpBNss3Y9FCCCcdBwDo16KqavJi8N5P0M_6Wnmoz6Y,2320
@@ -13,12 +13,12 @@ ai_edge_torch/_convert/to_channel_last_io.py,sha256=_31phf7TYgZY2ftpNbrdlB1RhDiu
 ai_edge_torch/_convert/fx_passes/__init__.py,sha256=NVe-eGcm7j8jZpP2pcMhC8j5dVjgR1pPzyXhHdvKH4E,1267
 ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py,sha256=doaww8KqrgRTD5LotBVAIRFsEqzPn9R5lcGehBJOczA,9098
 ai_edge_torch/_convert/fx_passes/build_interpolate_composite_pass.py,sha256=qb4JBDi4Xca14JJUIcaaZQIJiyqKyHJF49jsRCIFCVA,4335
-ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=IlZuK42kfVcRqAWZp4j2k_81T2uWo9T2558U_GPJAlU,2327
+ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=NxT-iCOHq3r3jeZ8qhNoPXV5w8l2eRMu4yEcBri3NxY,2398
 ai_edge_torch/_convert/fx_passes/remove_non_user_outputs_pass.py,sha256=f1IUVWyhioOClsMiZzLyynoW2R17U83vA-7Q-3pGPM4,2126
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/__init__.py,sha256=lxnoH-WGLeiQIF8XjMGodjiZEFTxucl7g05N7MR9OPk,796
-ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py,sha256=hDsl9AHzmyuSWsdHOSO114l4nBUgUdAOUWafMTipMgA,7629
+ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py,sha256=S_Bniv6jY16oOoFUzlyECQ0I2HDjG2D1MOI-QYPk3jQ,8061
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_mark.py,sha256=4RyGUwR22bZqkn_TnptenFJodc_Q43f4_SBG7gmTbos,1621
-ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py,sha256=NW37V6QYdPOZOVhqLcmssVk-VAeO4ECk_CrbEBh4B0E,12740
+ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py,sha256=zoAZ2TXKvxUnWnT11U4tx2uF0J5kkNXydgaW7JzfkXI,13811
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/op_func_registry.py,sha256=bsYnudRlXp1PJlu4GF25KSogSkBGQPSaecBrUTONKaw,1031
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/pass_body.py,sha256=t94Am3iPbYQekg-rrtc-jS_aDWtEgAAj7pAKHrG0-9U,10563
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/utils.py,sha256=YLMttMg5PdvXTtQ8lxpKb434UGVvYVALV1-xeuH4UGc,2131
@@ -47,7 +47,7 @@ ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIX
 ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=GhwtQZ1xuMyKJl8qdxU6uKavQnlm5US9xhKJvdmgACc,2309
 ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=hsy4Gd7Inchi0p_Cc5yecH6vr9A7X4MvmQNfTt8N2sQ,2311
 ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=w8oWYibZzvEvCDyp39EYyAWmjgJljhzdYPyFCfAWxZA,3497
-ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=VTM2nO3TqK2d1DyEb2MiHc-Tyw2lMcUXyOhvg0H5ENY,10147
+ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=whQ6DEnmhmj9hd5OyaoEI-FUNJ4m302vY3Swo_IqQcA,9285
 ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=ip-Gmk4CI5f0GWSdAIdrectxQWJ0t328KCsA4nfHuGg,1736
 ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=IoBhEMwH07-tFm5-U6F2hpCsI8xynglhq1x9tIOdaPQ,1322
 ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=tR8RflXocDZqvuStyw9aFlzuiTllEC8rNnjrxms6_Is,5727
@@ -109,7 +109,7 @@ ai_edge_torch/generative/examples/t5/t5_attention.py,sha256=l01oYyJo77INzRwN4xqX
 ai_edge_torch/generative/examples/test_models/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/test_models/convert_toy_model.py,sha256=6-WaNHckq_LlXMVTh8x90MGWeWq2bu_T_XQd3w9FnGg,3261
 ai_edge_torch/generative/examples/test_models/toy_model.py,sha256=4113jZK-Hu3kYop__WTc8Bq-bG6YzQtADbxHtYPEB4w,5036
-ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=C9dzJFK3TybxKpM1vSdLjOKftkJ72DGjr8YR4H7vCe8,4664
+ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=WMl1iuCE8So9FDnxPV0OTMzuPngQUTO61g8rfnBLyB4,4664
 ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=5rgbTIxHoFg8sTnzrGA_ekT-HJEt9p7Dla7cIY874jU,2338
 ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=mhJ18rb9sxrYRzv1YSzhbNs97oUZck99avZDcUO2oV8,2800
@@ -117,14 +117,14 @@ ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=7Bk8z033M-BCXJ299f
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=jrzCB3ZyY_t5jJM1e2Czdt3DjAIL43R0_a-T-I7wOzw,1155
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=hhxSQvkDMv0isZJhmuLiod66ZODaJ8uSPSVTJVHBabQ,1931
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/attention.py,sha256=_OmamS3f0m_JtW73ljwGLwFPeMLL837JCLY-dJ3iRUg,12453
+ai_edge_torch/generative/layers/attention.py,sha256=aOoVM1hY7qjvzVQI1-m26p_f9qoTLzXXIy8dNtU8xC4,13199
 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=DhHIggaOQ2IAY4aRuMAuCLWZv1dBz5PYtmOEjkx9EQY,6291
 ai_edge_torch/generative/layers/model_config.py,sha256=viX51T_naJ9sPpPxPoMnSueBPYE2zxWNOD0xn0f-_bM,7510
-ai_edge_torch/generative/layers/normalization.py,sha256=h2btgRHMMjOcyLm8adEmcT0pG6imq4QcWblKJK5MYXA,7479
-ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=zbFTNgQdOT-tcKK1QaIX6fG-50syYwQX_ZbLhg2C98c,2691
+ai_edge_torch/generative/layers/normalization.py,sha256=MbwH-n80Fob5YvjBzdqDjBizMHLzSJGYRDdbD-rL5C0,6174
+ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=xxWtlVsGGJkEyXC6PwznubyhJnLPEfSpHOORE_hgxss,2670
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=gXxh3papKy4FBpGEX7VyZ7rZ1Js6aHK70Q6DKrVSckY,4154
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=9jKzOfiBQ66bp1ZnVIAoREIifVNFx4aTlQeYMAx2_pA,29062
@@ -149,12 +149,12 @@ ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5l
 ai_edge_torch/generative/utilities/converter.py,sha256=hIwWUWjgPvWLATtsYYG6RWbFQWhOr2RpPlMrd-4Am9U,5959
 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=S08WNqVKCmxd2QjtMlwETd7J97UnlME_bTKdz5LMkGU,6352
+ai_edge_torch/generative/utilities/model_builder.py,sha256=rfD6INxunvDVdiUfTUxD7yy0dRxL74W7kVmZsxUjpOQ,6379
 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
 ai_edge_torch/generative/utilities/transformers_verifier.py,sha256=8sp9m_FMcXn7nqOrochtu2jIANkJKhnhIBUmH0ZTDR4,1549
-ai_edge_torch/generative/utilities/verifier.py,sha256=awO-sQrEpsFxIkZw72ysWZenYEmkLOLOuj62o2c7XeQ,11994
+ai_edge_torch/generative/utilities/verifier.py,sha256=6lnBU9Cy5GanB8JWK3-2_VU3PxqunDWGe-SgSLba5Yw,12065
 ai_edge_torch/hlfb/__init__.py,sha256=sH4um75na-O8tzxN6chFyp6Y4xnexsE7kUQpZySv6dE,735
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=cjTprggj_cuktSCm7-A25e7Shop3k63ylp7sdZmtZ8o,4790
 ai_edge_torch/hlfb/mark_pattern/passes.py,sha256=pjkKcI1nHECPluAt87cFBrt1DP0f3ge7rHq1NhCkBIE,1936
@@ -166,14 +166,14 @@ ai_edge_torch/lowertools/_shim.py,sha256=Mbg16tnCVK0YsHowfbpqpNX1qySuMLvpGI_-I5S
 ai_edge_torch/lowertools/common_utils.py,sha256=Z7p-ivOHtddktpnHrlDm_dSoTxJOdEjFXIGQbzjgwQo,4504
 ai_edge_torch/lowertools/odml_torch_utils.py,sha256=dxg2pBuVhSZeY2Ouc0F6nHiZilWZmpSPA7I8kGqSkVI,8282
 ai_edge_torch/lowertools/test_utils.py,sha256=mdxTlhqHABZEQ_GEmPFCL8LIAWtqRtYZUGdSY1ieZjw,1949
-ai_edge_torch/lowertools/torch_xla_utils.py,sha256=tH5BW8-Up1uy5Iq1LdXiJInXBh4-YqNXJpSwwy3kwSg,9460
+ai_edge_torch/lowertools/torch_xla_utils.py,sha256=1EytIw2R6dthhLhf69wN1L9BaQTeybCD0wga-PhHcMI,9518
 ai_edge_torch/lowertools/translate_recipe.py,sha256=ymkBpFqAUiupRWqrPOWiVphKcXR1K5vHK0RjgBFtxlE,5652
 ai_edge_torch/odml_torch/__init__.py,sha256=S8jOzE9nLof-6es3XDiGJRN-9H_XTxsVm9dE7lD3RWo,812
 ai_edge_torch/odml_torch/_torch_future.py,sha256=AJ0klpsbu2ZBTfiZlqSOoaYzBVITt40a1fYN8xKkEPw,3044
 ai_edge_torch/odml_torch/_torch_library.py,sha256=Lw1gqL2HWNRspdTwNhIkYAHDyafHedHtkXyKKxn-Wss,805
-ai_edge_torch/odml_torch/export.py,sha256=QzOPmcNPB7R-KhhPEP0oGVbDRgGPptIxRSoz3S8py9I,13405
+ai_edge_torch/odml_torch/export.py,sha256=sqIMXmxK_qIuVC-_DNJ6wKlIWiXq4_WOCKbSqMRFudg,13293
 ai_edge_torch/odml_torch/export_utils.py,sha256=QeA37Irlty6AiIBuqmHmJgn3lqahBQ5xsh6IKRoKm1g,4774
-ai_edge_torch/odml_torch/tf_integration.py,sha256=lTFJPPEijLPFmn6qq2jbpVTQOo0YaOTK36kK6rCiyIE,5956
+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
 ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py,sha256=2Y52E_gLeoXpMcPpV-svXsgN3JbEIjnPVjm0xkpTUdQ,3319
@@ -203,8 +203,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.3.0.dev20241226.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241226.dist-info/METADATA,sha256=khQQVRgopWndD2IbqOblhMzGAlOzSS6f0SbpP1oZ5xw,1966
-ai_edge_torch_nightly-0.3.0.dev20241226.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20241226.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241226.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20250107.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20250107.dist-info/METADATA,sha256=p2F-coQaq7CbpMOkQLVnpFB01cCKqftVRGZ4dCVu8Ck,1966
+ai_edge_torch_nightly-0.3.0.dev20250107.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20250107.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20250107.dist-info/RECORD,,