You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
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)
0 commit comments