# Multiple GRU layers in TensorFlow 2

I have some TensorFlow 1 code which implements a GRU layer, and I am updating it to TensorFlow 2. So instead of the hand-written layer iterating over timesteps with a GRUCell I am using the in-built TensorFlow GRU layer. However I could not find a way to specify multiple layers using that class. The PyTorch GRU layer has a construction parameter num_layers which allows for stacking multiple GRUs:

dim = 1024
gru = torch.nn.GRU(
input_size=dim,
hidden_size=dim,
num_layers=2
)


TensorFlow 1 has the MultiRNNCell, which enables one to construct a 'stack' of rnn cells which then behave as a single cell:

num_rnn = 4
dim = 1024
gru_stack = tf.contrib.rnn.MultiRNNCell(
[tf.contrib.rnn.GRUCell(dim) for _ in range(num_rnn)])


In TensorFlow 2 I could use the StackedRNNCells class, but that's just a replacement for MultiRNNCell and does the same thing. And it still defines only a single timestep for an RNN layer.

I want to use the GRU layer because that has in-built support for CuDNN and better performance (and it requires less code). It seems the only option is to manually define multiple GRUs in a sequence. Am I right? I guess this is not a big deal, but it's slightly awkward if (as in my case) you wish to expose the number of GRU layers as a parameter for a custom layer.