Pre-calculating non-recurrent term is a good approach. You can use Tensor.unfold to create (16000-2,32,3) feature tensor. You could then also apply 3x128 map to it (with conv1d or matmul), but most rnn implementations do just that with input-to-hidden matrix.
Now, recurrent part is tricky. There are RNN implementations with independent hidden-to-hidden transitions - SRU, IndRNN among others, they could almost do what you want (i.e. they act as a stack of width 1 rnns), with some tweaks. But I’m not aware of implementations that look two steps back (maybe it is possible to emulate this somehow, not sure).



 PaulC:

May implementing in C++ (as it is the case for pytorch for the recurrence relation) become necessary to achieve good performance ?


And to this I would say yes. I’m sceptical about python loops with hundreds of steps already - timestep data slices are small, invocation overheads are huge, backward graph is a chain of small ops too. Actually, in my experience, such loops with GPU tensors are slower than with cpu tensors.