Printing iteration number in jitted functions/loops using host_callback.id_tap

[jeremiecoullon]

Hello! I'm writing MCMC samplers and optimisers in JAX, and I would like to be able to print the iteration number in jitted JAX functions. There are related issues (such as #196) which explain that you can't simply use Python's print function, but there's the experimental host_callback.id_tap which might be able to do this.

So I wrote a progress_bar function which only prints progress every few iterations and called it with host_callback.id_tap:

def progress_bar(i, n_iter, print_rate):
    if i%print_rate==0:
        print(f"Iteration {i}/{n_iter}")
    else:
        pass

If I run it in a normal Python loop it works as expected:

def my_python_loop():
    """
    Python loop that increments `a` 100 times
    """
    n_iter, print_rate = 100, 10
    a = 0
    for i in range(n_iter):
        host_callback.id_tap(progress_bar, i, n_iter=n_iter, print_rate=print_rate)
        a += 1
    return a

my_python_loop()

This outputs the expected result:

Iteration 0/100
Iteration 10/100
Iteration 20/100
Iteration 30/100
Iteration 40/100
Iteration 50/100
Iteration 60/100
Iteration 70/100
Iteration 80/100
Iteration 90/100

100

However if I jit the my_python_loop function, then it only prints out the iteration numbers the first time it's called (during compilation) but not every time after that.

Furthermore, if I modify the function in the following way:

@jit
def my_python_loop2(a):
    """
    Python loop that increments `a` 100 times
    """
    n_iter, print_rate = 100, 10
    for i in range(n_iter):
        a = host_callback.id_tap(progress_bar, a, n_iter=n_iter, print_rate=print_rate)
        a += 1
    return a

Then this function prints out the iterations. But this is only because I pass a into id_tap rather than the iteration index i. My guess as to what's happening is that here a is allowed to change even once the function is compiled, whereas in the previous case it is static. Note that if I add the decorator @partial(jit, static_argnums=(0,) in my_python_loop2 then the iteration numbers are not printed anymore once it's compiled.

So my question: is there a practical way to print out the iteration number in jitted functions (in particular using the scan function)? If you run scan within a jitted function then the iteration number will always need to be static so the trick above wouldn't work. I realise there are related issues about this (such as #196) but I was wondering if this specific use case is possible with the host_callback module.

[shoyer]

You definitely need to include a data dependency in order to ensure the callback runs at the correct time.

Here are two versions that work with jit, with and without scan:

from jax.experimental import host_callback
from jax import jit

def progress_bar(arg, transforms):
    i, n_iter, print_rate = arg
    if i % print_rate==0:
        print(f"Iteration {i}/{n_iter}")
    else:
        pass

@jit
def my_python_loop2(a):
    """
    Python loop that increments `a` 100 times
    """
    n_iter, print_rate = 100, 10
    for i in range(n_iter):
        a = host_callback.id_tap(progress_bar, (i, n_iter, print_rate), result=a)
        a += 1
    return a


@jit
def my_python_loop3(a):
    """
    Python loop that increments `a` 100 times
    """
    n_iter, print_rate = 100, 10
    a, _ = jax.lax.scan(
        lambda a, i: (
            host_callback.id_tap(progress_bar, (i, n_iter, print_rate), result=a) + 1,
            None),
        init=a,
        xs=jax.numpy.arange(n_iter),
    )
    return a

(Note that this is using the newer host_callback interface, you may need to upgrade to the newest version of JAX to use this)

[jeremiecoullon]

Hi Stephan, thanks very much for your answer!

It hadn't "clicked" for me how to use id_tap, but this makes sense! I upgraded my version of jax as you suggest.

Also, I made the following modification to the solution so that the Python print function doesn't get called at every iteration (and rather only gets called when it actually needs to print). This makes the loop run faster:

import jax
from jax.experimental import host_callback

def _print_consumer(arg, transrorm):
    i, n_iter = arg
    print(f"Iteration {i}/{n_iter}")

@jax.jit
def progress_bar(arg, result):
    "Print progress of loop only if iteration number is a multiple of the print_rate"
    i, n_iter, print_rate = arg
    result = jax.lax.cond(
        i%print_rate==0,
        lambda _: host_callback.id_tap(_print_consumer, (i, n_iter), result=result), 
        lambda _: result,
        operand=None)
    return result

@jax.jit
def jax_loop(a):
    """
    Jax loop that increments `a` 100 times
    """
    n_iter, print_rate = 100, 10
    def body(carry, x):
        carry = progress_bar((x, n_iter, print_rate), carry)
        carry += 1
        return carry, None
    carry, _ = jax.lax.scan(body, a, jax.numpy.arange(n_iter))
    return carry

Also, was this question more appropriate for the Discussions section rather than the issue tracker?

Thanks again.

[shoyer]

@jeremiecoullon using a conditional to pass back pass data into Python is a nice improvement!

Per your suggestion, I've converted this issue into a discussion.