jax-ml/jax · error · ValueError

mean does not have dimension {self.d}

Error message

mean does not have dimension {self.d}

What it means

In gaussian_kde.integrate_gaussian, mean is squeezed then atleast_1d'd and must end up with shape exactly (self.d,). Passing a mean with extra/missing components, or one that squeezes to a scalar when d == 1 but was passed with a spurious dimension (or vice versa), triggers this error.

Source

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

  def evaluate(self, points):
    """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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape mean explicitly to (kde.d,): mean = jnp.reshape(mean, (kde.d,))
2. Check kde.d via kde.dataset.shape[0] and build mean to match
3. Avoid batching here; use vmap over single calls if multiple means are needed

Example fix

// before
kde.integrate_gaussian(mu, Sigma)  # mu.shape == (1, d)
// after
kde.integrate_gaussian(jnp.squeeze(mu, axis=0), Sigma)

mean = jnp.reshape(jnp.asarray(mean), (kde.d,))

def mean_shape_ok(mean, kde) -> bool:
    return jnp.asarray(mean).shape == (kde.d,) or jnp.squeeze(mean).shape == (kde.d,)

When it happens

Trigger: integrate_gaussian(mean, cov) with mean of shape (d+1,), (), (1, 1) when d == 1 after squeeze, or a batched mean of shape (B, d).

Common situations: Reusing a mean vector computed for a different dataset dimensionality; passing nested single-element arrays from upstream math that squeeze unexpectedly.

Related errors

• `weights` input should be one-dimensional.
• covariance does not have dimension {self.d}
• KDEs are not the same dimensionality
• points have dimension {}, dataset has dimension {}
• points and xi must have same trailing dim

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