Home / JAX / Extending JAX with custom C++ and CUDA code Extending JAX with custom C++ and CUDA code Tutorial demonstrating the infrastructure required to provide custom ops in JAX. Package 403 stars GitHub Back to JAX