jax-ml/jax · error · NotImplementedError

allow_singular argument of multivariate_normal.logpdf

Error message

allow_singular argument of multivariate_normal.logpdf

What it means

jax.scipy.stats.multivariate_normal.logpdf raises NotImplementedError when allow_singular is set to anything other than None. JAX's implementation relies on a Cholesky factorization, which requires the covariance matrix to be positive definite, so the singular-covariance path that scipy supports is simply not implemented.

Source

Thrown at jax/_src/scipy/stats/multivariate_normal.py:54

     f(x) = \frac{1}{\sqrt{(2\pi)^k\det\Sigma}}\exp\left(-\frac{(x-\mu)^T\Sigma^{-1}(x-\mu)}{2} \right)

  where :math:`\mu` is the ``mean``, :math:`\Sigma` is the covariance matrix (``cov``), and
  :math:`k` is the rank of :math:`\Sigma`.

  Args:
    x: arraylike, value at which to evaluate the PDF
    mean: arraylike, centroid of distribution
    cov: arraylike, covariance matrix of distribution
    allow_singular: not supported

  Returns:
    array of logpdf values.

  See Also:
    :func:`jax.scipy.stats.multivariate_normal.pdf`
  """
  if allow_singular is not None:
    raise NotImplementedError("allow_singular argument of multivariate_normal.logpdf")
  x, mean, cov = promote_dtypes_inexact(x, mean, cov)
  if not mean.shape:
    return (-1/2 * jnp.square(x - mean) / cov
            - 1/2 * (jnp.log(2*np.pi) + jnp.log(cov)))
  else:
    n = mean.shape[-1]
    if not np.shape(cov):
      y = x - mean
      return (-1/2 * jnp_einsum.einsum('...i,...i->...', y, y) / cov
              - n/2 * (jnp.log(2*np.pi) + jnp.log(cov)))
    else:
      if cov.ndim < 2 or cov.shape[-2:] != (n, n):
        raise ValueError("multivariate_normal.logpdf got incompatible shapes")
      L = lax.linalg.cholesky(cov)
      y = jnp_vectorize.vectorize(
        partial(lax.linalg.triangular_solve, lower=True, transpose_a=True),
        signature="(n,n),(n)->(n)"
      )(L, x - mean)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Remove the allow_singular argument (or pass None) and ensure cov is non-singular, e.g. add jitter: cov + 1e-6 * jnp.eye(d)
Regularize the covariance beforehand with a shrinkage estimator or diagonal loading
If you truly need singular support, fall back to scipy.stats.multivariate_normal for that computation (non-JIT)

Example fix

// before
jax.scipy.stats.multivariate_normal.logpdf(x, mean, cov, allow_singular=True)

// after
cov_reg = cov + 1e-6 * jnp.eye(cov.shape[-1])
jax.scipy.stats.multivariate_normal.logpdf(x, mean, cov_reg)

Defensive patterns

Strategy: validation

Validate before calling

def safe_mvn_logpdf(x, mean, cov, jitter=1e-6):
    d = cov.shape[-1]
    cov = cov + jitter * jnp.eye(d, dtype=cov.dtype)
    return jax.scipy.stats.multivariate_normal.logpdf(x, mean, cov)  # no allow_singular

Try / catch

try:
    return mvn.logpdf(x, mean, cov)
except NotImplementedError as e:
    if 'allow_singular' in str(e):
        cov = cov + 1e-6 * jnp.eye(cov.shape[-1])
        return mvn.logpdf(x, mean, cov)
    raise

Prevention

Never pass allow_singular to JAX's multivariate_normal
Always add diagonal jitter to empirical covariances
Check cov condition number before logpdf in tests

When it happens

Trigger: Calling jax.scipy.stats.multivariate_normal.logpdf or .pdf with allow_singular=True (a valid scipy argument) in an attempt to replicate scipy behavior.

Common situations: Porting scipy code to JAX verbatim, or fitting models where the empirical covariance is singular (perfectly correlated features, n_samples < n_dimensions) and scipy's allow_singular=True was used as a workaround.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/16cab7020ca99a25. Report an issue: GitHub.