jax-ml/jax · error · ValueError

covariance does not have dimension {self.d}

Error message

covariance does not have dimension {self.d}

What it means

In gaussian_kde.integrate_gaussian, cov is atleast_2d'd and must have shape exactly (self.d, self.d) to be added to the KDE covariance. Non-square, wrong-size, or wrongly-oriented matrices raise this error.

Source

Thrown at jax/_src/scipy/stats/kde.py:150

    """Evaluate the Gaussian KDE on the given points."""
    check_arraylike("evaluate", points)
    points = self._reshape_points(points)
    result = _gaussian_kernel_eval(False, self.dataset.T, self.weights[:, None],
                                   points.T, self.inv_cov)
    return result[:, 0]

  def __call__(self, points):
    return self.evaluate(points)

  def integrate_gaussian(self, mean, cov):
    """Integrate the distribution weighted by a Gaussian."""
    mean = jnp.atleast_1d(jnp.squeeze(mean))
    cov = jnp.atleast_2d(cov)

    if mean.shape != (self.d,):
      raise ValueError(f"mean does not have dimension {self.d}")
    if cov.shape != (self.d, self.d):
      raise ValueError(f"covariance does not have dimension {self.d}")

    chol = linalg.cho_factor(self.covariance + cov)
    norm = jnp.sqrt(2 * np.pi)**self.d * jnp.prod(jnp.diag(chol[0]))
    norm = 1.0 / norm
    return _gaussian_kernel_convolve(chol, norm, self.dataset, self.weights,
                                     mean)

  @api.jit
  def integrate_box_1d(self, low, high):
    """Integrate the distribution over the given limits."""
    if self.d != 1:
      raise ValueError("integrate_box_1d() only handles 1D pdfs")
    if np.ndim(low) != 0 or np.ndim(high) != 0:
      raise ValueError(
          "the limits of integration in integrate_box_1d must be scalars")
    sigma = jnp.squeeze(jnp.sqrt(self.covariance))
    low = jnp.squeeze((low - self.dataset) / sigma)
    high = jnp.squeeze((high - self.dataset) / sigma)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Ensure cov.shape == (kde.d, kde.d)
  2. If you have per-dimension variances, wrap with jnp.diag(variances)
  3. If cov is a Cholesky factor, reconstruct the covariance via L @ L.T

Example fix

// before
kde.integrate_gaussian(mu, variances)  # shape (d,)
// after
kde.integrate_gaussian(mu, jnp.diag(variances))
Defensive patterns

Strategy: validation

Validate before calling

cov = jnp.atleast_2d(jnp.asarray(cov))
assert cov.shape == (kde.d, kde.d), f'cov must be {(kde.d, kde.d)}'

Type guard

def cov_shape_ok(cov, kde) -> bool:
    return jnp.asarray(cov).shape == (kde.d, kde.d)

Prevention

When it happens

Trigger: Passing cov of shape (d, d+1), (1,) (becomes (1,1) while d > 1), or a full covariance of a different model with mismatched dimensionality.

Common situations: Mixing KDEs of different dimensionality in a mixture model; passing a precision/Cholesky factor instead of the covariance; passing a vector of variances instead of a matrix (use jnp.diag(variances)).

Related errors


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