Jax-PR Visualisation

[zombie-einstein]

I wanted to shamelessly promote this library I made to visualise generated JAX-PRs. Docs and examples can be found here on the github repo.

The intention was to be able to create a higher level visualisation of JAX programs in comparison to the built in HLO visualisation to allow examination of program structure.

I built it to support another project I'm working on, but interested to see if people find this useful, and any feedback.

[ASEM000]

This is really cool,

For your reference other projects with the same spirit (jaxpr-level viz)

• haiku jaxpr info
• jaxpr to latex - I think a bit old though -

It would be interesting if the user can define/register his own rules, for example general convolution can be visualized like this.

I think it might also be useful to define a viz primitive that carries some information about certain variable, for example

def foo(x):
  y = tag(x, metadata="this is the first step")
  return y

where tag itself is identity, but you can make use of the metada at the jaxpr level to output it to graphviz.

[zombie-einstein]

Thanks @ASEM000

It would be interesting if the user can define/register his own rules

Do you mean as in user defined styling of the generated graph?

I think it might also be useful to define a viz primitive that carries some information about certain variable

Does jaxpr have the ability to carry additional metadata, or would this be something would mean making a change to jax-core?

[ASEM000]

Sorry for late reply

Do you mean as in user defined styling of the generated graph?

Yes

Does jaxpr have the ability to carry additional metadata, or would this be something would mean making a change to jax-core?

AFAIK yes, you can check jaxpr to latex link for more.

Hope this helps.