torch-mlir/lib/Conversion
Stephen Baione d49eabb3fc
Add Op for `torch.aten.unfold` (#3772)
# Description

Implementation of the op for `torch.aten.unfold`: [TorchToLinalg Op
Support #347](https://github.com/nod-ai/SHARK-ModelDev/issues/849)

Documentation of op can be found here: [PyTorch
Docs](https://pytorch.org/docs/stable/generated/torch.Tensor.unfold.html)

For this op, we apply a sliding window of some `size` along a single
`dimension`, with `step` in between iterations.

`Declaration: aten::unfold(Tensor(a) self, int dimension, int size, int
step) -> Tensor(a)`

The resulting `unfolded` tensor modifies the shape of `dimension` to be
equal to the number of blocks that the sliding windows extracts/inserts,
with an additional dimension of `size` appended (the number of cols of
the output tensor directly translates from the size of the sliding
window).

So if we had a tensor of rank 3 (A x B x C), with dimension = 1, size =
2 and step = 2:

    (A x B x C) |=> (A x (B - size) // step + 1 x C x size)

After extracting the window from the input tensor, we insert the (1 x
size) slice into the output tensor. We can make this simpler by mapping
the output indices from the input indices, like they do in the official
implementation:

[PyTorch
Code](https://github.com/pytorch/pytorch/blob/main/torch/_inductor/lowering.py#L1694)
2024-10-08 21:10:43 +00:00
..
TorchConversionToMLProgram [NFC] Remove unused header files (#3386) 2024-05-30 14:30:36 +08:00
TorchOnnxToTorch onnx.LSTM - bidirectional, layout attr (#3771) 2024-10-08 11:29:49 -07:00
TorchToArith [TorchToArith] Add a lowering for `torch.add.float_int` (#3594) 2024-08-07 11:55:27 -05:00
TorchToLinalg Add Op for `torch.aten.unfold` (#3772) 2024-10-08 21:10:43 +00:00
TorchToSCF [NFC] Change to *cast instead of .*cast variants (#3405) 2024-05-30 23:45:13 -07:00
TorchToStablehlo Bump to llvm/llvm-project@e813750354 (#3765) 2024-10-04 12:08:35 -07:00
TorchToTMTensor AtenCumprodOp (#3737) 2024-09-26 18:17:22 -04:00
TorchToTensor [NFC] Remove unused header files (#3386) 2024-05-30 14:30:36 +08:00
TorchToTosa [TOSA] Add legalization for fill, flip, and round (#3768) 2024-10-07 10:28:26 -07:00
Utils AtenCumprodOp (#3737) 2024-09-26 18:17:22 -04:00
CMakeLists.txt [torch] Improve shape inference for `torch-to-linalg` path for reshapes (#3055) 2024-03-26 12:41:40 -07:00
PassDetail.h Minor fixes for `ConvertTorchConversionToMLProgram`. (#1991) 2023-04-04 09:09:58 -07:00
Passes.cpp Clang format refresh (#2812) 2024-01-29 12:59:33 -05:00