Commit 1a27df1
authored
feat: implement LSTM and GRU operators for torchlib (#2674)
Implement aten_lstm and aten_gru operators to enable torch.onnx.export
for PyTorch LSTM and GRU layers. This addresses issue #2546.
Key features:
- Full support for multi-layer RNNs (num_layers > 1)
- Bidirectional support (forward and backward directions)
- Handles both biased and non-biased configurations
- batch_first parameter support with automatic transposition
- Dropout support between layers (nondeterministic seeded)
- Proper gate reordering for ONNX compatibility:
* LSTM: PyTorch [i,f,g,o] -> ONNX [i,o,f,g]
* GRU: PyTorch [r,z,n] -> ONNX [z,r,n]
Implementation details:
- Uses ONNX LSTM/GRU operators with proper parameter formatting
- Handles weight matrix transposition and reshaping
- Correctly concatenates biases using op.Concat
- Processes each layer independently with proper state management
- Returns outputs in PyTorch-compatible format
Closes: #2546
Also resolves:
- pytorch/pytorch#120626 (GRU export)
- pytorch/pytorch#123089 (LSTM export)
- pytorch/pytorch#164834 (GRU dynamo export)1 parent 93783ee commit 1a27df1
File tree
2 files changed
+496
-0
lines changed- onnxscript/function_libs/torch_lib/ops
- tests/function_libs/torch_lib
2 files changed
+496
-0
lines changed
0 commit comments