def f(carry_init):

    def scan_body(c, x):
        return ((c[1], c[2], c[3], 0.0), None)
    return lax.scan(scan_body, (0.0, 1.0, 2.0, carry_init), jnp.zeros(2))
def f(carry_init):

    def scan_body(c, x):
        return ((c[1], c[2], c[3], 0.0), None)
    return lax.scan(scan_body, (0.0, 1.0, 2.0, carry_init), jnp.zeros(2))