ai-edge-torch-nightly 0.8.0.dev20251206__py3-none-any.whl → 0.8.0.dev20260105__py3-none-any.whl

ai_edge_torch/fx_infra/__init__.py

@@ -30,3 +30,4 @@ run_passes = pass_base.run_passes
 
 CanonicalizePass = _canonicalize_pass.CanonicalizePass
 safe_run_decompositions = _safe_run_decompositions.safe_run_decompositions
+annotate_force_decomp = _safe_run_decompositions.annotate_force_decomp

ai_edge_torch/fx_infra/_safe_run_decompositions.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """ExportedProgram.run_decompositions wrapper to handle unexpected export behavior."""
 import operator
+from typing import Any, Callable
 import torch
 
 
@@ -59,6 +60,15 @@ def _require_decomp(
   return False
 
 
+_FORCE_DECOMP_ATTR = "_ai_edge_torch_force_decomp"
+
+
+def annotate_force_decomp(decomp: Callable[..., Any]):
+  """Annotates a decomp to force it to be run (at least shallowly) in safe_run_decompositions."""
+  setattr(decomp, _FORCE_DECOMP_ATTR, _FORCE_DECOMP_ATTR)
+  return decomp
+
+
 def safe_run_decompositions(exported_program, decomp_table=None, can_skip=True):
   """Wrapper for ExportedProgram.run_decompositions to handle unexpected export behavior."""
 
@@ -79,6 +89,14 @@ def safe_run_decompositions(exported_program, decomp_table=None, can_skip=True):
       # back to one aten.view.
       node.target = lambda self, size: torch.reshape(self.contiguous(), size)
 
+    # Torch may skip some decompositions even if target is in decomp_table.
+    # The following ensures the target is always run through the decompositions
+    # shallowly if it has _FORCE_DECOMP_ATTR.
+    if decomp_table and node.target in decomp_table:
+      decomp = decomp_table[node.target]
+      if hasattr(decomp, _FORCE_DECOMP_ATTR):
+        node.target = decomp
+
   exported_program = exported_program.run_decompositions(decomp_table)
 
   if hasattr(torch.ops.aten, "_assert_tensor_metadata"):

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/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.8.0.dev20251206"
+__version__ = "0.8.0.dev20260105"

{ai_edge_torch_nightly-0.8.0.dev20251206.dist-info → ai_edge_torch_nightly-0.8.0.dev20260105.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-edge-torch-nightly
-Version: 0.8.0.dev20251206
+Version: 0.8.0.dev20260105
 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.8.0.dev20251206.dist-info → ai_edge_torch_nightly-0.8.0.dev20260105.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=igYuweKQrisARFhIhK3eg0wnxeYvQtBYWgM2qt2tV7M,806
+ai_edge_torch/version.py,sha256=PMprc1uQTZmTOISS9DIrSA9E1I9zJZ2iLY3JxxfbY0w,806
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=JqGZZGbpTmYiT-ta07IQbJ9-gFm-3Vip2aSzW9ulIng,6117
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -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=V-vhvScNrE3nTenT1LIULbWIU-2FD_OHthFZCPuxtzk,3480
+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
@@ -259,7 +259,7 @@ ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=uJ-niilt1c-D6QJzLwgvCUf62l
 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/_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
@@ -276,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.8.0.dev20251206.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.8.0.dev20251206.dist-info/METADATA,sha256=-BM6OyCHwhI74Y8OYtnKNLt6dYiKGkEr1mSZK5hv9Lg,2399
-ai_edge_torch_nightly-0.8.0.dev20251206.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
-ai_edge_torch_nightly-0.8.0.dev20251206.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.8.0.dev20251206.dist-info/RECORD,,
+ai_edge_torch_nightly-0.8.0.dev20260105.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.8.0.dev20260105.dist-info/METADATA,sha256=GC2dZfzgEolabjHuKdFelKgOCCyeWKGSOur-JTxPT-g,2399
+ai_edge_torch_nightly-0.8.0.dev20260105.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
+ai_edge_torch_nightly-0.8.0.dev20260105.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.8.0.dev20260105.dist-info/RECORD,,