pymc-devs / pymc-devs/pytensor
Numba MvNormalRV raises LinAlgError on a NaN covariance: np.linalg calls inside the njit body
Nobody has claimed this yet.
- 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
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 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