Metadata-Version: 2.4
Name: visu-hlo
Version: 1.0.5
Summary: Displays the HLO representation of (un-)jitted JAX functions as SVG
Project-URL: Documentation, https://visu-hlo.readthedocs.io
Project-URL: Repository, https://github.com/CMBSciPol/visu-hlo
Project-URL: Issues, https://github.com/CMBSciPol/visu-hlo/issues
Author-email: Pierre Chanial <chanial@apc.in2p3.fr>
License-Expression: MIT
License-File: LICENSE
Classifier: Intended Audience :: Education
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Topic :: Scientific/Engineering :: Visualization
Classifier: Topic :: Software Development :: Debuggers
Requires-Python: >=3.10
Requires-Dist: graphviz>=0.18.2
Requires-Dist: jax>=0.4.0
Description-Content-Type: text/markdown

<div align="center">
<img src="https://raw.githubusercontent.com/CMBSciPol/visu-hlo/main/docs/source/_static/logo.svg" alt="VISU-HLO" width="100%">
</div>


# Visualizing HLO computation graphs

[![PyPI version](https://img.shields.io/pypi/v/visu-hlo?color=%232EBF4F)](https://pypi.org/project/visu-hlo)
[![Python versions](https://img.shields.io/pypi/pyversions/visu-hlo.svg?color=%232EBF4F)](https://pypi.org/project/visu-hlo)
[![Continuous integration](https://github.com/CMBSciPol/visu-hlo/actions/workflows/ci.yml/badge.svg)](https://github.com/CMBSciPol/visu-hlo/actions/workflows/ci.yml)
[![Coverage](https://codecov.io/gh/CMBSciPol/visu-hlo/graph/badge.svg)](https://codecov.io/gh/CMBSciPol/visu-hlo)
[![Documentation](https://readthedocs.org/projects/visu-hlo/badge/?version=latest)](https://visu-hlo.readthedocs.io/en/latest/?badge=latest)

<!-- Start common text with source/index.md -->

**visu-hlo** is a Python package that displays the HLO (High Level Operations) representation of JAX functions as SVG visualizations. This tool helps developers understand the computational graphs generated by XLA compilations.

## 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


## Quick Example

```python
import jax.numpy as jnp
from visu_hlo import show

# Display optimized HLO (default)
show(lambda x: 3 * x * 2, jnp.ones(10))
```

<div align="center">
<img src="https://raw.githubusercontent.com/CMBSciPol/visu-hlo/main/docs/source/_static/hlo_optimized.svg" alt="Optimized HLO" width="50%">
</div>

To display the non-optimized HLO:
```python
show(lambda x: 3 * x * 2, jnp.ones(10), jit=False)
```

<div align="center">
<img src="https://raw.githubusercontent.com/CMBSciPol/visu-hlo/main/docs/source/_static/hlo_not_optimized.svg" alt="Non-optimized HLO" width="50%">
</div>

To save as an SVG file:
```python
from visu_hlo import write_svg
write_svg('graph.svg', func, jnp.ones(10))
```

## Installation

```bash
pip install visu-hlo
```

**System dependency:** Install [Graphviz](https://graphviz.org/download/)

<!-- End common text with source/index.md -->

## Documentation

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

## License

MIT
