def loss(A):

    def step(x, i):
        return (jnp.matmul(A, x), None)
    init_x = jnp.zeros(A.shape[(- 1):])
    (last_x, _) = lax.scan(step, init_x, jnp.arange(10))
    return jnp.sum(last_x)
def loss(A):

    def step(x, i):
        return (jnp.matmul(A, x), None)
    init_x = jnp.zeros(A.shape[(- 1):])
    (last_x, _) = lax.scan(step, init_x, jnp.arange(10))
    return jnp.sum(last_x)