Customize JAX with GPU kernels
Writing custom kernels for JAX on GPUs.
Developers need customization of the JAX stack to customize their models and to optimize beyond what’s natively available in JAX/XLA stack.
NVIDIA offers a suite of tools, libraries and kernel DSLs, helping customization.