5 projects
jztree
A framework for Morton z-order trees for JAX with CUDA backend
jztree-cu12
CUDA 12 backend for jztree
jztree-cu13
CUDA 13 backend for jztree
jax-ffi-gen
Scripts to auto-generate jax's foreign function interface from CUDA kernels
pytest-jax-bench
Pytest plugin to profile jitted JAX functions (compile time, runtime, memory).
