aesara
aesara copied to clipboard
Folding graph representation with Textual
We may want to investigate using Textual or rich to pretty-print aesara graphs. This would make the aesara.dprint output interactive.
With textual, an option may be to have the graphs fold and/or offer filtering capability for different node types.
To illustrate, we're talking about making the following kinds of output interactive:
import aesara
import aesara.tensor as at
x = at.vector("x")
y = at.matrix("y")
z = x + y
aesara.dprint(z, depth=2)
# Elemwise{add,no_inplace} [id A] ''
# |InplaceDimShuffle{x,0} [id B] ''
# |y [id C]
For example, if a user selected/clicked the B node in this graph print-out, they would get the following expanded node:
aesara.dprint(z)
# Elemwise{add,no_inplace} [id A] ''
# |InplaceDimShuffle{x,0} [id B] ''
# | |x [id C]
# |y [id D]
I looked into it, and we want to wait at least until the css branch is merged on their side: https://github.com/Textualize/textual as many things are about to change.
I played around with this a little:

The features demonstrated there involve:
- Non-breaking, mid-session usage. In other words, it can be used as a debugging tool without ruining the debug session.
- Basic graph display and folding.
- Identical sub-graphs are only displayed once and closest to the "root" of the graph/tree. Unfolding non-root copies will jump to the their roots.
- Node search and selection. Selecting from the search results jumps to the corresponding node in the graph.
A few features that seem easy to add:
- Narrow by sub-graph, and the ability to re-expand/undo the narrowing
- Traversable jump history
- Interactive graph comparisons/diffs
It would be great to get this working with Emacs' comint mode, so that the graph explorer could be used within interactive debugging sessions. This is really the only thing preventing me from considering it as a real day-to-day debug tool.
It would be nice to be able to jump to the next occurence the current subgraph using a shortcut like n. Other than that I think it's great and is already worth being added to Aesara!