How to extract the carry values of every "iteration" when using nn.scan?

[Maexxel]

Is there any way to extract the carry values of every iteration when using nn.scan? For example, if I pass an entire sequence of values to the RNN, scan will iterate over that array and concatenate the output of the GRU cell at every entry of the sequence. Is there a way to also get the carry values in a similar fashion?

Currently, I'm simply passing one entry of the sequence at a time and returning the carry at every step. However, this approach is slow.

I have the following RNN using a GRU cell:

class SimpleGRU(nn.Module):
  """A simple unidirectional RNN."""

  hidden_size: int
  out_dim: int

  @nn.compact
  def __call__(self, carry, x, inspect=False):
    cell = nn.scan(nn.GRUCell,
                  variable_broadcast='params',
                  in_axes=1,
                  out_axes=1,
                  split_rngs={'params': False},
                  )
    new_carry, cell_out = cell(self.hidden_size)(carry, x)
    dense_out = nn.Dense(self.out_dim, use_bias=True)(cell_out)

    if inspect:
      return new_carry, dense_out
    return dense_out

  def initialize_carry(self, input_shape):
    # Use fixed random key since default state init fn is just zeros.
    return nn.GRUCell(self.hidden_size, parent=None).initialize_carry(
        jax.random.key(0), input_shape
    )

And this is my current solution:

gru_model = SimpleGRU(10, 2)

for seq in dataset.xs[:5]:
    print("Carry Reset")
    carry = gru_model.initialize_carry(np.array([[0]]).shape)
    for input in seq:
        carry, output = gru_model.apply({'params': trained_model_state.params}, carry, np.array([input]), True)
        print(f"For Input: {input.tolist()}, model_predicts: {jnp.argmax(output, axis=-1)}, weights: {carry}")

Any help would be greatly appreciated.