Eli Bendersky: Decorator JITs - Python as a DSL

Spend enough time looking at Python programs and packages for machine learning, and you'll notice that the "JIT decorator" pattern is pretty popular. For example, this JAX snippet:

Or the Triton language for writing GPU kernels directly in Python:

In both cases, the function decorated with jit doesn't get executed by the Python interpreter in the normal sense. Instead, the code inside is more like a DSL (Domain Specific Language) processed by a special purpose compiler built into the library (JAX or Triton). Another way to think about it is that Python is used as a meta language to describe computations.

In this post I will describe some implementation strategies used by libraries to make this possible.

https://eli.thegreenplace.net/2025/decorator-jits-python-as-a-dsl/