Extracting connectivity graph from FLAX or JAX

[mattiasmar]

Hi,
Looking for a way to extract the computational graph from Flax (or JAX). More precisely I want to gather the connectivity graph and the dimensions of the tensors communicated between operators. Any ideas / pointers?

[jheek]

I think what you are looking for is jax.make_jaxpr this gives you all the jax operations. It has a list of operations eqns each equation has invars/outvars and a var has a .aval field which will typically have a shape so you can do things like:

jaxpr = jax.make_jaxpr(fn)(*inputs)
inputs = jaxpr.eqns[i].invars
print([x.aval.shape for x in inputs])

[mattiasmar]

@marcvanzee I created a reduced copy of the original Colab here demonstrating the error when calling to_dot.

[marcvanzee]

The error states that the variable collection batch_stats should be mutable. However, by default only the params collection is mutable, so you should add that as a keyword argument to the apply functions. Since to_dot wraps model.apply inside a function it doesn't seem possible to pass kwargs to it, but the following works:

dot.to_dot(functools.partial(model.apply, mutable=['batch_stats']))(
    {'params': state.params, 'batch_stats': state.batch_stats},
    batch['image'][:128])

This outputs something that I suppose you should pass to graphviz.Source(out).

[mattiasmar]

Thanks @marcvanzee. Your code change makes a lot of sense. I had a missed to educate myself on that parameter.

I tried the notebook in the TPU Collab but it caused the whole runtime to crash and restart. The logs says only:
2022-01-04 13:13:25.808480: W tensorflow/core/framework/cpu_allocator_impl.cc:82] Allocation of 308281344 exceeds 10% of free system memory.

Locally though, on my own GPU system, with the code change I now get a different error than before (which is a good thing! )

RuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:
%cudnn-conv = (f32[32,112,112,64]{2,1,3,0}, u8[0]{0}) custom-call(f32[32,224,224,3]{2,1,3,0} %copy.3, f32[7,7,3,64]{1,0,2,3} %copy.4), window={size=7x7 stride=2x2 pad=3_3x3_3}, dim_labels=b01f_01io->b01f, custom_call_target="__cudnn$convForward", metadata={op_type="conv_general_dilated" op_name="jit(conv_general_dilated)/conv_general_dilated[\n  batch_group_count=1\n  dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 3, 1, 2), rhs_spec=(3, 2, 0, 1), out_spec=(0, 3, 1, 2))\n  feature_group_count=1\n  lhs_dilation=(1, 1)\n  lhs_shape=(32, 224, 224, 3)\n  padding=((3, 3), (3, 3))\n  precision=None\n  preferred_element_type=None\n  rhs_dilation=(1, 1)\n  rhs_shape=(7, 7, 3, 64)\n  window_strides=(2, 2)\n]" source_file="/usr/local/lib/python3.8/dist-packages/haiku/_src/dot.py" source_line=204}, backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"

Original error: UNIMPLEMENTED: DNN library is not found.

tf-docker ~ > echo $LD_LIBRARY_PATH
/usr/local/cuda-11.4/lib64:/usr/local/cuda-11.0/targets/x86_64-linux/lib:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/cuda/lib64:/usr/local/nvidia/lib:/usr/local/nvidia/lib64
tf-docker ~ > echo $PATH
/usr/local/cuda-11.4/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin

Is your guess that this error is related to JAX, Flax or Haiku? From other git issues with this error, setting the LD_LIBRARY_PATH has been suggested as a solution, but it doesn't seem to have done the trick for me.

[marcvanzee]

As the error says, the image is too big to fit in memory. Why don't you start with a simple model, like this?

from flax import linen as nn
from jax import random
import jax.numpy as jnp
from haiku._src import dot
import graphviz

class Foo(nn.Module):
  @nn.compact
  def __call__(self, x):
    for _ in range(3):
      x = nn.Dense(4)(x)
    return x

inputs = jnp.zeros((3, 4))
variables = Foo().init(random.PRNGKey(0), inputs)
graphviz.Source(dot.to_dot(Foo().apply)(variables, inputs))

Outputs:

[mattiasmar]

Excellent. I got it to work now, also for the ImageNet example. Thanks so much for your help.
One of course actually doesn't actually need the accelerator (TPU/GPU) at all to generate the graph and so it is easy to avoid the OOM error.
I created a minimalistic Collab for the ResNet example. Perhaps that could help someone in the future.

[avital]

Is it possible to visualize the computation graph using tensorboard, as

both jax/flax and tensorflow are from google, I suppose it will be much easier to solve the incompatible part. Actually, I believe I recall someone using jax2tf (which includes Flax metadata because of Flax's use of named_call, which become tf.name_scopes) and loading a graph into the TensorBoard visualization that looked pretty good. If you try this out and it works, please report back! (and we can add a documentation page about it)

…

On Wed, Dec 15, 2021 at 10:39 AM jheek ***@***.***> wrote: Tensorboard doesn't support a XLA computation graph it only shows a tf graph. The haiku visualization module should work on any JAX function. And I think it will also support the Flax metadata (so it will groups ops belonging to the same flax Module). — You are receiving this because you commented. Reply to this email directly, view it on GitHub <#1718 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAAJFURW3ULOZ2NWOXHVLF3URBO3RANCNFSM5J4OHXFA> . Triage notifications on the go with GitHub Mobile for iOS <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675> or Android <https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub>.

[mattiasmar]

I tried to catch up on the jax2tf suggestion and applied it on the /flax/examples/imagenet/imagenet.ipynb example. Code pasted below (used after cell number 16 in the notebook).
In order to get the graph from TF I need to run the model once. I could use a little help there. Last line below fails. Any help appreciated.

import sonnet as snt
from jax import  random
import tree
from jax.experimental import jax2tf
def create_variable(path, value):
    name = '/'.join(map(str, path)).replace('~', '_')
    return tf.Variable(value, name=name)

class JaxModule(snt.Module):
    def __init__(self, params, apply_fn, name=None):
        super().__init__(name=name)
        self._params = tree.map_structure_with_path(create_variable, params)
        self._apply = jax2tf.convert(lambda p, x: apply_fn(p, None, x))
        self._apply = tf.autograph.experimental.do_not_convert(self._apply)

    def __call__(self, inputs):
        return self._apply(self._params, inputs)

model_cls = getattr(train.models, config.model)
model = train.create_model( model_cls=model_cls, half_precision=config.half_precision)
variables = model.init(random.PRNGKey(0), train_batch['image'])
state, params = variables.pop('params')
net = JaxModule(params, model.apply)
batch = next(iter(train_ds))
net(batch )  #This line fails