{"record":{"id":"2e9625a6bcfb82c2","repo":"jax-ml/jax","slug":"multivariate-normal-logpdf-got-incompatible-shapes","errorCode":null,"errorMessage":"multivariate_normal.logpdf got incompatible shapes","messagePattern":"multivariate_normal\\.logpdf got incompatible shapes","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/stats/multivariate_normal.py","lineNumber":67,"sourceCode":"\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)\n      return (-1/2 * jnp_einsum.einsum('...i,...i->...', y, y) - n/2 * jnp.log(2*np.pi)\n              - jnp.log(L.diagonal(axis1=-1, axis2=-2)).sum(-1))\n\n\ndef pdf(x: ArrayLike, mean: ArrayLike, cov: ArrayLike) -> Array:\n  r\"\"\"Multivariate normal probability distribution function.\n\n  JAX implementation of :obj:`scipy.stats.multivariate_normal` ``pdf``.\n\n  The multivariate normal PDF is defined as\n\n  .. math::\n","sourceCodeStart":49,"sourceCodeEnd":85,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/stats/multivariate_normal.py#L49-L85","documentation":"When the covariance argument has shape (i.e. is treated as a full covariance matrix), multivariate_normal.logpdf requires its trailing two dimensions to be (n, n) where n = mean.shape[-1]. The raised ValueError indicates cov is not a square matrix matching the event dimension of x and mean.","triggerScenarios":"Calling logpdf/pdf with mean of shape (..., n) and cov whose ndim >= 2 but cov.shape[-2:] != (n, n), e.g. mean of dim 3 with a 2x2 covariance, or passing a cov of ndim < 2 when mean is a vector (the diagonal branch requires scalar cov).","commonSituations":"Mismatched feature dimension between data and covariance (changed embedding size, transposed matrices), or passing a variance vector where a full matrix or scalar is expected. Note the API only accepts a scalar variance or a full (n, n) matrix, not a per-dimension variance vector.","solutions":["Check mean.shape[-1] equals cov.shape[-1] == cov.shape[-2] before calling","If you meant per-dimension variances, expand to a diagonal matrix: jnp.diag(var_vector)","If you meant a single shared variance, pass a scalar cov (then the fast scalar path is used)"],"exampleFix":"// before\nmean = jnp.zeros(3)\nvar = jnp.array([1.0, 2.0, 3.0])\nmvn.logpdf(x, mean, var)  # shape (3,) != (3,3) -> raises\n\n// after\nmean = jnp.zeros(3)\ncov = jnp.diag(jnp.array([1.0, 2.0, 3.0]))\nmvn.logpdf(x, mean, cov)","handlingStrategy":"validation","validationCode":"def check_mvn_shapes(x, mean, cov):\n    n = mean.shape[-1]\n    assert x.shape[-1] == n, f\"x event dim {x.shape[-1]} != mean dim {n}\"\n    if cov.ndim >= 2:\n        assert cov.shape[-2:] == (n, n), f\"cov trailing shape {cov.shape[-2:]} != {(n, n)}\"\n    return True","typeGuard":null,"tryCatchPattern":"try:\n    mvn.logpdf(x, mean, cov)\nexcept ValueError as e:\n    if 'incompatible shapes' in str(e):\n        cov = jnp.diag(cov) if cov.ndim == 1 and cov.shape[-1] == mean.shape[-1] else cov\n        mvn.logpdf(x, mean, cov)\n    else: raise","preventionTips":["Assert cov.shape[-2:] == (d, d) right after building cov","Wrap per-dimension variances with jnp.diag before calling","Log shapes of x/mean/cov in debugging setup"],"tags":["jax","scipy","multivariate-normal","shape-mismatch"],"backgroundTag":"matrix-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}