Skip to main content

Jaxpr Visualisation Tool

Project description

Jaxpr-Viz

JAX Computation Graph Visualisation Tool

JAX has built-in functionality to visualise the HLO graph generated by JAX, but I've found this rather low-level for some use-cases.

The intention of this package is to visualise how sub-functions are connected in JAX programs. It does this by converting the JaxPr representation into a pydot graph. See here for examples.

NOTE: This project is still at an early stage and may not support all JAX functionality (or permutations thereof). If you spot some strange behaviour please create a Github issue.

Installation

Install with pip:

pip install jpviz

Dependent on your system you may also need to install Graphviz

Usage

Jaxpr-viz can be used to visualise jit compiled (and nested) functions. It wraps jit compiled functions, which when called with concrete values returns a pydot graph.

For example this simple computation graph

import jax
import jax.numpy as jnp

import jpviz

@jax.jit
def foo(x):
    return 2 * x

@jax.jit
def bar(x):
    x = foo(x)
    return x - 1

# Wrap function and call with concrete arguments
#  here dot_graph is a pydot object
dot_graph = jpviz.draw(bar)(jnp.arange(10))
# This renders the graph to a png file
dot_graph.write_png("computation_graph.png")

produces this image

bar computation graph

Pydot has a number of options for rendering graphs, see here.

NOTE: For sub-functions to show as nodes/sub-graphs they need to be marked with @jax.jit, otherwise they will just merged into thir parent graph.

Jupyter Notebook

To show the rendered graph in a jupyter notebook you can use the helper function view_pydot

...
dot_graph = jpviz.draw(bar)(jnp.arange(10))
jpviz.view_pydot(dot)

Visualisation Options

Collapse Nodes

By default, functions that are composed of only primitive functions are collapsed into a single node (like foo in the above example). The full computation graph can be rendered using the collapse_primitives flag, setting it to False in the above example

...
dot_graph = jpviz.draw(bar, collapse_primitives=False)(jnp.arange(10))
...

produces

bar computation graph

Show Types

By default, type information is included in the node labels, this can be hidden using the show_avals flag, setting it to False

...
dot_graph = jpviz.draw(bar, show_avals=False)(jnp.arange(10))
...

produces

bar computation graph

NOTE: The labels of the nodes don't currently correspond to argument/variable names in the original Python code. Since JAX unpacks arguments/outputs to tuples they do correspond to the positioning of arguments and outputs.

Examples

See here for more examples of rendered computation graphs.

Developers

Developer notes can be found here.

Project details


Download files

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

Source Distribution

jpviz-0.1.6.tar.gz (10.2 kB view details)

Uploaded Source

Built Distribution

jpviz-0.1.6-py3-none-any.whl (10.9 kB view details)

Uploaded Python 3

File details

Details for the file jpviz-0.1.6.tar.gz.

File metadata

  • Download URL: jpviz-0.1.6.tar.gz
  • Upload date:
  • Size: 10.2 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: poetry/1.5.0 CPython/3.10.6 Linux/6.4.6-76060406-generic

File hashes

Hashes for jpviz-0.1.6.tar.gz
Algorithm Hash digest
SHA256 9712e1cc376e8fe1969b6a70d7765737d1b80cb61a5b0429d3492a67922f99c8
MD5 28a7ce73be7177c22894eeed57c62586
BLAKE2b-256 155ea0762397841393d274a479c9466aa113cf0107fa35180f0bd82ec87cd5f3

See more details on using hashes here.

File details

Details for the file jpviz-0.1.6-py3-none-any.whl.

File metadata

  • Download URL: jpviz-0.1.6-py3-none-any.whl
  • Upload date:
  • Size: 10.9 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: poetry/1.5.0 CPython/3.10.6 Linux/6.4.6-76060406-generic

File hashes

Hashes for jpviz-0.1.6-py3-none-any.whl
Algorithm Hash digest
SHA256 d438b40c90e2d61799f08a3631ac7ca2b729a4f551dc4b4988d488e68ebcbeff
MD5 05d0cc8c5849b91dbbe2113f87a4758f
BLAKE2b-256 7c1fbd9e59c687c34a51fbc52f4a7ac10687eed64b530487c1180f506608ed54

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page