jax-ml/jax · error · ValueError

KDEs are not the same dimensionality

Error message

KDEs are not the same dimensionality

What it means

gaussian_kde.integrate_kde(other) computes the convolution of two Gaussian KDEs, which requires both KDEs to have the same dimensionality (other.d == self.d). Mismatched dimensionalities — e.g. a 1-D KDE against a 2-D KDE — raise this error before the covariance sum.

Source

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

                                     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)
    return jnp.sum(self.weights * (special.ndtr(high) - special.ndtr(low)))

  def integrate_kde(self, other):
    """Integrate the product of two Gaussian KDE distributions."""
    if other.d != self.d:
      raise ValueError("KDEs are not the same dimensionality")

    chol = linalg.cho_factor(self.covariance + other.covariance)
    norm = jnp.sqrt(2 * np.pi)**self.d * jnp.prod(jnp.diag(chol[0]))
    norm = 1.0 / norm

    sm, lg = (self, other) if self.n < other.n else (other, self)
    result = api.vmap(partial(_gaussian_kernel_convolve, chol, norm, lg.dataset,
                              lg.weights),
                      in_axes=1)(sm.dataset)
    return jnp.sum(result * sm.weights)

  @api.jit(static_argnames=("shape",))
  def resample(self, key, shape=()):
    r"""Randomly sample a dataset from the estimated pdf

    Args:
      key: a PRNG key used as the random key.
      shape: optional, a tuple of nonnegative integers specifying the result

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Verify d with kde.dataset.shape[0] for both objects and make the input layouts consistent
2. Rebuild KDEs from consistently shaped data (1-D vectors for 1-D KDEs)
3. If the dimensionalities genuinely differ, compare marginals by fitting KDEs on the same feature subset

Example fix

// before
p = kde_a.integrate_kde(kde_b)  # d=1 vs d=2
// after
kde_b1 = gaussian_kde(pts_b[0])  # marginal of first dim
p = kde_a.integrate_kde(kde_b1)

assert kde_a.d == kde_b.d, 'KDE dimensionality mismatch'

def same_dimensionality(a, b) -> bool:
    return a.d == b.d

When it happens

Trigger: kde1.integrate_kde(kde2) where kde1 was built from shape (1, n) data and kde2 from (2, m) data.

Common situations: Comparing densities of features with different dimensionalities in a two-sample test; one dataset accidentally kept an extra leading axis so atleast_2d produced a different d.

Related errors

• `weights` input should be one-dimensional.
• mean does not have dimension {self.d}
• covariance does not have dimension {self.d}
• integrate_box_1d() only handles 1D pdfs
• points have dimension {}, dataset has dimension {}

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