Python's JIT Decorators: Three Implementation Strategies

This article delves into the popular JIT decorator pattern in Python, particularly its use in JAX and Triton libraries. The author implements three JIT decorators from scratch using a simplified example: AST-based, bytecode-based, and tracing-based. The AST-based approach directly manipulates the Abstract Syntax Tree; the bytecode-based approach leverages Python's bytecode interpreter; and the tracing-based approach builds an expression IR by tracing function execution at runtime. The article details the advantages and disadvantages of each approach and uses JAX and Numba as examples to illustrate their strategies in real-world applications.