Skip to content

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

Closed

Description

@jessegrabowski

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions