ai-edge-torch-nightly 0.7.0.dev20251007__py3-none-any.whl → 0.8.0.dev20251225__py3-none-any.whl

ai_edge_torch/odml_torch/lowerings/_decomp_registry.py

@@ -14,13 +14,72 @@
 # ==============================================================================
 """Torch export decompositions to run before lowering."""
 
+import functools
 from ai_edge_torch import fx_infra
 import torch
 
 
+# Fork from pytorch/torch/_decomp/decompositions.py
+def upsample_compute_output_size(input_size, output_size, scale_factors):
+  spatial_dimensions = len(input_size) - 2
+  if output_size is not None:
+    torch._check(
+        scale_factors is None,
+        lambda: "Must specify exactly one of output_size and scale_factors",
+    )
+    torch._check(len(output_size) == spatial_dimensions, lambda: "")
+    return output_size
+  if scale_factors is not None:
+    # NB: this isn't necessary lol
+    torch._check(
+        output_size is None,
+        lambda: "Must specify exactly one of output_size and scale_factors",
+    )
+    torch._check(len(scale_factors) == spatial_dimensions, lambda: "")
+    output_size = []
+    for i, s in enumerate(scale_factors):
+      if int(s) == s:
+        output_size.append(input_size[i + 2] * int(s))
+      else:
+        output_size.append(torch.sym_int(input_size[i + 2] * s))
+    return output_size
+  torch._check(
+      False, lambda: "Must specify exactly one of output_size and scale_factors"
+  )
+
+
+# Fork from pytorch/torch/_decomp/decompositions.py
+def _compute_upsample_nearest_indices(input, output_size, scales, exact=False):
+  indices = []
+  num_spatial_dims = len(output_size)
+  offset = 0.5 if exact else 0.0
+
+  for d in range(num_spatial_dims):
+    osize = output_size[d]
+    isize = input.shape[-num_spatial_dims + d]
+    scale = (
+        isize / (isize * scales[d]) if scales[d] is not None else isize / osize
+    )
+
+    output_indices = torch.arange(
+        osize, dtype=torch.float32, device=input.device
+    )
+    input_indices = ((output_indices + offset) * scale).to(torch.int64)
+    for _ in range(num_spatial_dims - 1 - d):
+      input_indices = input_indices.unsqueeze(-1)
+    indices.append(input_indices)
+  return tuple(indices)
+
+
+# Fork from pytorch/torch/_decomp/decompositions.py
+def _upsample_nearest2d_common(input, h_indices, w_indices):
+  result = torch.ops.aten.index(input, (None, None, h_indices, w_indices))
+  result = result.contiguous()
+  return result
+
+
 fx_infra.decomp.update_pre_lower_decomp(
     torch._decomp.get_decompositions([
-        torch.ops.aten.upsample_nearest2d,
         torch.ops.aten._native_batch_norm_legit.no_stats,
         torch.ops.aten._native_batch_norm_legit_functional,
         torch.ops.aten._adaptive_avg_pool2d,
@@ -35,11 +94,44 @@ fx_infra.decomp.update_pre_lower_decomp(
         torch.ops.aten.replication_pad2d,
         torch.ops.aten.replication_pad3d,
         torch.ops.aten.upsample_bilinear2d.vec,
-        torch.ops.aten.upsample_nearest2d.vec,
         torch.ops.aten.addmm,
     ])
 )
 
+
+@functools.partial(
+    fx_infra.decomp.add_pre_lower_decomp,
+    torch.ops.aten.upsample_nearest2d.default,
+)
+@fx_infra.annotate_force_decomp
+def upsample_nearest2d(input, output_size, scales_h=None, scales_w=None):
+  h_indices, w_indices = _compute_upsample_nearest_indices(
+      input, output_size, (scales_h, scales_w)
+  )
+  return _upsample_nearest2d_common(input, h_indices, w_indices)
+
+
+def get_scale_value(scales, idx):
+  if scales is None:
+    return None
+  return scales[idx]
+
+
+@functools.partial(
+    fx_infra.decomp.add_pre_lower_decomp,
+    torch.ops.aten.upsample_nearest2d.vec,
+)
+@fx_infra.annotate_force_decomp
+def upsample_nearest2d_vec(input, output_size, scale_factors):
+  osize = upsample_compute_output_size(input.size(), output_size, scale_factors)
+  scale_h = get_scale_value(scale_factors, 0)
+  scale_w = get_scale_value(scale_factors, 1)
+
+  return torch.ops.aten.upsample_nearest2d.default(
+      input, osize, scale_h, scale_w
+  )
+
+
 fx_infra.decomp.remove_pre_lower_decomp(torch.ops.aten.roll)
 
 # Torch's default einsum impl/decompositions is less efficient and

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -21,6 +21,7 @@ from ai_edge_torch.odml_torch.lowerings import registry
 import jax
 import jax.numpy as jnp
 from jax._src.lib.mlir import ir
+import numpy as np
 import torch
 import torch_xla2.ops.jaten  # Import to load torch_xla2 ops
 import torch_xla2.ops.ops_registry  # Import to load torch_xla2 ops
@@ -71,8 +72,6 @@ lower_by_torch_xla2(torch.ops.aten._cdist_forward)
 lower_by_torch_xla2(torch.ops.aten._local_scalar_dense)
 lower_by_torch_xla2(torch.ops.aten._local_scalar_dense)
 lower_by_torch_xla2(torch.ops.aten._log_softmax)
-lower_by_torch_xla2(torch.ops.aten._native_batch_norm_legit)
-lower_by_torch_xla2(torch.ops.aten._native_batch_norm_legit_no_training)
 lower_by_torch_xla2(torch.ops.aten._pdist_forward)
 lower_by_torch_xla2(torch.ops.aten._softmax)
 lower_by_torch_xla2(torch.ops.aten._unsafe_index)
@@ -158,10 +157,8 @@ 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.max)
-lower_by_torch_xla2(torch.ops.aten.max_pool2d_with_indices)
 lower_by_torch_xla2(torch.ops.aten.max_pool2d_with_indices_backward)
 lower_by_torch_xla2(torch.ops.aten.max_pool2d_with_indices_backward)
