jaxplorer

A compiler explorer TUI for JAX. Put a jitted function on the left, watch its jaxpr, StableHLO and optimized HLO on the right, and see all three change as you type.

┌─ source ──────────────────┬─ Jaxpr │ StableHLO │ Optimized HLO │ Analysis │ Errors ─┐
│ def f(x, w):              │ { lambda ; a:f32[8,16] b:f32[16,4]. let                 │
│     return jnp.tanh(x @ w)│     c:f32[8,4] = dot_general[...] a b                   │
│                           │     d:f32[8,4] = tanh c                                 │
└───────────────────────────┴────────────────────────────────────────────────────────┘
 cpu · jax 0.11.0 · 134 ms · ok

Installation

uv sync                             # in a checkout
uv tool install "jaxplorer[jax]"    # standalone, with its own jax
uv tool install jaxplorer           # standalone, borrowing a project's jax (see below)

Needs Python 3.12+ and textual >= 6.0. jax >= 0.9 is an extra, since the jax worth inspecting is usually a project's own. Any CPU-only jax install is enough; a GPU or TPU backend is only needed to compile for one.

Usage

uv run jaxplorer                      # scratch buffer
uv run jaxplorer examples/mlp.py      # open a snippet and edit it in place
uv run jaxplorer examples/mlp.py --watch   # keep editing in your own editor; jaxplorer reloads on save

Click any HLO instruction to select the source line that produced it.

Other options: --python PATH (compile under another environment's jax, see below), --platform cpu|gpu|tpu, --x64, --timeout SECONDS, --stages jaxpr,stablehlo,... (stopping before optimized_hlo skips XLA, which is much faster on a large model), and --passes to collect per-pass HLO and LLVM IR from the start.

Against your own project's jax

jaxplorer compiles in a subprocess, and that subprocess can be your project's interpreter:

cd my-project
uv run jaxplorer model.py                     # lookup via uv (VIRTUAL_ENV)
jaxplorer model.py --python .venv/bin/python   # or name the interpreter outright.

The interpreter is chosen in this order: --python, then $VIRTUAL_ENV, then the one running jaxplorer. uv run and a plain activate both export VIRTUAL_ENV, so inside a project the flag is rarely needed.

Anatomy of a snippet

A snippet is an ordinary Python module that defines a callable f and a tuple args of example inputs. Nothing needs to import jaxplorer.

import jax
import jax.numpy as jnp


def f(x, w):
    return jnp.tanh(x @ w).sum()


args = (
    jax.ShapeDtypeStruct((8, 16), jnp.float32),
    jax.ShapeDtypeStruct((16, 4), jnp.float32),
)

args may hold concrete arrays or jax.ShapeDtypeStruct specs. jaxplorer only traces and compiles f, never runs it, so shape specs are enough. Optionally define kwargs, static_argnums, static_argnames or donate_argnums; they are passed to jax.jit. See examples/ for an MLP, a lax.scan loop, and causal attention with a static argument.

How it works

Each pane is one step of JAX's public lowering chain:

Stages are reported independently, so a lowering failure still leaves you a valid jaxpr to read, and a buffer that does not even parse keeps the last IR that did compile on screen.

Compilation happens in a subprocess (python -m jaxplorer.worker) that stays warm between edits: JAX takes seconds to boot, XLA can abort the process outright, and the platform and x64 flags can only be set before JAX is imported.

For the rest of the pipeline (per-pass HLO dumps, LLVM IR, object code, and comparing two XLA builds against each other) see docs/xla-introspection.md.

jaxplorer executes the buffer. It is your own code in your own environment, but a snippet is run at module level on every recompile, so treat it the way you would treat python snippet.py.

Development

Contributions welcome. Here's the general testing and formatting workflow of the repo:

uv run --group test pytest
uvx prek run --all-files
uv run --group typing ty check