Tom Hennigan

icon indicating a website link https://tom.gd icon indicating an email [email protected]

icon company @google-deepmind icon location London, UK icon location Software Engineer at @google-deepmind. Previously: @deepmind, @google, @duedil-ltd.

Results 1 issues of Tom Hennigan

Using a device function inside the host function of host_callback fails confusingly

6 comment

Using a JAX device function in the host fn of `host_callback` causes hangs/long compile/oom (breakage changes with the backend). This not being supported is probably WAI, but we could provide...

bug
P1 (soon)
NVIDIA GPU