KroneckerNormal using deprecated batched_dot

[ricardoV94]

pymc/pymc/distributions/multivariate.py

Line 2128 in dc7cfee

quad = pt.batched_dot(sqrt_quad.T, sqrt_quad.T)

I'm not sure if this is supposed to be a matmul or vecdot. I think the later. If someone is familiar with the distribution please weigh in.

CC @bwengals or perhaps @jessegrabowski (in case we can simplify this with PyTensor magic)

[jessegrabowski]

All of the kron machinery should be depreciated and replaced with linalg rewrites. All this machinery does is lift the eigh into kron:

eigh(kron(A, B))[0] -> outer(eigh(A)[0], eigh(B)[0]).ravel()
eigh(kron(A,B))[1] -> kron(eigh(A)[1], eigh(B)[1])) 

It's also doing kron(A, B) @ C -> kron(A @ C, B @ C), which we already do. So we're just missing the eigenvalue stuff, then we can eliminate all this.

As for the specific question, the quad line should be a an inner product (so you get back one scalar per batch dim). It's the (x - mu) ** 2 / sigma term in the logp.

The code has this extra logic to make sure we're always working with column vectors (in the one-sample case), which I think can also be done away with?

[ricardoV94]

Can we just pass a Kron covariance matrix to a regular MvNormal once we're done?

[jessegrabowski]

Yes, that's what I'm saying

[ricardoV94]

Maybe not. We also have diagonal sigma on top of the covariances, we would need to lift that through the Kronecker if that's also a thing.

[jessegrabowski]

import pytensor.tensor as pt
import pymc as pm

sigma = pt.dscalar('sigma')
mu, value = pt.dvectors('mu', 'value')
A, B = pt.dmatrices("A", "B")

cov = pt.linalg.kron(A, B) # + pt.eye(A.shape[0] * B.shape[0]) * sigma
lp = pm.logp(pm.MvNormal.dist(mu=mu, cov=cov), value)

Setting aside the sigma, here's what we currently rewrite down into:

Check{posdef covariance} [id A] 'value_logprob' 18
 ├─ Composite{(((-0.9189385332046727 * cast{float64}(i1)) - (0.5 * i0)) - i2)} [id B] 17
 │  ├─ Squeeze{axis=0} [id C] 16
 │  │  └─ CAReduce{Composite{(i0 + sqr(i1))}, axis=1} [id D] 15
 │  │     └─ Transpose{axes=[1, 0]} [id E] 14
 │  │        └─ SolveTriangular{unit_diagonal=False, lower=True, check_finite=True, b_ndim=2, overwrite_b=False} [id F] 13
 │  │           ├─ Switch [id G] 12
 │  │           │  ├─ ExpandDims{axes=[0, 1]} [id H] 11
 │  │           │  │  └─ All{axis=0} [id I] 5
 │  │           │  │     └─ Gt [id J] 4
 │  │           │  │        ├─ ExtractDiag{offset=0, axis1=0, axis2=1, view=True} [id K] 3
 │  │           │  │        │  └─ KroneckerProduct{inline=False} [id L] 2
 │  │           │  │        │     ├─ Cholesky{lower=True, check_finite=False, on_error='raise', overwrite_a=False} [id M] 1
 │  │           │  │        │     │  └─ A [id N]
 │  │           │  │        │     └─ Cholesky{lower=True, check_finite=False, on_error='raise', overwrite_a=False} [id O] 0
 │  │           │  │        │        └─ B [id P]
 │  │           │  │        └─ [0] [id Q]
 │  │           │  ├─ KroneckerProduct{inline=False} [id L] 2
 │  │           │  │  └─ ···
 │  │           │  └─ [[1]] [id R]
 │  │           └─ Sub [id S] 10
 │  │              ├─ ExpandDims{axis=1} [id T] 9
 │  │              │  └─ value [id U]
 │  │              └─ ExpandDims{axis=1} [id V] 8
 │  │                 └─ mu [id W]
 │  ├─ Shape_i{0} [id X] 7
 │  │  └─ value [id U]
 │  └─ CAReduce{Composite{(i0 + log(i1))}, axes=None} [id Y] 6
 │     └─ ExtractDiag{offset=0, axis1=0, axis2=1, view=True} [id K] 3
 │        └─ ···
 └─ All{axis=0} [id I] 5
    └─ ···

The important thing here is that we can totally eliminate the KroneckerProduct. Our ExractDiag rewrite (rewrite_diag_kronecker) isn't being triggered here for some reason, so that ExtractDiag branch should be an outer product instead of a full Kron. We also aren't doing anything with solve(kron(a, b), c), which can be rewritten to solve(a_val.T, solve(b_val, c.reshape((n, m))).T).T.ravel(). Both of those eliminate the kron completely, which is what we want.

For elemwise operations, we should be pushing them inside the kron as well. Something like kron(a, b) + eye((n*m)) * sigma should end up as kron(a + eye(n) * sigma, b + eye(m) * sigma), which we can then keep working with.

I haven't benchmarked it, but computing the logp with cholesky should be strictly faster than using the full eigenvalue decomposition, because it exploits more of the special structure. I am basing this guess on the method argument in MvNormal, where cholesky is faster than eigh. If I'm wrong though, we should be doing the eigh method in the base case.

[ricardoV94]

It's not diagonal off, it's on top of

pymc/pymc/distributions/multivariate.py

Line 1957 in dc7cfee

cov = cov + sigma**2 * pt.eye(cov.shape[-2])

Anyway besides that I don't think we can just use regular logp because ours is defined on the cholesky, with solve triangulars.

May need something like the PrecisionMvNormal rewrite