Get vector input to debug_callback inside vmap

[dapatil211]

Hi,

I was wondering if something like this pattern was possible:

def log(metrics):
   # Here usually when using vmap, metrics are deaggregated and log is called one at a time 
   # on each of the values.
   # Instead I would like to aggregate them here and log them to wandb
   wandb.log(jax.tree.map(lambda x: x.mean(), metrics))
def computation(x):
   # Some stuff
   jax.debug.callback(log, metrics)

jax.vmap(computation)(inputs)

[jakevdp]

No, there's no way to make debug.callback work in the way you have in mind.

However, that sort of thing is supported by the batching semantics of pure_callback via the vmap_method argument. It may be possible for us to add similar batching behavior control to other callback types.

[dapatil211]

I see, thanks. As far as I understand, there's no way to use pure_callback for logging purposes, right? If not, being able to do that type of batching in either io_callback or debug.callback would be a nice feature to have when trying to run experiments vmapped over multiple seeds.