def step(params, state, input):
    (W_trans, W_out) = params
    stacked = jnp.concatenate([state, input])
    output = jnp.tanh(jnp.dot(W_out, stacked))
    next_state = jnp.tanh(jnp.dot(W_trans, stacked))
    return (next_state, output)
def step(params, state, input):
    (W_trans, W_out) = params
    stacked = jnp.concatenate([state, input])
    output = jnp.tanh(jnp.dot(W_out, stacked))
    next_state = jnp.tanh(jnp.dot(W_trans, stacked))
    return (next_state, output)