pytensor icon indicating copy to clipboard operation
pytensor copied to clipboard

Implement MvNormal as cholesky(cov) @ normal

Open ricardoV94 opened this issue 1 year ago • 5 comments

Description

This is much faster, and even more in PyMC models that are usually parametrized with a direct prior on the cholesky.

import pytensor
import pytensor.tensor as pt

srng = pt.random.RandomStream()

x = srng.multivariate_normal([0, 0], [[1, 0.5], [0.5, 1]])
fn = pytensor.function([], x)
%timeit fn()  # 510 µs ± 81.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# Decompose cholesky in graph (numpy probably does this under the hood)
A = pt.linalg.cholesky([[1, 0.5], [0.5, 1]])
x = A @ srng.normal(size=(2,))
fn = pytensor.function([], x)
%timeit fn()  # 27.4 µs ± 3.27 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In general we should probably reduce the number of pure RV Ops we have. This allows more optimizations and makes it easier to implement different backends.

We should implement the MvNormal as an OpFromGraph that gets inlined after canonicalization (not as early as the ones with inline=True)

ricardoV94 avatar Dec 11 '24 18:12 ricardoV94