Debug an Aesara graph (advanced)
The graph compiles, I get an ugly stack trace!
E IndexError: index out of bounds E Apply node that caused the error: Subtensor{int64, ::, int64}(BroadcastTo.0, ScalarConstant{0}, ScalarConstant{39}) E Toposort index: 492 E Inputs types: [TensorType(float64, (None, None, None)), ScalarType(int64), ScalarType(int64)] E Inputs shapes: [(2401, 2, 2), (), ()] E Inputs strides: [(32, 16, 8), (), ()] E Inputs values: ['not shown', 0, 39] E Outputs clients: [[InplaceDimShuffle{x,0}(Subtensor{int64, ::, int64}.0)]]
- Re-run in a debugger
- Go up (
u
) the stack trace fgraph
(the function graph where the error occurs) should be avaible in the name space.- In the error message you should see indications such as
Toposort index: 492
, and this means that you can access the node vianode = fgraph.toposort()[492]
- You can then inspect the resulting node, and see who created it by inspecting the tag
node.tag
.
Clear the cache
Sometimes it can be useful to clear the cache manually between different runs
aesara-cache clear
Pretty printing the graph
Aeasara can return at a pretty-printed version of the graph:
aesara.dprint(fgraph)
This dprint
function has a couple of interesting keyword arguments:
print_op_info=True
Gives annotation next to the Scan inputsprint_type_info=True
Rewrites applied to the graph
We can change the Logging properties using a context manager, fun.
with aesara.config.(rewrite_verbose=True): Compile
Disable optimizations
After examining the graph, and how the nodes were created you may suspect that a bug was introduced by a rewrite. The best way to confirm this hypothesis is usually to disable all rewrites at once. And then exclude individual rewrites to narrow it down.
mode=Mode(optimizer="None")
FASTRUN.excluding("<name>")
If this is the inner-graph of a Scan
Inner-graphs in Scan
need to ave their set of optimizations set separately. It is then easier to use the global aesara.config
. This is better done in conftest.py
(the config needs to be updated before Aesara is initialized).
import os def pytest_sessionstart(session): os.environ["AESARA_FLAGS"] = ",".join( [ os.environ.setdefault("AESARA_FLAGS", ""), "warn__ignore_bug_before=all,on_opt_error=raise,on_shape_error=raise,optimizer_excluding=remove_constants_and_unused_inputs", ] )
We're looking for the optimizer_excluding
flag.
TODO Rich for graph folding project
TODO Add rewrite trajectory to a node's tag
in the fully rewritten graph
Funny, isn't that something that we could easily deduct if the rewrites were described in an e-graph?
Profile Aesara code
# Enable profiling for compilation of the outer-function f_doubling = aesara.function([A, B, tol, max_iter], gamma[-1], profile=True) f_doubling.profile.summary()