jax-ml/jax · error · ValueError

multivariate_normal.logpdf got incompatible shapes

Error message

multivariate_normal.logpdf got incompatible shapes

What it means

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.

Source

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

  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)
      return (-1/2 * jnp_einsum.einsum('...i,...i->...', y, y) - n/2 * jnp.log(2*np.pi)
              - jnp.log(L.diagonal(axis1=-1, axis2=-2)).sum(-1))


def pdf(x: ArrayLike, mean: ArrayLike, cov: ArrayLike) -> Array:
  r"""Multivariate normal probability distribution function.

  JAX implementation of :obj:`scipy.stats.multivariate_normal` ``pdf``.

  The multivariate normal PDF is defined as

  .. math::

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Check mean.shape[-1] equals cov.shape[-1] == cov.shape[-2] before calling
  2. If you meant per-dimension variances, expand to a diagonal matrix: jnp.diag(var_vector)
  3. If you meant a single shared variance, pass a scalar cov (then the fast scalar path is used)

Example fix

// before
mean = jnp.zeros(3)
var = jnp.array([1.0, 2.0, 3.0])
mvn.logpdf(x, mean, var)  # shape (3,) != (3,3) -> raises

// after
mean = jnp.zeros(3)
cov = jnp.diag(jnp.array([1.0, 2.0, 3.0]))
mvn.logpdf(x, mean, cov)
Defensive patterns

Strategy: validation

Validate before calling

def check_mvn_shapes(x, mean, cov):
    n = mean.shape[-1]
    assert x.shape[-1] == n, f"x event dim {x.shape[-1]} != mean dim {n}"
    if cov.ndim >= 2:
        assert cov.shape[-2:] == (n, n), f"cov trailing shape {cov.shape[-2:]} != {(n, n)}"
    return True

Try / catch

try:
    mvn.logpdf(x, mean, cov)
except ValueError as e:
    if 'incompatible shapes' in str(e):
        cov = jnp.diag(cov) if cov.ndim == 1 and cov.shape[-1] == mean.shape[-1] else cov
        mvn.logpdf(x, mean, cov)
    else: raise

Prevention

When it happens

Trigger: 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).

Common situations: 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.

Related errors

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