LSTM and GRU code

1

Hi there,

I was looking through the PyTorch source code for LSTM and GRU, and I don’t see where their equations are in the code. Can anyone point me in the right direction?

Thanks!

2

They map to C++ code in ATen/native/RNN.cpp, the same in the cudnn subdirectory or native cuda kernels in the cuda subdirectory’s RNN.cu.
You can see by following the same method as described in my Selective excursion into PyTorch internals blog post, with a little detour at the beginning from torch.nn.LSTM.forward via torch.nn.modules.rnn._rnn_impls.

Best regards

Thomas