Stacked RNN with dynamic batch size for training and deployment

The following code is from TF's stacked RNN link.

batch_size = 3
sentence_max_length = 5
n_features = 2
new_shape = (batch_size, sentence_max_length, n_features)
x = tf.constant(np.reshape(np.arange(30), new_shape), dtype = tf.float32)

rnn_cells = [tf.keras.layers.LSTMCell(128) for _ in range(2)]
stacked_lstm = tf.keras.layers.StackedRNNCells(rnn_cells)
lstm_layer = tf.keras.layers.RNN(stacked_lstm)

result = lstm_layer(x)

How can I make varying batch size with None. Need to use batch_size = 8 for training and batch_size = 1 for deployment.

If I set new_shape as

new_shape = (None, sentence_max_length, n_features)

I have error in training. How can I set dynamic batch size in stacked RNN?

You can put your layers inside a tf model and then fit and predict with different batch size parameters:

sentence_max_length = 5
n_features = 2
# the -1 here adapts to the other numbers to get the right shape 
new_shape = (-1, sentence_max_length, n_features)  
x = tf.constant(np.reshape(np.arange(30), new_shape), dtype = tf.float32)

rnn_cells = [tf.keras.layers.LSTMCell(128) for _ in range(2)]
stacked_lstm = tf.keras.layers.StackedRNNCells(rnn_cells)
lstm_layer = tf.keras.layers.RNN(stacked_lstm)
model = tf.keras.Sequential([lstm_layer])  # you can add other layers to the sequential model
model.compile()  # <-- with optimizer and loss here
model.fit(x, y, batch_size=8)
...
y_pred = model.predict(another_x, batch_size=1)