jaxify 0.0.3

Write Python. Run JAX

jaxify

Write Python. Run JAX.

⚠️ jaxify is an experimental project under development
Feel free to test out and report any issues. Do not use in production.

---

jaxify lets you apply JAX transformations (like @jax.jit and/or @jax.vmap) to functions with Python control flow that JAX normally cannot compile, like if/elif/else statements depending on input values.

Installation

pip install jaxify

Getting started

import jax
import jax.numpy as jnp
from jaxify import jaxify

@jax.jit
@jax.vmap
@jaxify  # <-- Just add a @jaxify decorator
def absolute_value(x):
    if x >= 0:  # <-- If conditional in a JIT-compiled function!
        return x
    else:
        return -x

xs = jnp.arange(-1000, 1000)
ys = absolute_value(xs)  # <-- Runs at JAX speed!
print(ys)

How it works

@jaxify is a decorator that transforms Python functions by rewriting their abstract syntax tree (AST) to replace unsupported control flow constructs with JAX-compatible alternatives. It currently supports if/elif/else statements depending on input values, allowing you to write more natural Python code while still benefiting from JAX's performance boost.

When you decorate a function with @jaxify, it analyzes the function's source code, identifies control flow constructs, and rewrites them to use JAX's functional control flow primitives (like jax.lax.cond). The transformed function is then traceable by JAX, enabling you to apply JAX transformations like @jax.jit and @jax.vmap seamlessly.

Compatibility status

The following Python control flow constructs are currently supported within @jaxify-decorated functions:

Python construct	Support status	Notes
if / elif / else	✅	Should mostly work
if-else expressions	⚠️	Static values only
and / or	⚠️	Static values only. For dynamic values, use & or jnp.logical_and / | or jnp.logical_or instead
for loops	❌	Use jax.lax.fori_loop, jax.lax.scan, or jax.lax.while_loop instead
while loops	❌	Use jax.lax.while_loop instead
match-case	⚠️	Static values only. For dynamic values, use an if-elif-else chain or jax.lax.switch instead