jaxify 0.0.2

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 JIT-compile functions (using JAX) that @jax.jit cannot handle. With jaxify, you can compile functions with e.g. Python if/elif/else statements (with support for other control flow structures planned for the future) that might be affected by the values of inputs.

jaxify's @jitx decorator works exclusively on the decorated function and intervenes only at tracing/compilation time; it does not have any effect at actual runtime besides the code it emits for JAX.

Installation

pip install jaxify

Getting started

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

@jax.vmap
@jitx
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)