Comparison of Callback Functions

[ntenenz]

There seem to be a few options for host callback functions which return data:

• jax.make_array_from_callback(...)
• jax.pure_callback(...)
• jax.experimental.host_callback.call(...)

Could someone with knowledge of these methods possibly compare/contrast them, highlighting which (if any) are preferred?

[mattjj]

Thanks for the question! These are all a bit different:

• jax.make_array_from_callback is just about array creation, like an analogue of jnp.array but more general to work in distributed multi-controller scenarios. It's different from the other two.
• jax.pure_callback is a way to call a functionally pure Python function within staged-out (e.g. jitted) code. It's one of the new set of callbacks, along with e.g. jax.debug.debug_callback.
• jax.experimental.host_callback is essentially deprecated in the sense that we don't recommend using it if you can avoid it, and instead we suggest using jax.debug.debug_callback or jax.pure_callback where possible.

[ntenenz]

Are there others (available in nightly) which I missed / are not yet documented that I should be aware of @jakevdp?

[jakevdp]

I think it's basically the ones @mattjj mentioned, along with host_callback.id_tap (which is deprecated I believe), and debug.print, which IIUC is replacing host_callback.id_print. Each individual function is documented, but what's not documented is a guide to when and where you'd use each one.

[mattjj]

cc @sharadmv for docs planning

[jakevdp]

I wasn't going to tag him because I just asked him about this again yesterday 😆

[mattjj]

Oops I'm the tag nag...

[zhaopku]

is jax.pure_callback compatible with pjit?

[zhaopku]

Thanks a lot for your super fast reply! Two more questions:

1. For pjit, in a multi-process setting, the host_callback.call() function can only execute in one process. Does pure_callback() still have this constraint?

2. I am migrating from host_callback.call() to pure_callback(), what are the practical differences between these two functions that I should be aware of when I use them with pjit?

[jakevdp]

Regarding question 1: callbacks only execute on the host, not on the multiple devices.

Regarding question 2: the new documentation in #15440 should help answer some questions. host_callback shares some attributes of pure_callback, ``io_callback, and debug.callback`. Depending on what you want to do, one of those should be a good replacement.

[zhaopku]

Sorry for confusing in question 1. Correct me if wrong: even if I have multiple hosts for a single pjit program, the pure_callback() still only executes in one of the hosts. In other words, the following still holds (quote from host_callback() docs):

"When using pjit.pjit() with multiple devices on multiple hosts, only the host for the device 0 (w.r.t. the mesh) will receive the callback"

[jakevdp]

I think that's correct. @yashk2810 or @sharadmv would know for sure.

[sharadmv]

They actually behave differently. host_callback, io_callback and debug.callback are always executed on one device (usually device 0). Pure callback currently only works in jax.jit-ted computations with only one device, and under a shard_map in computations with multiple devices. We'll be sure to include these details in our guides and docs but they are currently still not set in stone (wrt to shard_map, since it's experimental).

[BillHuang2001]

What does "guaranteed execution" mean in the document?
What could cause the callback not being executed?
Or does it simply mean that the callback is guaranteed to be executed immediately.

And what is the correct way to do async callback, like the old id_tap's behavior?
So my computation can continue going without needing to wait for the callback on the host.

[jakevdp]

"guaranteed execution" means that the compiler will not remove the callback from the program.

In general, when the ouput of some primitive does not affect the output of the program, the compiler is free to ignore it. Here's a simple example:

import jax

def f(x):
  for i in range(100):
    y = jax.lax.sin(x)
  return x

print(jax.jit(f).lower(1.0).compile().as_text())
# HloModule jit_f, entry_computation_layout={(f32[])->f32[]}, allow_spmd_sharding_propagation_to_output={true}

# ENTRY %main.2 (Arg_0.1: f32[]) -> f32[] {
#   %Arg_0.1 = f32[] parameter(0), sharding={replicated}
#   ROOT %copy = f32[] copy(f32[] %Arg_0.1)
# }

The compiler never executes our 100 lax.sin evaluations, because they're not used in the output. If you did the same with pure_callback, the compiler would be free to remove it. If you did the same with io_callback, the compiler would never remove it because execution is guaranteed, under the assumption that the function has side-effects that are important to the user.

[BillHuang2001]

But why debug.callback doesn't guarantee execution? The document gives an example of using it to call print, but since it's return value isn't used, the callback might be ignored. Right?

Another question is that are callbacks executed asynchronously?
For example, if I have a callback that saves some arrays on the disk for logging purpose. Is there a way to run my callback on the host while continue executing the JAX program on the device?

[jakevdp]

But why debug.callback doesn't guarantee execution?

I believe it's because the intent of debug.callback is to not modify the execution semantics of the program. @sharadmv might be able to say more.

[sharadmv]

The document gives an example of using it to call print, but since it's return value isn't used, the callback might be ignored.

We currently prevent XLA from DCEing it but we reserve the right to DCE it from JAX if it makes sense to do so.

Another question is that are callbacks executed asynchronously? For example, if I have a callback that saves some arrays on the disk for logging purpose. Is there a way to run my callback on the host while continue executing the JAX program on the device?

I think it depends on the platform. My understanding is that callbacks are executed on the main thread on CPUs but not on TPUs.

If you wanted to guarantee executing something asynchronously, you could use a jax.experimental.io_callback that kicks off a thread that does the work you want.

[BillHuang2001]

Thanks, that helps a lot. But I must say, the example of calling print using debug.callback is confusing. It seems io_callback is the right tool for that job, not debug.callback.

[sharadmv]

Debug print will also reveal some of the details of how functions are transformed in Jax. Io callback is more restrictive in how it can be transformed. For example it cannot be naively vmapped. This makes it less useful as a drop-in print statement since it may not work in all circumstances.

[mattjj]

I'd think of an io_callback-based print as a "real print" and a debug.callback-based print (aka debug.print) as a print statement for debugging purposes. The former has cleaner semantics, like you'd want for something producing valuable output of your program, and those more well-defined semantics are why it comes with more restrictons on its use as Sharad mentioned. The latter isn't for producing program output so much as just inspection tooling. That's why we put "debug" in the name!

It sounds like perhaps you have a "real print" in mind, in which case I'd agree that io_callback fits your expectations better.