pymc-devs / pymc-devs/pytensor
Implement MvNormal as cholesky(cov) @ normal
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
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)
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by tracing the RandomStream.multivariate_normal entry point and the existing cholesky and normal operations. Review how OpFromGraph implementations are inlined during canonicalization, then verify that MvNormal is represented through those operations and preserves the stated sampling behavior and optimization benefit.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100