-lower_by_torch_xla2(torch.ops.aten.max_pool3d_with_indices)
 lower_by_torch_xla2(torch.ops.aten.maximum)
 lower_by_torch_xla2(torch.ops.aten.mean)
 lower_by_torch_xla2(torch.ops.aten.min)
@@ -175,7 +172,6 @@ lower_by_torch_xla2(torch.ops.aten.nonzero)
 lower_by_torch_xla2(torch.ops.aten.outer)
 lower_by_torch_xla2(torch.ops.aten.permute)
 lower_by_torch_xla2(torch.ops.aten.permute_copy)
-lower_by_torch_xla2(torch.ops.aten.pixel_shuffle)
 lower_by_torch_xla2(torch.ops.aten.pow)
 lower_by_torch_xla2(torch.ops.aten.prod)
 lower_by_torch_xla2(torch.ops.aten.reciprocal)
@@ -240,6 +236,249 @@ lower_by_torch_xla2(torch.ops.prims.broadcast_in_dim)
 lower_by_torch_xla2(torch.ops.prims.var)
 
 
+def _ceil_mode_padding(
+    padding: list[int],
+    input_shape: list[int],
+    kernel_size: list[int],
+    stride: list[int],
+    dilation: list[int],
+    ceil_mode: bool,
+):
+  """Creates low and high padding specification for ceil mode.
+
+  This is for the given padding (which is symmetric). Additional high padding
+  could be required when ceil mode is set.
+  """
+  ceil_mode_padding = []
+  for i in range(len(padding)):
+    left_padding = padding[i]
+    right_padding = left_padding
+
+    input_size = input_shape[2 + i]
+    output_size_rem = (
+        input_size + 2 * left_padding - (kernel_size[i] - 1) * dilation[i] - 1
+    ) % stride[i]
+    if ceil_mode and output_size_rem != 0:
+      extra_padding = stride[i] - output_size_rem
+      new_output_size = (
+          input_size
+          + left_padding
+          + right_padding
+          + extra_padding
+          - (kernel_size[i] - 1) * dilation[i]
+          - 1
+          + stride[i]
+          - 1
+      ) // stride[i] + 1
+      # Ensure that the last pooling starts inside the image.
+      size_to_compare = input_size + left_padding
+
+      if (new_output_size - 1) * stride[i] < size_to_compare:
+        right_padding += extra_padding
+
+    ceil_mode_padding.append((left_padding, right_padding))
+  return ceil_mode_padding
+
+
+def max_pool(
+    inputs,
+    kernel_size,
+    strides=None,
+    padding=0,
+    dilation=1,
+    ceil_mode=False,
+    with_index=False,
+):
+  num_spatial_dims = len(kernel_size)
+  num_batch_dims = inputs.ndim - num_spatial_dims - 1
+  kernel_size_tup = tuple(kernel_size)
+  # Default stride is kernel_size
+  strides_tup = tuple(strides) if strides else kernel_size_tup
+  if isinstance(padding, int):
+    padding_list = [padding for _ in range(num_spatial_dims)]
+  elif not padding:  # padding can be [], meaning all zeros.
+    padding_list = [0 for _ in range(num_spatial_dims)]
+  else:
+    padding_list = padding
+
+  if isinstance(dilation, int):
+    dilation_tup = tuple(dilation for _ in range(num_spatial_dims))
+  elif not dilation:
+    dilation_tup = tuple(1 for _ in range(num_spatial_dims))
+  elif isinstance(dilation, list):
+    dilation_tup = tuple(dilation)
+  else:
+    dilation_tup = dilation
+
+  input_shape_for_ceil = inputs.shape
+  if num_batch_dims == 0:
+    input_shape_for_ceil = [1, *input_shape_for_ceil]
+  padding_pairs = _ceil_mode_padding(
+      padding_list,
+      input_shape_for_ceil,
+      kernel_size_tup,
+      strides_tup,
+      dilation_tup,
+      ceil_mode,
+  )
+
+  assert len(kernel_size_tup) == len(
+      strides_tup
+  ), f"len({kernel_size_tup=}) must equal len({strides_tup=})"
+  assert len(kernel_size_tup) == len(
+      dilation_tup
+  ), f"len({kernel_size_tup=}) must equal len({dilation_tup=})"
+
+  is_single_input = False
+  if num_batch_dims == 0:
+    inputs = inputs[None]
+    is_single_input = True
+
+  reduce_window_strides = (1,) * (inputs.ndim - num_spatial_dims) + strides_tup
+  reduce_window_dims = (1,) * (inputs.ndim - num_spatial_dims) + kernel_size_tup
+  reduce_window_dilation = (1,) * (
+      inputs.ndim - num_spatial_dims
+  ) + dilation_tup
+
+  assert inputs.ndim == len(
+      reduce_window_dims
+  ), f"len({inputs.shape}) != len({reduce_window_dims})"
+  if not isinstance(padding_pairs, str):
+    padding_pairs_tup = tuple(padding_pairs)
+    assert all(
+        [len(x) == 2 for x in padding_pairs_tup]
+    ), f"each entry in padding {padding_pairs_tup} must be length 2"
+    padding_lax = ((0, 0),) * (
+        inputs.ndim - len(padding_pairs_tup)
+    ) + padding_pairs_tup
+  else:
+    padding_lax = padding_pairs
+
+  indices = jnp.arange(
+      np.prod(inputs.shape[-num_spatial_dims:]), dtype=jnp.int64
+  )
+  indices = indices.reshape(inputs.shape[-num_spatial_dims:])
+  indices_shape = (1,) * (inputs.ndim - indices.ndim) + indices.shape
+  indices = jnp.broadcast_to(indices.reshape(indices_shape), inputs.shape)
+
+  return_dtype = inputs.dtype
+  if jnp.issubdtype(inputs.dtype, jnp.integer):
+    init_val = jnp.int32(jnp.iinfo(jnp.int32).min)
+    inputs = inputs.astype(jnp.int32)
+  else:
+    init_val = jnp.float32(-jnp.inf)
+    inputs = inputs.astype(jnp.float32)
+
+  if not with_index:
+    y = jax.lax.reduce_window(
+        inputs,
+        init_val,
+        jax.lax.max,
+        reduce_window_dims,
+        reduce_window_strides,
+        padding_lax,
+        window_dilation=reduce_window_dilation,
+    )
+    if is_single_input:
+      y = jnp.squeeze(y, axis=0)
+    return y.astype(return_dtype)
+  else:
+
+    def reduce_fn(a, b):
+      ai, av = a
+      bi, bv = b
+      which = av >= bv
+      return jnp.where(which, ai, bi), jnp.where(which, av, bv)
+
+    indices, y = jax.lax.reduce_window(
+        (indices, inputs),
+        (jnp.int64(0), init_val),
+        reduce_fn,
+        reduce_window_dims,
+        reduce_window_strides,
+        padding_lax,
+        window_dilation=reduce_window_dilation,
+    )
+    if is_single_input:
+      indices = jnp.squeeze(indices, axis=0)
+      y = jnp.squeeze(y, axis=0)
+    y = y.astype(return_dtype)
+    return y, indices
+
+
+@lower_by_jax(torch.ops.aten.max_pool2d_with_indices)
+def _aten_max_pool2d_with_indices(
+    self, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False
+):
+  stride = stride if stride is not None else []
+  y = max_pool(
+      self,
+      kernel_size,
+      strides=stride,
+      padding=padding,
+      dilation=dilation,
+      ceil_mode=ceil_mode,
+      with_index=False,
+  )
+  # TFLite's reduce_window kernel doesn't support multiple inputs/outputs,
+  # so we emit reduce_window with a single output and return dummy indices.
+  return y, jnp.zeros_like(y, dtype=jnp.int64)
+
+
+@lower_by_jax(torch.ops.aten.max_pool3d_with_indices.default)
+def _aten_max_pool3d_with_indices(
+    self, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False
+):
+  stride = stride if stride is not None else []
+  y = max_pool(
+      self,
+      kernel_size,
+      strides=stride,
+      padding=padding,
+      dilation=dilation,
+      ceil_mode=ceil_mode,
+      with_index=False,
+  )
+  # TFLite's reduce_window kernel doesn't support multiple inputs/outputs,
+  # so we emit reduce_window with a single output and return dummy indices.
+  return y, jnp.zeros_like(y, dtype=jnp.int64)
+
+
+@lower_by_jax(torch.ops.aten.pixel_shuffle)
+def _aten_pixel_shuffle(x, upscale_factor):
+  """PixelShuffle implementation in JAX lowering.
+
+  Args:
+    x: Input tensor. Typically a feature map.
+    upscale_factor: Integer by which to upscale the spatial dimensions.
+
+  Returns:
+    Tensor after PixelShuffle operation.
+  """
+
+  batch_size, channels, height, width = x.shape
+
+  if channels % (upscale_factor**2) != 0:
+    raise ValueError(
+        "Number of channels must be divisible by the square of the upscale"
+        " factor."
+    )
+
+  new_channels = channels // (upscale_factor**2)
+  new_height = height * upscale_factor
+  new_width = width * upscale_factor
+
+  x = x.reshape(
+      batch_size, new_channels, upscale_factor, upscale_factor, height, width
+  )
+  x = jnp.transpose(
+      x, (0, 1, 4, 2, 5, 3)
+  )  # Move channels to spatial dimensions
+  x = x.reshape(batch_size, new_channels, new_height, new_width)
+
+  return x
+
+
 @lower_by_jax(torch.ops.aten.unbind)
 def _aten_copy(self, *args, **kwargs):
   return _TORCH_XLA2_IMPLS[torch.ops.aten.unbind_copy](self, *args, **kwargs)
