> For clean Markdown content of this page, append .md to this URL. For the complete documentation index, see https://docs.nvidia.com/jax-toolbox/llms.txt

# Customize JAX with GPU kernels

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.

- [**Writing High-Performance CuTe DSL kernels in JAX**](https://docs.jax.dev/en/latest/notebooks/cute_dsl_jax.html)