ai-edge-torch-nightly 0.7.0.dev20251020__py3-none-any.whl → 0.7.0.dev20251022__py3-none-any.whl

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import logging
 import math
 import operator
 from typing import Optional, Union
@@ -101,9 +102,12 @@ def _hann_window_impl(
 def _aten_hann_window_default(
     lctx: LoweringContext,
     size: int,
-    *,
-    dtype: Optional[torch.dtype] = None,
+    **kwargs,
 ) -> ir.Value:
+  dtype = kwargs.pop("dtype", None)
+  layout = kwargs.pop("layout", torch.strided)
+  if layout != torch.strided:
+    logging.warning("hann_window only supports torch.strided layout.")
   return _hann_window_impl(lctx, size, True, dtype)
 
 
@@ -114,9 +118,12 @@ def _aten_hann_window_periodic(
     lctx: LoweringContext,
     size: int,
     periodic: bool,
-    *,
-    dtype: Optional[torch.dtype] = None,
+    **kwargs,
 ) -> ir.Value:
+  dtype = kwargs.pop("dtype", None)
+  layout = kwargs.pop("layout", torch.strided)
+  if layout != torch.strided:
+    logging.warning("hann_window only supports torch.strided layout.")
   return _hann_window_impl(lctx, size, periodic, dtype)
 
 
@@ -318,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.dev20251020"
+__version__ = "0.7.0.dev20251022"

{ai_edge_torch_nightly-0.7.0.dev20251020.dist-info → ai_edge_torch_nightly-0.7.0.dev20251022.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.7.0.dev20251020
+Version: 0.7.0.dev20251022
 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
@@ -37,17 +37,7 @@ 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"
-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
+Requires-Dist: torch-xla>=2.4.0; extra == "torch-xla"
 
 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.dev20251020.dist-info → ai_edge_torch_nightly-0.7.0.dev20251022.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=7_nrKPXeG7PcqOTpJYoiFjOUpDJDKrRA6WU8ZIEVPDY,806
+ai_edge_torch/version.py,sha256=EHjxIt8UozzU1M_dt9HhtAkPXLZ4-nFSjG3rQegVtSo,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=-zKZtqvOx3COBhjCDtiZWMn5fY-boktkWXjp5Kepiro,14716
+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.dev20251020.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.7.0.dev20251020.dist-info/METADATA,sha256=8b4KHY8SsXzytP_fKt5POyMHaxcoxbk5MOxToS4SPHo,2297
-ai_edge_torch_nightly-0.7.0.dev20251020.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
-ai_edge_torch_nightly-0.7.0.dev20251020.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.7.0.dev20251020.dist-info/RECORD,,
+ai_edge_torch_nightly-0.7.0.dev20251022.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.7.0.dev20251022.dist-info/METADATA,sha256=0QfuVTBKI9hx8RFCAn0FCq5HdeZed6x4rI8A3iROHzA,2074
+ai_edge_torch_nightly-0.7.0.dev20251022.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.7.0.dev20251022.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.7.0.dev20251022.dist-info/RECORD,,

{ai_edge_torch_nightly-0.7.0.dev20251020.dist-info → ai_edge_torch_nightly-0.7.0.dev20251022.dist-info}/WHEEL

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