Skip to content

Commit 1a27df1

Browse files
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

2 files changed

+496
-0
lines changed

0 commit comments

Comments
 (0)