Custom RNN stateful model BATCH SIZE training - tensorflow

Question

Willing to train a custom, stateful RNN, but failing to do it so over batches. I followed documentation of tf.keras.layers.RNN, but still workless.
Stand-alone code:
import tensorflow as tf
class RecurrentCell(tf.keras.layers.Layer):
def __init__(self,
units,
**kwargs):
self.units = units
self.state_size = units
self.output_size = 5
super(RecurrentCell, self).__init__(**kwargs)
def build(self, input_shape):
self.training_params = self.add_weight(shape=(1, self.output_size),
initializer='uniform',
name='kernel')
self.built = True
def call(self, inputs, states):
new_states = states
output = inputs + self.training_params
return output, [new_states]
def get_initial_state(self, inputs=None, batch_size=10,dtype=np.float32):
return tf.zeros( tuple([batch_size]) + tuple([self.state_size]), dtype=np.float32)
class RecurrentModel(tf.keras.Model):
def __init__(self, **kwargs):
super(RecurrentModel,self).__init__()
self.recurrent_layer =tf.keras.layers.RNN(RecurrentCell(units=5),
return_sequences=True,
stateful=True,
)
def call(self, inputs):
return self.recurrent_layer(inputs)
class CustomCallback(tf.keras.callbacks.Callback):
def on_train_batch_begin(self, batch, logs=None):
keys = list(logs.keys())
print("...Training: start of batch {}; got log keys: {}".format(batch, keys))
model = RecurrentModel()
model.compile(optimizer="adam", loss="mse")
batch_size = 2
time_index = 10
features = 5
inputs = tf.random.uniform((batch_size, time_index, features))
outputs = tf.random.uniform((batch_size, time_index, features))
model.fit(x=inputs, y=outputs, batch_size=2, shuffle=False, epochs=10, callbacks=[CustomCallback()])
In particular, documentation seems to be also inconsistent with Sequential model (haven't tried the functional API).
Thanks a lot!
matias