Better jax.debug.print

[yliu120]

Problems we want to solve

jax.debug.print currently works under most environments but it is pretty costly.

1. Always all gather data under GSPMD mode
   Under the GSPMD mode, it only supports maximal sharding as inherited from io_callback so that when a program has debug.print, it will be way more inefficient as lots of communications are added.

2. Going across Python interpreter/Numpy
   Basically these are overheads that lead to slow printing.

3. Unstable IR and disrupting AOT
   The current impl is based on io_callback which puts the opaque pointer of the callback in the stablehlo IR so that each machine ends up with something different. (Related to [Lowering] Stable IR #25123)

How to make it better

1. 3 is WIP.
2. We can implement debug.print callback with FFI in XLA so that it will become a precompiled callback in the jaxlib. The custom call will have sharding propagated and will be partitioned as of the inputs. Here is what it will look like if programming in python,

def print(x):
  x = ffi_call.print(x) ==> custom call following the same sharding as x.
  return x

[yliu120]

This issue reveals an offline conversation with Peter @hawkinsp and @dfm

[dfm]

Also pinging @danielsuo who has been thinking about this.

[yliu120]

A few code samples scattering around:

1. Having custom_partitioning calls working with debug.print ([io_callback] Adds test for io_callback being used inside custom partitioning #23620)
2. Having debug.print (or io_callback) working inside partial auto ([Shmap/PartialAuto] Temporary solution for debug.print inside a partial-auto shard map. #25705)