How to print with JAX

user_13394817 0
Viewed 1067

I have a JAX Boolean array and want to print a statement combined with sum of Trues:

import jax
import jax.numpy as jnp
from jax.experimental.host_callback import id_print

@jax.jit
def overlaps_jax():
    mask_cp = jnp.array([True, False, False, True, False, True, False, True, True])
    id_print(jnp.sum(mask_cp))

overlaps_jax()

There are 5 Trues in the mask_cp; I want to print as:

With jax accelerator
There are 5 true bools

As this function is jitted, I tried to print this by using id_print, but I couldn't. id_print(jnp.sum(mask_cp)) will print 5, but I couldn't use this with the strings. I have tried the followings:

id_print(jnp.sum(mask_cp))
# print:
# 5

id_print("\nWith jax accelerator\nThere are " + jnp.sum(mask_cp) + " true bools\n")
# error:
# TypeError: can only concatenate str (not "DynamicJaxprTracer") to str

print("\nWith jax accelerator\nThere are {} true bools\n".format(jnp.sum(mask_cp)))
# print:
# With jax accelerator
# There are Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=0/1)> true bools

How could I print such statements in this code?

jax jit
1 Answers

Please note that id_print is experimental, and its API and capabilities are subject to change. That said, I don't believe id_print has the capability to add text like this, but you can do it via a more general host_callback.call:

import jax
import jax.numpy as jnp
from jax.experimental.host_callback import call

@jax.jit
def overlaps_jax():
    mask_cp = jnp.array([True, False, False, True, False, True, False, True, True])
    call(lambda x: print(f"There are {x} true bools"), jnp.sum(mask_cp))

overlaps_jax()

The output is

There are 5 true bools
user_29378310
Related