jaxify 0.0.1

Write Python. Run JAX

jaxify

Write Python. Run JAX.

---

⚠️ jaxify is an experimental project under development
Right now, only some if statements may work. Use at your own risk.

Installation

pip install jaxify

Getting started

import jax.numpy as jnp
from jaxify import jitx

@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)