[host_callback] Add support for pmap and for passing the device to tap

[gnecula]

• Adds support for jit of pmap and pmap of pmap.
• Also add a tap_with_device optional argument to id_print and
  id_tap, to have the tap function invoked with a device keyword argument.

In presence of pmap there will be multiple devices sending data to the host. Each such
tap will be processed separately (and there may be interleaving between devices). To make
the output more understandable, we added an option to pass the device to the tap function.
For backwards-compatibility, the device is passed as a keyword argument to id_tap and only if
the tap_with_device optional flag is passed to id_tap or id_print.

Issue: #5134
Fixes: #5169

[gnecula]

I made these changes of zip -> safe_zip because while debugging my changes zip was failing silently.

[shoyer]

What do you think about treating pmap as just another transformation from the perspective of host_callback? E.g., adding an entry to the sequence of transformations, like ('pmap', {'device': ...})?

From a user perspective, I guess this might be a little more expected and would avoid needing to change the signature of host callback. On the other hand, device is part of how a computation is run (you can't nest pmap with different devices).

[gnecula]

I thought about adding pmap to the list of transforms but it is not easy at all, because pmap is not really a transformation. It is more like jit than like vmap. One of the test cases I have added is a jvp of pmap of vmap, specifically to explore this. By the time I can intervene the code is transformed (with jvp of vmap). (The same would happen if we had jvp + jit + vmap).

Perhaps you are suggesting to just add a pmap at the end of transforms just as a way to carry the device. I thought about that too, but the device on which a computation runs has significance even without pmap. That is why I think it makes sense to add the device to the host callback API.

[sbodenstein]

@gnecula: for the logging usecase, we would like the data in the form @shoyer mentioned. But there is not enough info to get it into this form with the current setup, even with the device information. Consider:

@jax.pmap
@jax.vmap
def f(x):
  x = host_callback.id_tap(print, x)
  return x**2

If we switch the vmap and pmap, the output of the function is the same. But in the one case the prints are

[0 0] (('batch', {'batch_dims': (0,)}),)
[1 1] (('batch', {'batch_dims': (0,)}),)

and for the other it is

[0 1] (('batch', {'batch_dims': (0,)}),)
[0 1] (('batch', {'batch_dims': (0,)}),)

I do realize that this is not a common way of using pmap and vmap, and perhaps its not worth worrying about. Other than this sort of scenario, does one have enough info to reconstruct the log as though it ran on a single device?

[gnecula]

Unfortunately, I cannot surface precisely in the printing that the user had a vmap(pmap(f)), because that information is lost by the time we get to execute the code. What is printed is what actually executes on each device. Probably most users would not understand that vmap(pmap(f)) is identical to pmap(vmap(f), in_axes=1). In other words, pmap is always moved to the top level of the transformation stack, taking care of adjusting in_axes. The new change to printing can be used to experiment and understand that. In fact, I have added a test_pmap_vmap and test_vmap_pmap.

[shoyer]

Right, pmap is really just sugar for launching the same program on multiple devices at the same time. I guess this distinction is lost by the time host_callback runs?

[shoyer]

is this a typo? It looks like it's always 2 CPU devices?

[gnecula]

A typo indeed, fixed. Thanks!