def f(x):

    def body(x, _):
        summary.summary(x, name='x', mode='append')
        return ((x + 1), ())
    return lax.scan(body, x, jnp.arange(10.0))[0]
def f(x):

    def body(x, _):
        summary.summary(x, name='x', mode='append')
        return ((x + 1), ())
    return lax.scan(body, x, jnp.arange(10.0))[0]