Add support for `jax.lax.map`, `jax.pmap` in `kmap` and `smap`

ASEM000 / kernex

Stencil computations in JAX

MIT License

66 stars 3 forks source link

Add support for `jax.lax.map`, `jax.pmap` in `kmap` and `smap` #9

Closed ASEM000 closed 1 year ago

ASEM000 commented 1 year ago

Enable the kernel map backend. The default now is vmap, however it is beneficial to enable the user to choose among other variants like pmap and lax.map

ASEM000 commented 1 year ago

closed by #10