visu-hlo 0.1

Displays the HLO representation of (un-)jitted functions as SVG

visu-hlo

visu-hlo displays the HLO representation of JAX functions as SVG visualizations.

Quick Example

import jax.numpy as jnp
from jax import jit
from visu_hlo import show

def func(x):
    return 3 * x * 2

# Visualize original function
show(func, jnp.ones(10))

# Visualize jitted function
show(jit(func), jnp.ones(10))

Installation

pip install visu-hlo

System dependency: Install Graphviz

Features

• 🎯 Easy Visualization: Display HLO graphs with a single function call
• ⚡ JIT Support: Works with both regular and jitted JAX functions
• 🖼️ SVG Output: High-quality vector graphics that scale perfectly
• 🖥️ Cross-Platform: Supports Linux, macOS, and Windows
• 📦 Lightweight: Minimal dependencies - just JAX and Graphviz

Documentation

Full documentation: https://visu-hlo.readthedocs.io/

License

MIT