def rnn(params, inputs):
    init_state = jnp.zeros(n_hid)
    (_, outputs) = lax.scan(partial(step, params), init_state, inputs)
    return outputs
def rnn(params, inputs):
    init_state = jnp.zeros(n_hid)
    (_, outputs) = lax.scan(partial(step, params), init_state, inputs)
    return outputs