{"record":{"id":"16cab7020ca99a25","repo":"jax-ml/jax","slug":"allow-singular-argument-of-multivariate-normal-log","errorCode":null,"errorMessage":"allow_singular argument of multivariate_normal.logpdf","messagePattern":"allow_singular argument of multivariate_normal\\.logpdf","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/stats/multivariate_normal.py","lineNumber":54,"sourceCode":"     f(x) = \\frac{1}{\\sqrt{(2\\pi)^k\\det\\Sigma}}\\exp\\left(-\\frac{(x-\\mu)^T\\Sigma^{-1}(x-\\mu)}{2} \\right)\n\n  where :math:`\\mu` is the ``mean``, :math:`\\Sigma` is the covariance matrix (``cov``), and\n  :math:`k` is the rank of :math:`\\Sigma`.\n\n  Args:\n    x: arraylike, value at which to evaluate the PDF\n    mean: arraylike, centroid of distribution\n    cov: arraylike, covariance matrix of distribution\n    allow_singular: not supported\n\n  Returns:\n    array of logpdf values.\n\n  See Also:\n    :func:`jax.scipy.stats.multivariate_normal.pdf`\n  \"\"\"\n  if allow_singular is not None:\n    raise NotImplementedError(\"allow_singular argument of multivariate_normal.logpdf\")\n  x, mean, cov = promote_dtypes_inexact(x, mean, cov)\n  if not mean.shape:\n    return (-1/2 * jnp.square(x - mean) / cov\n            - 1/2 * (jnp.log(2*np.pi) + jnp.log(cov)))\n  else:\n    n = mean.shape[-1]\n    if not np.shape(cov):\n      y = x - mean\n      return (-1/2 * jnp_einsum.einsum('...i,...i->...', y, y) / cov\n              - n/2 * (jnp.log(2*np.pi) + jnp.log(cov)))\n    else:\n      if cov.ndim < 2 or cov.shape[-2:] != (n, n):\n        raise ValueError(\"multivariate_normal.logpdf got incompatible shapes\")\n      L = lax.linalg.cholesky(cov)\n      y = jnp_vectorize.vectorize(\n        partial(lax.linalg.triangular_solve, lower=True, transpose_a=True),\n        signature=\"(n,n),(n)->(n)\"\n      )(L, x - mean)","sourceCodeStart":36,"sourceCodeEnd":72,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/stats/multivariate_normal.py#L36-L72","documentation":"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.","triggerScenarios":"Calling jax.scipy.stats.multivariate_normal.logpdf or .pdf with allow_singular=True (a valid scipy argument) in an attempt to replicate scipy behavior.","commonSituations":"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.","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)"],"exampleFix":"// before\njax.scipy.stats.multivariate_normal.logpdf(x, mean, cov, allow_singular=True)\n\n// after\ncov_reg = cov + 1e-6 * jnp.eye(cov.shape[-1])\njax.scipy.stats.multivariate_normal.logpdf(x, mean, cov_reg)","handlingStrategy":"validation","validationCode":"def safe_mvn_logpdf(x, mean, cov, jitter=1e-6):\n    d = cov.shape[-1]\n    cov = cov + jitter * jnp.eye(d, dtype=cov.dtype)\n    return jax.scipy.stats.multivariate_normal.logpdf(x, mean, cov)  # no allow_singular","typeGuard":null,"tryCatchPattern":"try:\n    return mvn.logpdf(x, mean, cov)\nexcept NotImplementedError as e:\n    if 'allow_singular' in str(e):\n        cov = cov + 1e-6 * jnp.eye(cov.shape[-1])\n        return mvn.logpdf(x, mean, cov)\n    raise","preventionTips":["Never pass allow_singular to JAX's multivariate_normal","Always add diagonal jitter to empirical covariances","Check cov condition number before logpdf in tests"],"tags":["jax","scipy","multivariate-normal","not-implemented","covariance"],"backgroundTag":"unsupported-argument-not-implemented","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}