JAX and CUDA are not two options on the same shelf. JAX is a Python-level array and automatic-differentiation framework compiled through XLA; CUDA is NVIDIA’s device-level compute API and toolchain. On NVIDIA hardware, a JAX program still lowers to CUDA kernels and vendor libraries such as cuBLAS and cuDNN underneath. Choosing JAX does not remove the NVIDIA dependency — it moves the question one layer down, where most teams stop looking.
That matters because the category error has a cost. Teams read “JAX vs CUDA” as a portability decision, adopt JAX, and record the lock-in question as closed. It is not closed. It has only been made harder to see.
Is JAX an alternative to CUDA, or does it run on top of it?
It runs on top of it, when the backend is NVIDIA. The useful framing separates two decisions that the phrase “JAX vs CUDA” collapses into one:
- Which framework expresses the workload — JAX, PyTorch, or something lower-level. This governs how you write models, how gradients are obtained, and which compiler sees your graph.
- Which compute API the backend targets — CUDA, SYCL, OpenCL, or a vendor runtime. This governs what silicon can execute the result.
The first decision is yours to make freely. The second is largely inherited from the hardware you buy, and it is the one that determines what a migration costs later.
Where the abstraction stops absorbing the problem
The framework holds up until one of two things happens. Either the workload needs performance outside XLA’s fusion patterns — an operator XLA will not fuse well, an unusual memory layout, a kernel you end up writing by hand — or the workload needs to run on a non-NVIDIA accelerator. Past either point, the framework abstraction stops absorbing the difference and the API-level choice becomes the real one.
Memory-layout and kernel assumptions baked in at the framework level do not port for free. That is the same compounding cost that CUDA-specific code carries, arriving by a less obvious route.
Quick answer: what is actually NVIDIA-bound in a JAX workload
| Layer | NVIDIA-bound? | What to check |
|---|---|---|
JAX Python code, jit, grad, vmap |
No | Portable in principle across backends |
| XLA-generated kernels | Backend-specific, generated | Re-generated per backend; performance not guaranteed to transfer |
| cuBLAS / cuDNN calls under XLA | Yes | Which ops dispatch to vendor libraries |
| Hand-written CUDA kernels and custom calls | Yes | Count them; each is a rewrite on a new backend |
| Profiling, tooling and CI assumptions | Often yes | Nsight-based workflows, container images, driver pinning |
Three numbers replace the unquantified assumption that JAX is portable: the share of runtime spent in XLA-generated kernels versus vendor-library or hand-written kernels, the number of custom kernels that would need rewriting on a non-NVIDIA backend, and the estimated engineering days for a backend switch. We ask for those three before treating any portability claim as decided.
Which of the two decisions is the binding constraint is measurable, not a matter of opinion — it depends on where runtime actually sits. The API-level trade-offs behind the second decision, including how CUDA-specific code accumulates over a multi-year hardware plan, are covered in our GPU engineering and acceleration work.
If your three-year plan assumes a non-NVIDIA option stays open, the honest test is not which framework you picked — it is how many kernels you would rewrite tomorrow, and whether anyone has counted them.
