if you want to implement a custom RNN, it is better to do it by yourself in a separate file, rather than modifying the internals. Here’s an example of people implementing custom LSTMs in a separate file: Implementation of Multiplicative LSTM