Advice or tips on improving compilation for large(!) O(100k) graph operations

[kratsg]

(note: happy to walk through this offline if I cannot explain super clearly here)

    dist_expression: TensorVar = model.log_prob # this is a very large computational graph

    from pytensor.compile.function import function
    from pytensor.graph.traversal import applys_between, explicit_graph_inputs
    inputs = [
        var
        for var in explicit_graph_inputs([dist_expression])
        if var.name is not None
    ]
    compilation_mode = model.mode

    applies = list(applys_between(inputs, [dist_expression]))
    op_types: dict[str, int] = {}
    for apply in applies:
        op_name = type(apply.op).__name__
        op_types[op_name] = op_types.get(op_name, 0) + 1

    print(f'Input variables: {len(inputs)}')
    print(f'Graph operations: {len(applies)}')
    print(f'Operation types: {dict(sorted(op_types.items()))}')

whic has:

Input variables: 184
Graph operations: 322159
Operation types: {'DimShuffle': 10965, 'Elemwise': 310452, 'Reshape': 364, 'Sum': 378}

Now, the log_prob is a likelihood which is the weighted sum of individual channels, that is, likelihood = [channel1, channel2, channel3, ..., cN], weights=[w1, ..., wN] is going to give us c1*w1 + c2*w2 + ... + cN*wN.

The log_prob above is this entire sum. We can also extract out those individual subgraphs, each of which takes ~1 minute to compile, but the total likelihood log prob takes over an hour to compile. Before I go down the route of throwing TONS of code, I'd like to start with the high-level thing first and understand if there's obvious things I can do to help debug or manage large graphs like these.

Note, by compilation, I mean

log_prob = function(
    inputs=inputs, 
    outputs=dist_expression, 
    mode=compilation_mode, 
    on_unused_input="ignore", 
    name=model._likelihood.name, 
    trust_input=True
)

[ricardoV94]

You need to write it in a vectorized manner like you would with numpy, not with 1000s of individual scalars or a python loop or whatever generated so many nodes.

[kratsg]

There are no for-loops. We truly do have that many pieces in these larger models that we do build that are composed together (and this is just one of the dozens of statistical workspaces we would eventually combine into a single one with many parameters correlated across them all):

>>> Counter([dist.type for dist in ws.distributions])
Counter({'crystalball_doublesided_dist': 350, 'gaussian_dist': 148, 'product_dist': 14, 'mixture_dist': 14, 'generic_dist': 14})
>>> Counter([func.type for func in ws.functions])
Counter({'product': 9709, 'interpolation': 420, 'generic_function': 28})

and this is already after constant-folding in ~300,000 individual parameters down to about 170 free parameters. All of these are typically added or multiplied together, in addition to having lots of individual distributions being normalized too. There could be graph rewrites that I'm not seeing, but I'm trying to understand how to identify low-hanging fruits for optimizing this. Would profiling be suggested?

[ricardoV94]

My guess is you're starting from a bad place, and before thinking about rewrite/compile performance you should check you're not doing some xyproblem https://xyproblem.info/

What do you have and what do you need in the end? 170 distinct parameter variables sounds insane, usually we work with a single raveled vector (which can be much larger than 170 entries without issues) and split internally. But rarely do we need 170 splits. more like first 20 entries go into a vectorized logp density, next 20 to a halfnormal, then another prior variable and so on, not 170 completely distinct branches.

That's why pymc pushes strongly on shape/dims to express batches of same family priors.

[ricardoV94]

Insane is a strong word, but never saw a python function being called with 170 params, I guess they would have structure and be in a dict / class, which pytensor doesn't support ofc.

But the question is the same, where do they come from, where do they go, how hard did you go on redefining the graph to look vectorized from the start. PyTensor doesn't do stuff like we discussed in #1984, but really would love to start from a place like that.