greedy_variance_init builds full N×N gram(s) just to take the kernel diagonal

[jessegrabowski]

inducing.py:602 computes the candidate diagonal as pt.diag(kernel(X_sym)), which materializes the full K(X, X) — for a K-factor ProductKernel, K separate N×N grams — before slicing the diagonal. Peak memory is O(K·N²) instead of O(N).

Current:

import pytensor
import pytensor.tensor as pt

from ptgp.kernels import ExpQuad

kernel = ExpQuad(1, 1.0) * ExpQuad(1, 1.0) * ExpQuad(1, 1.0)  # 3-factor product kernel
X = pt.matrix("X", shape=(None, 1))


def count_2d_nodes(fn):
    return sum(out.type.ndim == 2 for node in fn.maker.fgraph.toposort() for out in node.outputs)


f = pytensor.function([X], pt.diag(kernel(X)))  # as in inducing.py:602
print(count_2d_nodes(f))  # 5

Fix — use the existing O(N) kernel.diag:

g = pytensor.function([X], kernel.diag(X))
print(count_2d_nodes(g))  # 0

[jessegrabowski]

I also want to double check why this isn't being rewritten into a the more memory efficient form.

[jessegrabowski]

pymc-devs/pytensor#2302 gives a deeper fix, I think this is what we were actually being crushed by.

Correction: This is why we have a big memory blowup, but it wouldn't fix the issue by itself. We would need additional rewrites to push the ExtractDiag inside a kernels' elemwise Ops, which we currently don't have. I think this PR is valuable on its own.