Question 1

How do I add a custom CUDA kernel to JAX?

Accepted Answer

This tutorial shows writing a CUDA kernel, exposing it via pybind11 as an XLA custom call, and registering it with JAX's primitive system. Key steps include handling the 'opaque' parameter for data serialization and launching kernels with proper grid/block dimensions.

Question 2

What's the performance benefit of custom JAX ops vs pure Python implementation?

Accepted Answer

Custom ops can drastically speed up iterative algorithms or special functions by leveraging optimized C++/CUDA, as highlighted in the cost/benefit analysis. However, for many tasks, JAX's JIT compiler is sufficient, so benchmark first to avoid unnecessary complexity.

Question 3

How to handle multiple outputs in a JAX custom operation?

Accepted Answer

The project demonstrates this using XLA's tuple output mechanism, where the 'out' parameter is treated as a void** for multiple buffers. Code examples in cpu_ops.cc and kernels.cc.cu show how to parse and assign outputs for both CPU and GPU.

Question 4

JAX custom op vs using jax.grad for differentiation, which is better?

Accepted Answer

Custom ops require manual JVP rule definitions, essential for ops not composable from existing JAX functions, like Kepler's equation. jax.grad works automatically but may be inefficient for low-level code, so custom ops offer control at the cost of added development.

Question 5

How to build and package a JAX extension with CMake?

Accepted Answer

The tutorial provides a full CMakeLists.txt that integrates with scikit-build-core, handling both CPU and optional GPU builds. It shows compiling C++/CUDA into a Python module and registering custom call targets via pybind11 capsules.

Question 6

Is this tutorial still relevant with JAX's official C++ support?

Accepted Answer

While JAX now has a supported interface, this project offers a detailed, working example that can be easier to follow for specific use cases. The README recommends checking the official docs first, but this repo serves as a practical reference for deeper understanding.

Extending JAX with custom C++ and CUDA code by Dan Foreman-Mackey

What is Extending JAX with custom C++ and CUDA code by Dan Foreman-Mackey?

Overview

Use Cases

Best For

Related Projects

Found a gem we're missing?

Not Ideal For

Pros & Cons

Pros

Cons

Frequently Asked Questions