pymc-devs / pymc-devs/pytensor

Numba MvNormalRV raises LinAlgError on a NaN covariance: np.linalg calls inside the njit body

Open
#2,421 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement linalg numba
Dominant language
Python
Stars
644
Forks
208
Avg merge
2d 14h
Merged PRs (30d)
16

Description

pt.random.multivariate_normal compiled with the numba backend raises LinAlgError when the covariance holds a NaN, for every method. core_MvNormalRV in pytensor/link/numba/dispatch/random.py factors the covariance with np.linalg.cholesky, np.linalg.svd, or np.linalg.eigh inside the njit body, and numba's implementations of those check finiteness and raise. In a draw batched over many covariances, one non-finite covariance fails the whole batch.

import numpy as np
import pytensor
import pytensor.tensor as pt

cov = pt.matrix("cov")
rng = pytensor.shared(np.random.default_rng(0))
bad_cov = np.full((2, 2), np.nan)

draw = pt.random.multivariate_normal(np.zeros(2), cov, method="svd", rng=rng)
print(pytensor.function([cov], draw, mode="NUMBA")(bad_cov))  # LinAlgError: Array must not contain infs or NaNs.

Fix: replace the three np.linalg calls with the LAPACK-backed overload targets in pytensor/link/numba/dispatch/linalg/decomposition/, which skip the finiteness check and return NaN: _cholesky from cholesky.py, _svd_gesdd_full from svd.py, and _eigh from eigen.py.

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 in pytensor/link/numba/dispatch/random.py at core_MvNormalRV, then inspect the overload targets in pytensor/link/numba/dispatch/linalg/decomposition/cholesky.py, svd.py, and eigen.py. Run the provided NUMBA reproduction for each method and verify that a covariance containing NaN returns NaN rather than raising LinAlgError, including for batched covariances.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
backend
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
84/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.