pymc-devs / pymc-devs/pytensor

Implement MvNormal as cholesky(cov) @ normal

Open
#1,115 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Op implementation random variables
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.