dsr:

    if states is None:
      states = self.gru.get_initial_state(x)
    x, states = self.gru(x, initial_state=states, training=training)
    x, states = self.gru2(x, initial_state=states, training=training)
    x = self.dense(x, training=training)

    if return_state:
      return x, states



You treat GRU states as an intermediate output processed sequentially (you create some states, pass them to the first layer, get back new states, pass them to the second layer, then retrieve the final states…). But that’s not what RNN states are! They’re parallel, not sequential! The state of GRU2 is completely independent from the state of GRU1.
If you want to do stateful processing, then you should be loading both states independently at the start of each batch (and also returning both states at the end).
The good news is, a GRU has only has state vector, so the code will be a bit easier than if you were dealing with a LSTM, which has multiple state vectors. You’d do something like:
    if states is None:
      state1 = self.gru1.get_initial_state(x)
      state2 = self.gru2.get_initial_state(x)
    else:
       state1, state2 = states
    x, state1 = self.gru1(x, initial_state=state1, training=training)
    x, state2 = self.gru2(x, initial_state=state2, training=training)
    x = self.dense(x, training=training)

    if return_state:
      return x, [state1, state2]