Closed mfschubert closed 7 months ago
Update the eigensolve to use jax.pure_callback instead of jax.experimental.host_callback.call. This allows vmap over fmmax calculations. Adds a test that exercises vmap.
jax.pure_callback
jax.experimental.host_callback.call
Wow was just looking at this today, thanks! This should better enable pmap and xmap too, right?
pmap
xmap
Haha coincidence. Yeah I expect these should work seamlessly now.
Update the eigensolve to use
jax.pure_callback
instead ofjax.experimental.host_callback.call
. This allows vmap over fmmax calculations. Adds a test that exercises vmap.