zombie-einstein / jaxpr-viz Goto Github PK

Hello!

I'm just trying to run the quickstart but I'm having some issues:

My setup:

• python 3.9.16
• jax==0.4.14
• jaxlib==0.4.14
• I installed the package with: pip install jpviz

scripts

I created a script with the following (from the quickstart)

import jax
import jax.numpy as jnp

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

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

import jpviz

jpviz.draw(bar)(jnp.arange(10))

But no output is created.

jupyter notebook

I tried the same in a notebook, but following the readme I ran

dot = jpviz.draw(bar, collapse_primitives=True)(jnp.arange(10))
jpviz.view_pydot(dot)

Here I get the error FileNotFoundError: [Errno 2] "dot" not found in path.

Thanks!