jaxprvis 0.0.1

A visualizer for JaxPrs

JaxPr Visualizer

This project provides interactive visualizations of JAX computations. Composite nodes can be expanded using left-click and contracted using right-click. Space pauses the simulation.

Usage Example

import jax.numpy as jnp

from jax import make_jaxpr
from jax.core import ClosedJaxpr

import jaxprvis as jvis

def visualize_loop_unroll_jaxpr(iter: int) -> ClosedJaxpr:
    example_args = jnp.zeros((2, 2))

    def f(x, y):
        z = x
        for i in range(iter):
            z = (z + y) / x
        return z

    closed_jaxpr: ClosedJaxpr = make_jaxpr(f)(example_args, example_args)
    jvis.visualize(closed_jaxpr)

Installation

For Users

Install from Pypi:

#Create and activate virtual environment (if not present yet)
python -m venv --upgrade-deps ./venv
source ./venv/bin/activate
pip install --upgrade build
# Install your JAX version
#e.g. pip install --upgrade "jax[cpu]"

pip install --upgrade jaxprvis

Install from source:

# Clone
git clone https://github.com/PSilvestre/jaxprvis
cd jaxprvis

#Create and activate virtual environment
python -m venv --upgrade-deps ./venv
source ./venv/bin/activate
pip install --upgrade build

python -m build
pip install -e .

For Developers

Similar to the above, but instead install the dev requirements. Run the following to install:

# Clone
git clone https://github.com/PSilvestre/jaxprvis
cd CORAL

#Create and activate virtual environment
python -m venv --upgrade-deps ./venv
source ./venv/bin/activate
pip install --upgrade build

python -m build
pip install -e ".[dev]"

#Run the tests
pytest