ENH rewrite: vectorizing ops automatically (e.g. multiple distributions transformed to a single multivariate distribution)

[kratsg]

In a general sense, the following are roughly equivalent (with the latter two requiring a flattening to match the original dimensions):

Gauss(x1|m1,s1) * Gauss(x2|m2,s2) * Gauss(x3|m3,s3)
Gauss([x1,x2,x3], [m1,m2,m3], [s1,s2,s3]) # with broadcast semantics
MultiDimGauss([x1,x2,x3], [m1,m2,m3], [[s1,0,0],[0,s2,0],[0,0,s3]) # i.e. happens to have diagonal covariance

I intuitively imagine that the computational graph is more efficient with the latter two than the former, given the number of ops is reduced in favor of increasing the tensor dimensionality. However, it seems from a quick investigation (using Claude) scipp-atlas/pyhs3#18 that this is not the case.

I'm unsure of existing prior work. I'm told that this is something that Distributions.jl handles automatically.

[ricardoV94]

what do you mean with gauss() * gauss()? In general it depends. Concatenating the parameters requires copy. If the parameters came from a shared array (that is at some point split), then it could make sense to "delay" the split. Or if you control the inputs, but pytensor at compile time does not.

We can still have rewrites that are opt-in. They come in handy in different workflows. This is not specific to RVs, Any operation with batch semantics (Elemwise, Blockwise..., and where you can guarantee inputs would go nicely in a square tensor) could fit the mold.

[kratsg]

this would be like a joint probability density when writing gauss() * gauss(). although at least in pytensor-land, I suspect that writing Gauss() * Gauss() would need one to rewrite the Elemwise in such cases? I think in these cases, we have full control over the inputs so it should be true that the inputs are regular (or go into a square tensor).

[ricardoV94]

IIRC you're suggesting something like:

x = pt.vector("x")
y = pt.vector("y")

def gauss_density(x):
  return pt.exp(x)  # For example

x_density = gauss_density(x)
y_density = gauss_density(y)
joint_density = x_density * y_density

And you want to rewrite as:

xy = pt.stack([x, y], axis=0)
xy_density = gauss_density(x, y)
joint_density = xy_density.prod(axis=0)

What I meant is that pt.stack will require a copy of x, y into a common buffer, which may be costly.

It may also not always be possible to operate on the joint tensor the whole way. In this case we could do prod with an axis, but other cases may not have a clean solution.

[ricardoV94]

As a practical advice if you have an optimization idea, you should try to write equivalent numpy code that uses it vs does not use it and see if it actually comes ahead.

[ricardoV94]

import numpy as np

rng = np.random.default_rng(42)

def gauss_density(x):
    return np.exp(-0.5 * x**2)

n = 10_000
x = rng.standard_normal(n)
y = rng.standard_normal(n)

# Verify equivalence
xy = np.stack([x, y], axis=0)
assert np.allclose(gauss_density(x) * gauss_density(y), gauss_density(xy).prod(axis=0))

%timeit gauss_density(x) * gauss_density(y)  # 133 μs ± 4.1 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
%timeit gauss_density(np.stack([x, y], axis=0)).prod(axis=0)  # 160 μs ± 8.07 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

[kratsg]

This helps a lot - or at least gives me enough pointers that I will follow up myself on my end and understand. I wasn't thinking about the buffering of the inputs before but that can be costly (I agree) and might not be that optimal at the end of the day, or rather you lose out on potential benefits if it ends up being a wash.