Sharad Vikram
Sharad Vikram
Seems like callbacks do not work right now on Cloud TPU VM because they are using the older `stream_executor` runtime but will soon switch to a newer runtime that does...
> Hi, for a non PL/ Compiler person, can you say a bit about what this means? We're exploring implementing a new control-flow primitive (`for`) that generalizes scan and is...
> Can we expect some performance improvements by re-implementing scan in terms of for ? Probably not. However, some patterns are not efficiently expressed as a scan but are as...
> This is a bit of a side discussion question, so I can make a discussion if you'd like to move it over there. Yes, I'd prefer this discussion to...
Have you tried using `jax.named_scope`? I suspect the issue is that the trace annotations don't work in a `jit`. BTW you can share Perfetto links by clicking `Share` (not necessary...
#12632 proposes a new JAX API for custom ops. Anyone interested should give us feedback!
#12632 proposes a new JAX API for custom ops. Anyone interested should give us feedback!
> Sorry for the bump, but what's the current status of the second update? The custom call on GPU is now landed but not used in host callback quite yet....
It is [`jax.pure_callback`](https://jax.readthedocs.io/en/latest/_autosummary/jax.pure_callback.html#jax.pure_callback).
For reference, the "callback for a debugging side-effect (like printing)" is [`jax.debug.callback`](https://jax.readthedocs.io/en/latest/_autosummary/jax.debug.callback.html)