How to visualize the computing graph of the mlx models?

[0x1orz]

How to visualize the computing graph of the mlx models? Is there a tool like tensorboard for mlx?

[angeloskath]

For any array one can export the computation graph to .dot format using mx.export_to_dot("filename.dot", x, y, z). Then the dot file can be made into a png or pdf using the dot command line tool. My goto is dot -Tpng filename.dot -o filename.png.

There are a couple of things to pay attention to:

• This API is kind of an internal one that we sometimes use for debugging. I don't think it is going anywhere any time soon but your mileage may vary.
• Evaluated arrays almost always have no computation graph. So you should call export to dot before calling eval.
• The graph is the whole graph. Not with respect to the input to the model. For instance if the weights have not been evaluated then the graph will include the weight initialization.

Here is an example with the corresponding png generated:

import mlx.core as mx
x = mx.random.normal(shape=(10, 10))
y = mx.random.normal(shape=(10, 10))
mx.eval(x, y)  # <--- important in order to not have the x, y initialization in the graph
z = ((x @ y).square() + 1).sum()
mx.export_to_dot("/tmp/example.dot", z)

then running dot -Tpng /tmp/example.dot -o /tmp/example.png produces