Skip to content

Commit baa6b59

Browse files
jakesabathia2jakesabathia2
committed
try to make lstm, gru and rnn composite op
1 parent a43cc84 commit baa6b59

9 files changed

Lines changed: 2012 additions & 1 deletion

File tree

‎coreai_torch/_aten_to_core.py‎

Lines changed: 949 additions & 0 deletions
Large diffs are not rendered by default.

‎coreai_torch/_decomp.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,18 @@
1616
_COMPOSITE_OPS: list = [
1717
torch.ops.aten.hardsigmoid.default,
1818
torch.ops.aten.hardswish.default,
19+
torch.ops.aten.gru.input,
1920
torch.ops.aten.instance_norm.default,
21+
torch.ops.aten.lstm.input,
2022
torch.ops.aten.pixel_shuffle.default,
2123
torch.ops.aten.reflection_pad1d.default,
2224
torch.ops.aten.reflection_pad2d.default,
2325
torch.ops.aten.reflection_pad3d.default,
2426
torch.ops.aten.replication_pad1d.default,
2527
torch.ops.aten.replication_pad2d.default,
2628
torch.ops.aten.replication_pad3d.default,
29+
torch.ops.aten.rnn_relu.input,
30+
torch.ops.aten.rnn_tanh.input,
2731
torch.ops.aten.scaled_dot_product_attention.default,
2832
torch.ops.aten.silu.default,
2933
]
@@ -40,6 +44,7 @@ def get_decomp_table() -> dict:
4044
4145
* ``torch.ops.aten.hardsigmoid.default``
4246
* ``torch.ops.aten.instance_norm.default``
47+
* ``torch.ops.aten.lstm.input`` (lowered to a ``"lstm"`` composite op)
4348
* ``torch.ops.aten.pixel_shuffle.default``
4449
* ``torch.ops.aten.scaled_dot_product_attention.default``
4550

