Print Value in a non-Jitted Function?

[jacob-jacob-jacob]

Hi, I want to debug a function. I am passing in Jax objects and using lax.scan. I am not jitting the function though.

Still, when I print an array, I only get something like Traced<ShapedArray(float32[2])>with<DynamicJaxprTrace(level=1/0)>. How can I print the array contents when not jitting?

[YouJiacheng]

It is because that lax.scan will compile the body function.
I suggest that you can define a mock scan when you debug.
I think following code has nearly the same semantics as lax.scan.

DEBUG = True
def scan(f, init, xs, length=None, **kwargs):
    if not DEBUG:
        return lax.scan(f, init, xs, length, **kwargs)

    carry = init
    xs_flat = jax.tree_flatten(xs)[0]
    lengths = [x.shape[0] for x in xs_flat]
    if length is None:
        length = lengths[0] # not support xs is None when length is None
    assert all(length == l for l in lengths)
    def xs_generator():
        for i in range(length):
            yield jax.tree_map(lambda x: x[i], xs)
    ys = []
    for x in xs_generator():
        carry, y = f(carry, x)
        ys.append(y)
    
    return carry, jax.tree_map(lambda *x: jnp.stack(x), *ys)

[jakevdp]

scan will JIT compile the body function by default. If you want to disable this behavior, you can execute it in a disable_jit context; for example:

import jax
import jax.numpy as jnp

def f(x):
  def body_fun(carry, x):
    print(x)
    return carry, x
  return jax.lax.scan(body_fun, 1, x)

f(jnp.arange(4))
# Output:
# Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>

with jax.disable_jit():
  f(jnp.arange(4))
# Output:
# 0
# 1
# 2
# 3

[YouJiacheng]

Thanks for the exact solution ! I overlooked the if config.jax_disable_jit: branch in lax.scan (I didn't look into it, and just think XLA while need a jaxpr body)! 🤣🤣🤣

[YouJiacheng]

Maybe the document of disable_jit should be updated.

For debugging it is useful to have a mechanism that disables jit() everywhere in a dynamic context.

The document says that disable_jit will disable jit(), not mentions that it will also affect behavior of control flow operators such as lax.scan.