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 via node = 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 inputs
print_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")

FAST_{RUN.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()

Links to this note

Aesara