jaxify

Write Python. Run JAX.

⚠️ jaxify is an experimental project under development
You're welcome to try it out and report any issues!

jaxify lets you apply JAX transformations (like @jax.jit and/or @jax.vmap) to functions with common Python constructs that JAX cannot itself handle, such as if conditions that depend 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 decorate your function with @jaxify
def absolute_value(x):
    if x >= 0:  # <-- If block 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

The @jaxify decorator transforms Python functions using static analysis to replace unsupported Python constructs with JAX-compatible alternatives. After the transformations, the functions become traceable by JAX, enabling you to apply functional JAX transformations like @jax.jit and @jax.vmap in a seamless manner.

Compatibility status

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

🔀 Conditionals

Construct	Works?	Notes
`if` statements	✅	Fully supported including `elif` and `else` clauses. Translated to calls to `jax.lax.cond`
`if` expressions (e.g. `a if b else c`)	✅	Translated to `jax.lax.cond`

⚖️ Comparisons

Construct	Works?	Notes
`==`, `!=`, `<`, `>`, `<=`, `>=`	✅	Chained comparisons (e.g. `x < y <= z`) are supported by translation to the equivalent chain of individual comparisons

1️⃣ Logical operators

Construct	Works?	Notes
`and` / `or`	✅	Short-circuiting of traced values supported via translation to `jax.lax.cond` calls
`not`	✅	Translates to `jnp.logical_not` for traced single values

🔄 Loops

Construct	Works?	Notes
`for` loops	❌	Currently unsupported. Use `jax.lax.fori_loop`, `jax.lax.scan`, or `jax.lax.while_loop` instead
`while` loops	❌	Currently unsupported. Use `jax.lax.while_loop` instead

🎯 Pattern matching

Construct	Works?	Notes
`match`-`case`	✅⚠️	Static values only. For traced values, use an `if`-`elif`-`else` chain or `jax.lax.switch` instead

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

jaxify-0.0.5.tar.gz (6.0 kB view details)

Uploaded Aug 5, 2026 Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

jaxify-0.0.5-py3-none-any.whl (7.0 kB view details)

Uploaded Aug 5, 2026 Python 3

File details

Details for the file jaxify-0.0.5.tar.gz.

File metadata

Download URL: jaxify-0.0.5.tar.gz
Upload date: Aug 5, 2026
Size: 6.0 kB
Tags: Source
Uploaded using Trusted Publishing? Yes
Uploaded via: uv/0.12.1 {"installer":{"name":"uv","version":"0.12.1","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}

File hashes

Hashes for jaxify-0.0.5.tar.gz
Algorithm	Hash digest
SHA256	`d5ef08e7db71b2095913e6794acbe17c342710b38e7998003d5dd0403b309c7b`
MD5	`01ba4703a7d65ec7d8a8b3ff9e5b4db9`
BLAKE2b-256	`cb154865281715fedd7b5234f4ee127fd3b0bc2086d7d5b3833748aa2721dd69`

See more details on using hashes here.

File details

Details for the file jaxify-0.0.5-py3-none-any.whl.

File metadata

Download URL: jaxify-0.0.5-py3-none-any.whl
Upload date: Aug 5, 2026
Size: 7.0 kB
Tags: Python 3
Uploaded using Trusted Publishing? Yes
Uploaded via: uv/0.12.1 {"installer":{"name":"uv","version":"0.12.1","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}

File hashes

Hashes for jaxify-0.0.5-py3-none-any.whl
Algorithm	Hash digest
SHA256	`2108f701d66a705bbc4ca36584e4552c91c216f2d73a8c01a7df9f96330a8922`
MD5	`452c833ea19f39c91dcd8ac63b1c296f`
BLAKE2b-256	`f3b1d6a2a465ce58135da6e04ec8fd3249a81ee78f486e08ce5e658b7fbc811c`