‎docs/api/composite-ops/aten-derived.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,15 @@ Composite ops recognized automatically from the ATen nodes (`fx.Node`s) in your
77
88
batch-norm
99
group-norm
10+
gru
1011
hard-sigmoid
1112
instance-norm
1213
layer-norm
1314
linalg-vector-norm
1415
log-softmax
16+
lstm
1517
pixel-shuffle
18+
rnn
1619
```
1720

1821
For the ATen ops these are derived from, see {doc}`../supported-aten-ops`.

‎docs/api/composite-ops/gru.md‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# gru
2+
3+
Gated Recurrent Unit recurrent layer. Preserved as a named `gru` composite op so the compiler and delegates can recognize the whole recurrent layer instead of a decomposed, unrolled sequence of primitive ops.
4+
5+
Per timestep, for each layer and direction:
6+
7+
$$
8+
\begin{aligned}
9+
r_t &= \sigma(W_{ir} x_t + b_{ir} + W_{hr} h_{t-1} + b_{hr}) \\
10+
z_t &= \sigma(W_{iz} x_t + b_{iz} + W_{hz} h_{t-1} + b_{hz}) \\
11+
n_t &= \tanh(W_{in} x_t + b_{in} + r_t \odot (W_{hn} h_{t-1} + b_{hn})) \\
12+
h_t &= (1 - z_t) \odot n_t + z_t \odot h_{t-1}
13+
\end{aligned}
14+
$$
15+
16+
**ATen source:** `aten.gru.input`
17+
18+
## `reset_after` and the MIL GRU
19+
20+
This composite models PyTorch's GRU, which applies the reset gate **after** the hidden matrix multiply, *including* its bias: `r_t ⊙ (W_hn h_{t-1} + b_hn)`. The CoreML MIL `gru` op instead applies the reset gate before the hidden bias (`r_t * W_ho h_{t-1} + b_ho`) — the `reset_after = false` convention. Because of this difference:
21+
22+
- coremltools does **not** lower a torch GRU to the native MIL `gru` op; it hand-builds the recurrence.
23+
- the hidden bias `b_hn` **cannot** be folded into a single combined bias, so this composite takes the input and hidden biases separately (`bias_ih`, `bias_hh`).
24+
- the composite carries a `reset_after` attribute (always `true` for `torch.nn.GRU`) so a delegate knows which convention to apply.
25+
26+
## Inputs
27+
28+
A single composite is emitted per layer; stacked layers chain one composite each, and bidirectional layers use one composite that also takes the backward weights (states packed on the hidden axis: `[:, :H]` forward, `[:, H:]` reverse).
29+
30+
| Name | Shape | Description |
31+
|---|---|---|
32+
| `x` | `(S, B, I)` | Input sequence (time-major; `batch_first` inputs are transposed first) |
33+
| `initial_h` | `(B, H)` or `(B, 2H)` | Initial hidden state |
34+
| `weight_ih` | `(3H, I)` | Input-hidden weights, PyTorch `[r, z, n]` gate layout |
35+
| `weight_hh` | `(3H, H)` | Hidden-hidden weights, `[r, z, n]` layout |
36+
| `bias_ih` | `(3H,)` | Input bias (zeros when the layer has no bias) |
37+
| `bias_hh` | `(3H,)` | Hidden bias (zeros when the layer has no bias) |
38+
| `weight_ih_back`, `weight_hh_back`, `bias_ih_back`, `bias_hh_back` | | Backward-direction inputs (bidirectional only) |
39+
40+
## Attributes
41+
42+
| Name | Type | Description |
43+
|---|---|---|
44+
| `direction` | `str` | `"forward"` or `"bidirectional"` |
45+
| `output_sequence` | `bool` | Always `true` |
46+
| `recurrent_activation` | `str` | Reset/update gate activation (`"sigmoid"`) |
47+
| `activation` | `str` | New-gate activation (`"tanh"`) |
48+
| `reset_after` | `bool` | Reset gate applied after the hidden bias (`true` for torch) |
49+
| `version` | `int` | Composite op version |
50+
51+
## Outputs
52+
53+
| Name | Shape | Description |
54+
|---|---|---|
55+
| `output` | `(S, B, DH)` | Hidden state at every timestep (`D` = 2 if bidirectional) |
56+
| `h_n` | `(B, DH)` | Final hidden state |
57+
58+
The converted graph reassembles the standard PyTorch layout (`batch_first` reapplied to `output`; `h_n` shaped `(num_layers * D, B, H)`).
59+
60+
## Data types
61+
62+
`fp16`, `fp32`, `bf16`.
63+
64+
## Limitations
65+
66+
- The sequence length must be static (known at export time).
67+
- Packed / variable-length sequences (`pack_padded_sequence`) are not supported.
68+
- Training-time dropout (`train=True`, `dropout>0`, `num_layers>1`) is not supported; dropout configured on an `nn.GRU` in inference mode is a no-op and ignored.
69+
70+
## PyTorch example
71+
72+
```python
73+
import torch
74+
75+
gru = torch.nn.GRU(input_size=4, hidden_size=3, num_layers=2, batch_first=True).eval()
76+
x = torch.randn(2, 5, 4)
77+
h0 = torch.randn(2, 2, 3)
78+
output, h_n = gru(x, h0)
79+
```
80+
81+
## Reference
82+
83+
[`torch.nn.GRU`](https://docs.pytorch.org/docs/stable/generated/torch.nn.GRU.html)

‎docs/api/composite-ops/lstm.md‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
# lstm
2+
3+
Long Short-Term Memory recurrent layer. Preserved as a named `lstm` composite op (mirroring the CoreML MIL `lstm` interface) so the compiler and delegates can recognize the whole recurrent layer instead of a decomposed, unrolled sequence of primitive ops.
4+
5+
Per timestep, for each layer and direction:
6+
7+
$$
8+
\begin{aligned}
9+
i_t &= \sigma(W_{ii} x_t + b_i + W_{hi} h_{t-1}) \\
10+
f_t &= \sigma(W_{if} x_t + b_f + W_{hf} h_{t-1}) \\
11+
g_t &= \tanh(W_{ig} x_t + b_g + W_{hg} h_{t-1}) \\
12+
o_t &= \sigma(W_{io} x_t + b_o + W_{ho} h_{t-1}) \\
13+
c_t &= f_t \odot c_{t-1} + i_t \odot g_t \\
14+
h_t &= o_t \odot \tanh(c_t)
15+
\end{aligned}
16+
$$
17+
18+
**ATen source:** `aten.lstm.input`
19+
20+
A single composite is emitted per layer. Stacked (multi-layer) LSTMs chain one composite per layer; bidirectional layers use a single composite that also takes the backward weights. Weights are reordered from PyTorch's `[i, f, g, o]` gate layout to MIL's `[i, f, o, g]`, and the input/hidden biases are summed into one `bias`, so the composite inputs match the MIL `lstm` op directly.
21+
22+
## Inputs
23+
24+
| Name | Shape | Description |
25+
|---|---|---|
26+
| `x` | `(S, B, I)` | Input sequence (time-major; `batch_first` inputs are transposed before the composite) |
27+
| `initial_h` | `(B, H)` or `(B, 2H)` | Initial hidden state (forward+reverse packed on the hidden axis when bidirectional) |
28+
| `initial_c` | `(B, H)` or `(B, 2H)` | Initial cell state |
29+
| `weight_ih` | `(4H, I)` | Input-hidden weights, `ifog` gate layout |
30+
| `weight_hh` | `(4H, H)` | Hidden-hidden weights, `ifog` gate layout |
31+
| `bias` | `(4H,)` | Combined input+hidden bias (zeros when the layer has no bias) |
32+
| `weight_ih_back` | `(4H, I)` | Backward-direction input-hidden weights (bidirectional only) |
33+
| `weight_hh_back` | `(4H, H)` | Backward-direction hidden-hidden weights (bidirectional only) |
34+
| `bias_back` | `(4H,)` | Backward-direction combined bias (bidirectional only) |
35+
36+
## Attributes
37+
38+
| Name | Type | Description |
39+
|---|---|---|
40+
| `direction` | `str` | `"forward"` or `"bidirectional"` |
41+
| `output_sequence` | `bool` | Always `true` — the full per-timestep hidden sequence is returned |
42+
| `recurrent_activation` | `str` | Gate activation (`"sigmoid"`) |
43+
| `cell_activation` | `str` | Cell activation (`"tanh"`) |
44+
| `activation` | `str` | Output activation (`"tanh"`) |
45+
| `version` | `int` | Composite op version |
46+
47+
## Outputs
48+
49+
| Name | Shape | Description |
50+
|---|---|---|
51+
| `output` | `(S, B, DH)` | Hidden state at every timestep (`D` = 2 if bidirectional) |
52+
| `h_n` | `(B, DH)` | Final hidden state |
53+
| `c_n` | `(B, DH)` | Final cell state |
54+
55+
The final `output`, `h_n`, and `c_n` returned by the converted graph follow the standard PyTorch `nn.LSTM` layout (`batch_first` reapplied to `output`; states shaped `(num_layers * D, B, H)`).
56+
57+
## Data types
58+
59+
`fp16`, `fp32`, `bf16`.
60+
61+
## Limitations
62+
63+
- The sequence length must be static (known at export time).
64+
- Packed / variable-length sequences (`aten.lstm.data`, `pack_padded_sequence`) are not supported.
65+
- Training-time dropout (`train=True`, `dropout>0`, `num_layers>1`) is not supported; a dropout rate configured on an `nn.LSTM` evaluated in inference mode is a no-op and safely ignored.
66+
67+
## PyTorch example
68+
69+
```python
70+
import torch
71+
72+
lstm = torch.nn.LSTM(input_size=4, hidden_size=3, num_layers=2, batch_first=True).eval()
73+
x = torch.randn(2, 5, 4)
74+
h0 = torch.randn(2, 2, 3)
75+
c0 = torch.randn(2, 2, 3)
76+
output, (h_n, c_n) = lstm(x, (h0, c0))
77+
```
78+
79+
## Reference
80+
81+
[`torch.nn.LSTM`](https://docs.pytorch.org/docs/stable/generated/torch.nn.LSTM.html)

‎docs/api/composite-ops/rnn.md‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
# rnn
2+
3+
Elman (simple) recurrent layer. Preserved as a named `rnn` composite op — mirroring the CoreML MIL `rnn` op — so the compiler and delegates can recognize the whole recurrent layer instead of a decomposed, unrolled sequence of primitive ops.
4+
5+
Per timestep, for each layer and direction:
6+
7+
$$h_t = \text{activation}(W_{ih} x_t + b_{ih} + W_{hh} h_{t-1} + b_{hh})$$
8+
9+
where `activation` is `tanh` or `relu`.
10+
11+
**ATen source:** `aten.rnn_tanh.input` (tanh), `aten.rnn_relu.input` (relu)
12+
13+
Both biases are additive and are combined into a single `bias = b_ih + b_hh`, matching the MIL `rnn` op. A single composite is emitted per layer; stacked layers chain one composite each, and bidirectional layers use one composite that also takes the backward weights (states packed on the hidden axis).
14+
15+
## Inputs
16+
17+
| Name | Shape | Description |
18+
|---|---|---|
19+
| `x` | `(S, B, I)` | Input sequence (time-major; `batch_first` inputs are transposed first) |
20+
| `initial_h` | `(B, H)` or `(B, 2H)` | Initial hidden state |
21+
| `weight_ih` | `(H, I)` | Input-hidden weights |
22+
| `weight_hh` | `(H, H)` | Hidden-hidden weights |
23+
| `bias` | `(H,)` | Combined bias `b_ih + b_hh` (zeros when the layer has no bias) |
24+
| `weight_ih_back`, `weight_hh_back`, `bias_back` | | Backward-direction inputs (bidirectional only) |
25+
26+
## Attributes
27+
28+
| Name | Type | Description |
29+
|---|---|---|
30+
| `direction` | `str` | `"forward"` or `"bidirectional"` |
31+
| `output_sequence` | `bool` | Always `true` |
32+
| `activation` | `str` | `"tanh"` or `"relu"` |
33+
| `version` | `int` | Composite op version |
34+
35+
## Outputs
36+
37+
| Name | Shape | Description |
38+
|---|---|---|
39+
| `output` | `(S, B, DH)` | Hidden state at every timestep (`D` = 2 if bidirectional) |
40+
| `h_n` | `(B, DH)` | Final hidden state |
41+
42+
The converted graph reassembles the standard PyTorch layout (`batch_first` reapplied to `output`; `h_n` shaped `(num_layers * D, B, H)`).
43+
44+
> **Note:** coremltools lowers torch RNN to the native MIL `rnn` op but supports **uni-directional** only. This composite additionally supports bidirectional RNN.
45+
46+
## Data types
47+
48+
`fp16`, `fp32`, `bf16`.
49+
50+
## Limitations
51+
52+
- The sequence length must be static (known at export time).
53+
- Packed / variable-length sequences (`pack_padded_sequence`) are not supported.
54+
- Training-time dropout (`train=True`, `dropout>0`, `num_layers>1`) is not supported; dropout configured on an `nn.RNN` in inference mode is a no-op and ignored.
55+
56+
## PyTorch example
57+
58+
```python
59+
import torch
60+
61+
rnn = torch.nn.RNN(input_size=4, hidden_size=3, num_layers=2, nonlinearity="tanh", batch_first=True).eval()
62+
x = torch.randn(2, 5, 4)
63+
h0 = torch.randn(2, 2, 3)
64+
output, h_n = rnn(x, h0)
65+
```
66+
67+
## Reference
68+
69+
[`torch.nn.RNN`](https://docs.pytorch.org/docs/stable/generated/torch.nn.RNN.html)

‎docs/api/supported-aten-ops.md‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ This page lists every PyTorch ATen operator that `TorchConverter` lowers to Core
66

77
- Op names use the FX qualified-name form `op_name.overload` (e.g. `add.Tensor`, `mean.dim`). When PyTorch's decomposition pipeline produces a different overload than the one listed, that overload is not supported.
88
- A few names appear without an overload suffix (e.g. `add`, `mul`, `getitem`) — these match plain Python-operator FX nodes that have no `.default` overload.
9-
- Three ops — `instance_norm.default`, `pixel_shuffle.default`, and `scaled_dot_product_attention.default` — are deliberately preserved by `get_decomp_table()` and emitted as composite ops in the lowered IR.
9+
- A few ops — `gru.input`, `instance_norm.default`, `lstm.input`, `pixel_shuffle.default`, `rnn_relu.input`, `rnn_tanh.input`, and `scaled_dot_product_attention.default` — are deliberately preserved by `get_decomp_table()` and emitted as composite ops in the lowered IR.
1010
- All ops below are resolved through the registry in `coreai_torch._aten_to_core`. To override a built-in lowering with your own, pass `allow_override=True` to `register_torch_lowering()`.
1111

1212
## ATen ops
@@ -75,6 +75,7 @@ This page lists every PyTorch ATen operator that `TorchConverter` lowers to Core
7575
| `gelu.default` | |
7676
| `getitem` | |
7777
| `gt.Scalar`, `gt.Tensor` | |
78+
| `gru.input` | Preserved as composite by `get_decomp_table()` |
7879
| `hardsigmoid.default` | Lowered as a composite |
7980
| `hardswish.default` | |
8081
| `hardtanh.default` | |
@@ -95,6 +96,7 @@ This page lists every PyTorch ATen operator that `TorchConverter` lowers to Core
9596
| `logical_not.default` | |
9697
| `logical_or.default` | |
9798
| `logical_xor.default` | |
99+
| `lstm.input` | Preserved as composite by `get_decomp_table()` |
98100
| `lt.Scalar`, `lt.Tensor` | |
99101
| `max.default`, `max.dim` | |
100102
| `max_pool2d_with_indices.default` | |
@@ -122,6 +124,7 @@ This page lists every PyTorch ATen operator that `TorchConverter` lowers to Core
122124
| `remainder.Tensor` | |
123125
| `repeat.default` | |
124126
| `replication_pad1d.default`, `replication_pad2d.default`, `replication_pad3d.default` | Lowered to `coreai.pad` with `replicate` mode |
127+
| `rnn_relu.input`, `rnn_tanh.input` | Preserved as composite by `get_decomp_table()` |
125128
| `round.default`, `round.decimals` | |
126129
| `rsqrt.default` | |
127130
| `scalar_tensor.default` | |

0 commit comments

Comments
 (0)