ai-edge-torch-nightly 0.7.0.dev20251021__py3-none-any.whl → 0.7.0.dev20251023__py3-none-any.whl

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -325,6 +325,85 @@ def _aten_cat(lctx: LoweringContext, tensors, dim=0):
   return stablehlo.concatenate(non_empty_tensors, dim)
 
 
+# Schema:
+#   - aten::unfold(Tensor self, int dim, int size, int step) -> Tensor
+# Torch Reference:
+#   - https://pytorch.org/docs/stable/generated/torch.Tensor.unfold.html
+@lower(torch.ops.aten.unfold.default)
+def _aten_unfold(lctx, x: ir.Value, dim: int, size: int, step: int):
+  x_shape = x.type.shape
+  rank = len(x_shape)
+  if dim < 0:
+    dim += rank
+
+  num_windows = (x_shape[dim] - size) // step + 1
+  batch_shape = list(x_shape[:dim]) + [num_windows] + list(x_shape[dim + 1 :])
+
+  # Create start_indices for gather.
+  # The shape of start_indices will be batch_shape + [rank].
+  # start_indices[b_0,...,b_{rank-1}] will be [p_0,...,p_{rank-1}] where
+  # p_j = b_j for j != dim and p_dim = b_dim * step.
+  indices_parts = []
+  i64 = ir.IntegerType.get_signless(64)
+  for i in range(rank):
+    bshape = [1] * rank
+    bshape[i] = batch_shape[i]
+    dim_len = batch_shape[i]
+
+    iota = stablehlo.IotaOp(
+        ir.RankedTensorType.get([dim_len], i64),
+        iota_dimension=ir.IntegerAttr.get(i64, 0),
+    ).result
+    if i == dim:
+      iota = stablehlo.multiply(iota, utils.splat(step, i64, [dim_len]))
+
+    iota_reshaped = stablehlo.reshape(
+        ir.RankedTensorType.get(bshape, i64), iota
+    )
+    indices_parts.append(
+        stablehlo.broadcast_in_dim(
+            ir.RankedTensorType.get(batch_shape, i64),
+            iota_reshaped,
+            ir.DenseI64ArrayAttr.get(list(range(rank))),
+        )
+    )
+
+  # For each dimension i, indices_parts[i] contains the i-th coordinate
+  # of start_indices. We unsqueeze each part to shape batch_shape + [1]
+  # and concatenate along the new dimension to produce start_indices of
+  # shape batch_shape + [rank].
+  unsqueezed_parts = [
+      stablehlo.reshape(ir.RankedTensorType.get(batch_shape + [1], i64), part)
+      for part in indices_parts
+  ]
+  start_indices = stablehlo.concatenate(
+      unsqueezed_parts, ir.IntegerAttr.get(i64, rank)
+  )
+
+  slice_sizes_list = [1] * rank
+  slice_sizes_list[dim] = size
+  slice_sizes = ir.DenseI64ArrayAttr.get(slice_sizes_list)
+
+  collapsed_slice_dims_list = [i for i in range(rank) if i != dim]
+
+  dnums = stablehlo.GatherDimensionNumbers.get(
+      offset_dims=[rank],
+      collapsed_slice_dims=collapsed_slice_dims_list,
+      operand_batching_dims=[],
+      start_indices_batching_dims=[],
+      start_index_map=list(range(rank)),
+      index_vector_dim=rank,
+  )
+
+  return stablehlo.gather(
+      x,
+      start_indices,
+      dnums,
+      slice_sizes,
+      indices_are_sorted=ir.BoolAttr.get(False),
+  )
+
+
 # Schema:
 #   - aten::slice_scatter(Tensor self, Tensor src, int dim=0, SymInt?
 #       start=None, SymInt? end=None, SymInt step=1) -> Tensor

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.dev20251021"
+__version__ = "0.7.0.dev20251023"

{ai_edge_torch_nightly-0.7.0.dev20251021.dist-info → ai_edge_torch_nightly-0.7.0.dev20251023.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.7.0.dev20251021
+Version: 0.7.0.dev20251023
 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.7.0.dev20251021.dist-info → ai_edge_torch_nightly-0.7.0.dev20251023.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=OqJtg7_I5JrjrCcK_c026Q_w0JSHTPq1q3hHLlxAzt8,806
+ai_edge_torch/version.py,sha256=JI_4o5ccLYryjzSht5hbYIej5kJkl1edsjryPo70EPE,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
@@ -250,7 +250,7 @@ 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=kUbicoJcYZOgYRhydXYIsLyB2lW_Y-39skvjvrqhevo,15031
+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
@@ -270,8 +270,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.dev20251021.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.7.0.dev20251021.dist-info/METADATA,sha256=wB295znY-1RX_QXKMB8e-ujDBHQg_u39kqx8Pk8qleg,2074
-ai_edge_torch_nightly-0.7.0.dev20251021.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.7.0.dev20251021.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.7.0.dev20251021.dist-info/RECORD,,
+ai_edge_torch_nightly-0.7.0.dev20251023.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.7.0.dev20251023.dist-info/METADATA,sha256=2E0SW9n44kghVciReTO-DMBxRxWgVfhmZctt_IFMV4c,2074
+ai_edge_torch_nightly-0.7.0.dev20251023.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.7.0.dev20251023.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.7.0.dev20251023.dist-info/RECORD,,