@@ -250,6 +489,17 @@ def _aten_copy(self, src, **kwargs):
   return _TORCH_XLA2_IMPLS[torch.ops.aten.copy](self, src)
 
 
+@lower_by_jax(torch.ops.aten.elu.default)
+def _aten_elu(self, alpha=1.0, scale=1.0, input_scale=1.0):
+  pos_coef = scale
+  neg_coef = alpha * scale
+  neg_input_coef = input_scale
+
+  pos_branch = self * pos_coef
+  neg_branch = jnp.expm1(self * neg_input_coef) * neg_coef
+  return jnp.where(self >= 0, pos_branch, neg_branch)
+
+
 @registry.lower(torch.ops.aten.add.Scalar)
 def _aten_add_scalar(lctx: LoweringContext, self, other):
   _log_usage(torch.ops.aten.add.Scalar)

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.7.0.dev20251007"
+__version__ = "0.8.0.dev20251225"

{ai_edge_torch_nightly-0.7.0.dev20251007.dist-info → ai_edge_torch_nightly-0.8.0.dev20251225.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: ai-edge-torch-nightly
-Version: 0.7.0.dev20251007
+Version: 0.8.0.dev20251225
 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
@@ -13,6 +13,8 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Scientific/Engineering :: Mathematics
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
@@ -37,7 +39,17 @@ Requires-Dist: ai-edge-quantizer-nightly
 Requires-Dist: jax
 Requires-Dist: torch-xla2[odml]>=0.0.1.dev20241201
 Provides-Extra: torch-xla
-Requires-Dist: torch-xla>=2.4.0; extra == "torch-xla"
+Requires-Dist: torch_xla>=2.4.0; extra == "torch-xla"
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license-file
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 Library that supports converting PyTorch models into a .tflite format, which can
 then be run with TensorFlow Lite and MediaPipe.  This enables applications for

{ai_edge_torch_nightly-0.7.0.dev20251007.dist-info → ai_edge_torch_nightly-0.8.0.dev20251225.dist-info}/RECORD

@@ -2,9 +2,9 @@ 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=tU7CcypH5Q4v6aP0MF8Y1ZbZzaP8ChOX17X8eaE1v-I,806
+ai_edge_torch/version.py,sha256=EqYE0SbfgYjfk194m19-ExhokdnUqLLGwlHCgT7w_rM,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.py,sha256=JqGZZGbpTmYiT-ta07IQbJ9-gFm-3Vip2aSzW9ulIng,6117
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
 ai_edge_torch/_convert/converter.py,sha256=6MLKELzAwFoiXv-b7KRYi7gc7Z57XOeowcz9ArIl9TM,12100
 ai_edge_torch/_convert/signature.py,sha256=-YKJdLk-eNEHfhdPCtcQVtZf915SoVePEFxKXPPf16c,2572
@@ -41,9 +41,9 @@ ai_edge_torch/examples/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzN
 ai_edge_torch/examples/selfie_segmentation/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
 ai_edge_torch/examples/selfie_segmentation/model.py,sha256=5otCH1MzNgSP0fikYq53hgiO1F0ZN1SCVzOIo7cVAcA,17136
 ai_edge_torch/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/fx_infra/__init__.py,sha256=APjkSqEfwDxcnI8k53rGi3Ef-G2L-M8fdaPGpxXtuiI,1347
+ai_edge_torch/fx_infra/__init__.py,sha256=bseeaX7oDyyvl_oAIT2MfDYZWITmNpn660AlBOUcyUc,1418
 ai_edge_torch/fx_infra/_canonicalize_pass.py,sha256=GDRoDdPVQw--QQFTT5J_C3TVuphL31m6K6F1-67SE4s,1097
-ai_edge_torch/fx_infra/_safe_run_decompositions.py,sha256=ZbWheeZ8ydsxCk2aVGUgUynrkEkBOMjBCzPhS5uq4sU,2595
+ai_edge_torch/fx_infra/_safe_run_decompositions.py,sha256=3rGVQj7OgrIBd--olhnGpSyF2kvpjXMfADss6VPsvxQ,4167
 ai_edge_torch/fx_infra/decomp.py,sha256=S58SCgwMHYVFl_hJwlJxvu2wcI-AGNn82gel3qmTPrU,2500
 ai_edge_torch/fx_infra/graph_utils.py,sha256=nqGe-xIJ77RamSUh0UYyI2XHOsZqFDWax-vpRAtVR_E,2796
 ai_edge_torch/fx_infra/pass_base.py,sha256=Ic2AlhSoRFscz6l7gJKvWVNMDLQFfAw5kRf84-ZR9qM,2904
@@ -141,7 +141,7 @@ ai_edge_torch/generative/examples/smollm/verify_util.py,sha256=tHwmo8E474gLKAu4I
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=kDWG6MlIGa89zC5KSRcJlw2c4ITuw8KcchtfmF55f4g,3545
 ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=lSCRZsoLjH_kqasRMwCy5IogkhyJdwcHKsPEfyxsXCQ,6112
-ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=_yk6wVoZm1_FRMFJF5URaPZNNdmMR89fwmKz81BEyao,5601
+ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=Wba7-rEibvc7ii0uQhO87meVqdQC0_3mtYj5PRXXXbE,5584
 ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=afyHXc86h-ij5zTULmZnM1h313N9VWCyIVriH6pqeSo,16368
 ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=ylqXOZhYc6XFCaNBKQw0jAnYrCtRFFQKzQzEsFIntvo,34890
 ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=CAPsW84A8f00nS6fLFeh_XUjCPsDCA5UxHOUsMrLfSU,3450
@@ -151,7 +151,7 @@ ai_edge_torch/generative/examples/stable_diffusion/util.py,sha256=XIXIB0vCvQKOGy
 ai_edge_torch/generative/examples/stable_diffusion/samplers/__init__.py,sha256=uQWKzCD_49ackNFrt50H04dkDXxfAwUCtMWWQre5SVE,830
 ai_edge_torch/generative/examples/stable_diffusion/samplers/k_euler.py,sha256=wBBNM24waZ57M1rXonwesfUkKe9DqpqO3eW6BfZkrD0,2323
 ai_edge_torch/generative/examples/stable_diffusion/samplers/k_euler_ancestral.py,sha256=c89ldwtuQ2_yspGrGa7oh7fsvTt6A86Whxa6fBK9YOQ,2526
-ai_edge_torch/generative/examples/stable_diffusion/samplers/k_lms.py,sha256=ZE6HyOoBJrmTh54KVFf7DjNBnBS0pT4cgviYaq8HGMU,2801
+ai_edge_torch/generative/examples/stable_diffusion/samplers/k_lms.py,sha256=NDiJPreRFHW1l50Lmb1FHNQcs3wslFAhS4FEFTQKjMU,2805
 ai_edge_torch/generative/examples/stable_diffusion/samplers/sampler.py,sha256=RxR5rw0wFFm_5CfAY-3-EIz83vhM9EKye8Bb5zBb0Ok,1341
 ai_edge_torch/generative/examples/t5/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/t5/convert_to_tflite.py,sha256=HHtZTtUh3QgE4F74-ru_8n1pt6cqfbObw12xoaMJ7NQ,4596
@@ -169,49 +169,49 @@ ai_edge_torch/generative/examples/tiny_llama/verify_util.py,sha256=z6vPBXDWAL6gN
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=PFSMsA1vfBfrV9ssBCkYJNl8Hx_bLdWjN01iyjPM5jE,1094
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=myGjal5A8yIBoqgArd2k40rZmCgD1Ya369KR7182bhI,2129
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/attention.py,sha256=RaXENRRQo1MsLdt3U8h3kYTCmd6imHQ-aCXtmPXCh_o,13911
-ai_edge_torch/generative/layers/attention_test.py,sha256=9v8v96TLyFPdqxEylU1JOAeRFAp2s0YoDHZN83SFJJA,4764
-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=UiLOyvwd-bc0n5XcbVxi6JCn_qiKSC6zrDKSZT_TSDA,5030
+ai_edge_torch/generative/layers/attention.py,sha256=ZjU3vX-7gOq1KQb3xSZ1NT3xryOTXbYb_vkx_DlcizA,14524
+ai_edge_torch/generative/layers/attention_test.py,sha256=ON9jQRY1r2kFpVq-Qkg6b13Ob95fd4PqHo1hic3RbOQ,5057
+ai_edge_torch/generative/layers/attention_utils.py,sha256=3Ox1XjW_vaqz1-RuVG9RbzRKUqCberFW8P2BQcoNm7A,9659
+ai_edge_torch/generative/layers/attention_utils_test.py,sha256=IHIk39wqaPvxmkZtW27VD3_4xUpyFow_7mScf8OWdqU,3292
+ai_edge_torch/generative/layers/builder.py,sha256=5QL59CbOOW_mk3mlPdcdirGcAxdLee5atbZlnu5Z3ts,5079
 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
 ai_edge_torch/generative/layers/kv_cache.py,sha256=A0IFXZ1HD2ZHOWRLfsDO4almgE0KQfjyBOdBFZIGnAs,10893
 ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
-ai_edge_torch/generative/layers/model_config.py,sha256=HP-vu1UmAiTmdLlTyZGDUF3le0gji8a61mLCy966NZw,10261
-ai_edge_torch/generative/layers/normalization.py,sha256=WAhPcLbcC3SCEa6oIgdsojvN306_S8d90WyMQ7ZVP6I,7269
+ai_edge_torch/generative/layers/model_config.py,sha256=y6nyydMzm7JUV7AKttcFy3tvti-nE6tRXoVbBB9dyiM,10438
+ai_edge_torch/generative/layers/normalization.py,sha256=syasVh3dRDVp2Nwhl0x7zucL-chTnCqWgeV1mb87DFY,7435
 ai_edge_torch/generative/layers/normalization_test.py,sha256=zwurZly-TgFxdgVVdpzu9vCpcLbd5RYt_gKg9Lfg1jI,2248
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIrILJ7blceAcxrNqX1ZCN0ECKG1gz-bV8,2655
-ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=2_AgwENsaOgaxgiSqgoj0V0JzQ09dFtP_nBhX-lJK2g,5648
+ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=1zhOsJpI4CTn78weOs0uRwkRxYu6wGfBvYVFpGFl0qQ,6681
 ai_edge_torch/generative/layers/scaled_dot_product_attention_test.py,sha256=c6JBMQsq9XeMmR1XvGEIidNsoh-YIvichXo2LwVHgr4,3301
-ai_edge_torch/generative/layers/sdpa_with_kv_update.py,sha256=iw7D_46CFe9iRvU0UumbkIoqWQEhDroxm9ABcK-CLlM,3600
+ai_edge_torch/generative/layers/sdpa_with_kv_update.py,sha256=fK_h9M-03ai5dV8ZyQzvB0y84IKlNg9h-4bt9F6bU0g,3833
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=ZteHZXK6HKyxYji49DQ46sA9aIy7U3Jnz0HZp6hfevY,28996
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
 ai_edge_torch/generative/layers/unet/model_config.py,sha256=pPDwLawc23pfMaPVyMJlYmxVVusjMvx-l8wBwOYOH-c,9692
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/quantize/example.py,sha256=1lfVNUd2cEyRUnoZ7BLbRJ9IN-FTKiWBtZNPFUzAiWE,1747
-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=mmPsgQ2vloMWflqJ6ALmD1lANacTkWotEhurb5foK30,2771
-ai_edge_torch/generative/quantize/supported_schemes.py,sha256=TwR2FpQuBEORy6FshEyHNBMKARWlA2MVtTfX9tXV5aE,1488
+ai_edge_torch/generative/quantize/example.py,sha256=rRm5noeqnCJyaDYG6e1micwaYdA3ooXK0MBYl98neQQ,1742
+ai_edge_torch/generative/quantize/quant_attrs.py,sha256=THKS82G9sR25biEk8AVRFWwB1cjVXpbjZD5IBfNOOMw,2344
+ai_edge_torch/generative/quantize/quant_recipe.py,sha256=kL3UPNscWKzugiyunMn5XK6Lzpw-yXTdtM_NmCNbFUM,5330
+ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=bFvY-OnbH6fe5xhS6dfxkAl7NTGzOVtGJ_8417GwGgI,2478
+ai_edge_torch/generative/quantize/quant_recipes.py,sha256=nvADL7wObxJLp651HI1xv5_HeQZ2RaY2swoRC1EOK7k,2783
+ai_edge_torch/generative/quantize/supported_schemes.py,sha256=lEcqwxOIfxhTYCQUelODLky57933r372Uh7sHYNV0Ok,1712
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/test/test_custom_dus.py,sha256=MjIhTvkTko872M35XMciobvICcDWTcIDJ3rociko-wM,3267
-ai_edge_torch/generative/test/test_kv_cache.py,sha256=zVocQZTORIROwSYVMmy5KGPDl6kTuEGko7L1PTkua8g,5568
+ai_edge_torch/generative/test/test_kv_cache.py,sha256=kdoBpyzVOjQsZkEPnbDOh33gRE0RR8tq1XCgJkSYwU0,6157
 ai_edge_torch/generative/test/test_loader.py,sha256=9mQUeeZKOVApOWSWl2cN9c10axZjMKM1-0Zd823CCS4,3449
 ai_edge_torch/generative/test/test_lora.py,sha256=sKnBixmGIXHUtOwh3SA4MIFenPbjK2n-Xknwic_KMDQ,5046
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=T35zdzag2-nmy4qc6AifAjbDXAHU2vyLTE1QCabYBzk,6298
 ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=NkEwrjO8vIcde3XwanpFBhNIw1GSOyJFKNjlvSJmVMY,13271
-ai_edge_torch/generative/test/test_quantize.py,sha256=kKJ01wscTC2t_Ylr7huO5gNKES01gm3dT1gx52z15PA,7356
+ai_edge_torch/generative/test/test_quantize.py,sha256=Sh82oxDIc_vw-rCd991qMNQxzM3OkI3oiPj2g4hD5MQ,7088
 ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv9og0DDLHU,2656
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/converter.py,sha256=Bt-48O1wf-7YcGVof53eVKI7wJwNvrc1Bv5zE3JuFdk,20093
+ai_edge_torch/generative/utilities/converter.py,sha256=B2DNg5LcDoniHnh22KFDDSHqnmi8w87-dstKCae3CeM,23322
 ai_edge_torch/generative/utilities/export_config.py,sha256=5B15nYyqf96kjjYlHfPctUfsIdsBsh1f8rxKitJpwKQ,2384
-ai_edge_torch/generative/utilities/litertlm_builder.py,sha256=0cNuaqhc7cQcAa4NRalUXyoPQUQC9O3-aHAJEDV1Mps,4265
-ai_edge_torch/generative/utilities/loader.py,sha256=drgKBmNibuc3PCdc0kU0pVcp2Nt1_mjLYh67RyXOn7U,15952
+ai_edge_torch/generative/utilities/litertlm_builder.py,sha256=Jy6R4S18JCIk77ZwrfzeOCSaGozvRQfXSobGkwOEOKs,6447
+ai_edge_torch/generative/utilities/loader.py,sha256=QQeEu0cTC7gWnB7RkHonjWLdVGjMbDHd1lfYO_TcJyU,16047
 ai_edge_torch/generative/utilities/model_builder.py,sha256=xBvcTxihB9TN88UtQiXA9sAITQgf-pA77R-VZlLgUeU,6950
 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
@@ -232,11 +232,11 @@ 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=Oavs0ENKVnIeB-WidXvokTPqNlFfOuP0GMV0d5RK2Rg,6285
+ai_edge_torch/lowertools/translate_recipe.py,sha256=WEv8Kr0RZw0pBc5PoNUEPnQGBW_JUmuVWQ7Hf4bw0uY,6545
 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=WPsATMFsc62aETExBsb1v6Ec1fL3r0CEvxwE7r2C_CA,14931
+ai_edge_torch/odml_torch/export.py,sha256=PdNdFtEZcFpy4G8ikhElnil8yoYZUWHQ_180E1BGFV4,15347
 ai_edge_torch/odml_torch/export_utils.py,sha256=Eax4QUefIzpmVuQxo1y9FqJ6g0qXjg4C0IVZ5uYPscs,4899
 ai_edge_torch/odml_torch/optimization_barrier.py,sha256=2lmSiu5iXWLFWpupZHvsVeNYNzG5AVGSK3K_CNhS5Sk,2290
 ai_edge_torch/odml_torch/tf_integration.py,sha256=NN29WeXmHZ0S1RPDFHUnBi2DEjMvAtwczStPYIsQ1w8,4849
@@ -246,15 +246,21 @@ ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py,sha256=h6DQkYV
 ai_edge_torch/odml_torch/debuginfo/__init__.py,sha256=3A_lMyj-B-DOhLJG6WmjKvZK5te2rXje8FrfqOhZsN0,959
 ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=6Ns2rlfOilLJEk5cUxlkRwm2uxOgEF2-0S2DMcOqr6A,3319
 ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py,sha256=IvOBQyROI9WHS3umHRxsDW-1YElU9BPWzKtJA2eKWOI,1739
+ai_edge_torch/odml_torch/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
+ai_edge_torch/odml_torch/experimental/torch_tfl/__init__.py,sha256=MN-8wS8vZmSWr9TOoZivqzZGDPoNpy15dByM3arNnM8,968
+ai_edge_torch/odml_torch/experimental/torch_tfl/_decomps.py,sha256=KEPc8wZeAukFhlfKUA2TSmRdMozqG16msnEw7JsGrTg,13256
+ai_edge_torch/odml_torch/experimental/torch_tfl/_lowerings.py,sha256=9coRYoVQFtGhcWuMgZf7lRGIMpRqupbyueAqShx6ieQ,21065
+ai_edge_torch/odml_torch/experimental/torch_tfl/_ops.py,sha256=gRoAf6QjER-fb4oVB05vVGDV__HiXeJgEo1WXX7TnkM,10778
+ai_edge_torch/odml_torch/experimental/torch_tfl/torch_library_utils.py,sha256=9Oq3Dsc5YAkpBgPNFWTquT4TVQnOG8LOlwIabDIc10k,1236
 ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNitEeg-IoBUGNfUxsDSA,798
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=LqwZ1vCJTSOzgzvH8LUAN-sAkF-l_pGj1AMEIzAqHCA,6638
 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=sC4N5-7RS9yKecs97kM9J56enGvsZj1CJo7y79cuzRg,12784
+ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=HOTYfQWin8tqi1yakIyardxhRViZ6rhLV6ZomMSS7zA,17554
 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=6KZc0-SM0LupOnDAgnUjj47aHk7YJ0cC1ulAeWvV5_w,19289
+ai_edge_torch/odml_torch/lowerings/_decomp_registry.py,sha256=LmSj5RsZBi00EE7KfF3dI2U0e60LMHA6mDKc-TC2U0U,5486
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=YmyM-5HJeeYaIhmKTOnCjfX3_A1PPh1gPGUi1d8EBs8,26454
 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
@@ -270,8 +276,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.7.0.dev20251007.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.7.0.dev20251007.dist-info/METADATA,sha256=CmKVd9dJePv9sOzm2nGzlA67l5q8rQY0aHvH8feLfyI,2074
-ai_edge_torch_nightly-0.7.0.dev20251007.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.7.0.dev20251007.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.7.0.dev20251007.dist-info/RECORD,,
+ai_edge_torch_nightly-0.8.0.dev20251225.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.8.0.dev20251225.dist-info/METADATA,sha256=aFJe9FPOOmsDU7-axoCav2W0-n4fRxf0lApj08AZb0s,2399
+ai_edge_torch_nightly-0.8.0.dev20251225.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
+ai_edge_torch_nightly-0.8.0.dev20251225.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.8.0.dev20251225.dist-info/RECORD,,

{ai_edge_torch_nightly-0.7.0.dev20251007.dist-info → ai_edge_torch_nightly-0.8.0.dev20251225.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.45.1)
+Generator: setuptools (79.0.1)
 Root-Is-Purelib: true
 Tag: py3-none-any