jap is a map engine for Google's JAX.
It maps python functions to
JAX arrays or
pytree,
effectively turning python functions into accelerator kernels.
In addition to unifying JAX's vmap(), pmap(), and
xmap() interfaces, it also makes implementations of, e.g.,
integrators of ordinary differential equations with adaptive time
steps, more straightforward.
-
Notifications
You must be signed in to change notification settings - Fork 0
Map engine for JAX
License
adxsrc/jap
Folders and files
| Name | Name | Last commit message | Last commit date | |
|---|---|---|---|---|
Repository files navigation
About
Map engine for JAX
Resources
License
Stars
Watchers
Forks
Releases
No releases published
Packages 0
No packages published