Ordering of results printed with jax.debug.print does not match documentation

[hericks]

I recently went through the section about Runtime value debugging in JAX in the documentation. However, I didn't get the expected output for the very first example. Instead the two output lines are reordered.

I understand that distinct calls to jax.debug.print with arguments that don't depend on each other might be reordered when staged out, as in the example by jax.jit. However, in the example, the argument y of the second call does depend on x (the argument of the initial call).

The expected ordering can be restored by either using the ordered=True argument in the calls to jax.debug.print. Additionally, I noticed that the expected ordering is restored when returning x + y instead of just y.

Am I misunderstanding something about the printing dynamics here?

Example.

import jax
import jax.numpy as jnp

@jax.jit
def f(x):
  jax.debug.print("🤯 {x} 🤯", x=x)
  y = jnp.sin(x)
  jax.debug.print("🤯 {y} 🤯", y=y)
  return y
  
f(2.)

Expected output.

🤯 2.0 🤯
🤯 0.9092974662780762 🤯

Actual output.

🤯 0.9092974662780762 🤯
🤯 2.0 🤯

[sharadmv]

This is unusual...what platform are you running on?

[hericks]

I am working in a docker container running Ubuntu 20.04. More precisely, I am using nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04 and install jax 0.3.23 with jaxlib 0.3.22+cuda11.cudnn805 on top.

[sharadmv]

Unfortunately I wasn't able to reproduce your behavior, but seems plausible to me. We assumed data-dependence was sufficient to guarantee ordering of prints in an XLA program, but that assumption is probably wrong!

XLA is a functional compiler and is free to do its own optimizations, such as rematerializing intermediate values instead of storing them. In your program, therefore, XLA could choose to first execute print(sin(x)) and then execute print(x) instead of first storing y = sin(x).

These compiler optimizations should not delete the call to print but more aggressive reordering is possible due to rematerialization.

We should update our docs to reflect this potential behavior.

[hericks]

Ok, that's completely reasonable. The main takeaway is then that a user needs to expect the jax.debug.print statements to be called in any order; except when ordered = True.