What is the most Idiomatic way of dealing with temporal states in modules?

[PartiallyUntyped]

Hi,

My goal is to implement a stateful module where the state, a memory bank, exists only temporary and is reset after a certain amount of iterations, which will be indicated through a parameter. What is the most idiomatic way of approaching this? I would like to keep this limited to just jax and stax, so no Haiku or Flax.

The idea I came up with is basically to store a list as part of the parameters, and in that list to store the state, and replace it as needed.

def MemoryModule(bank_size, G, O):
    """
    bank_size is the number of rows in the memory bank.
    G is a module from input x bank to bank
    O is a module from input x bank to output
    """
    def init_fun(rng, input_shape):
        bank = jnp.zeros((bank_size, input_shape))
        _, g_params = G[0](rng, input_shape)
        out_dim, out_params = O[0](rng, input_shape)
        return out_dim, ([bank], g_params, out_params)

    def apply_fun(params, inputs, **kwargs):
        (bank,*_), g_params, out_params = params
        if kwargs.get("reset_bank", False):
            bank = jnp.zeros((bank_size, bank.shape[1]))
        bank = G[1](g_params, (inputs, bank), **kwargs)
        params[0][0] = bank
        return O[1](out_params, (inputs, bank), **kwargs)
    
    return init_fun, apply_fun