def loss(params, inputs, targets):
    predictions = rnn(params, inputs)
    return jnp.sum(((predictions - targets) ** 2))
def loss(params, inputs, targets):
    predictions = rnn(params, inputs)
    return jnp.sum(((predictions - targets) ** 2))