lax.scan converts JAX arrays into Numpy arrays

[ToshiyukiBandai]

Hi all,

I have a question regarding the behavior of lax.scan. In the following example code, the JAX array x is converted into a Numpy array inside the time_stepping function. Could anyone explain what is going on and how to avoid it?

def run():
    time_steps = jnp.ones(100)
    time = 0.0
    x = jnp.asarray([2.0])

    def time_stepping(carry, dt):
        time, x = carry
        time += dt
        jax.debug.breakpoint()
        jax.debug.print("The current time is {time}", time = time)
        x = 2*x
        carry = time, x
        return carry, x

    jax.debug.breakpoint()
    # initial time step
    init_carry = (time, x)
    carry, output = lax.scan(time_stepping, init_carry, time_steps)

    return carry, output

[jakevdp]

Hi - thanks for the question! I believe it's actually debug.callback and the utilities built on top of it (i.e. debug.print, debug.breakpoint) that is converting the JAX array to a NumPy array here, and this is by design in JAX v0.4.26 and earlier. But note that this will change in the next JAX release, once #20809 is part of the release.

In either case, the actual computation is not converting any inputs to NumPy arrays.

[ToshiyukiBandai]

Thank you for the clarification!