jax-ml/jax · error · ValueError

points and xi must have same trailing dim

Error message

points and xi must have same trailing dim

What it means

The internal _gaussian_kernel_eval (used by evaluate/logpdf via the non-fast path) requires the query points xi to have the same trailing (second) dimension as the kernel points: xi.shape[1] == points.shape[1] == d. A trailing-dim mismatch means xi encodes a different dimensionality than the kernels.

Source

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

                d, self.d))
    return points


def _gaussian_kernel_convolve(chol, norm, target, weights, mean):
  diff = target - mean[:, None]
  alpha = linalg.cho_solve(chol, diff)
  arg = 0.5 * jnp.sum(diff * alpha, axis=0)
  return norm * jnp.sum(jnp.exp(-arg) * weights)


@api.jit(static_argnums=0)
def _gaussian_kernel_eval(in_log, points, values, xi, precision):
  points, values, xi, precision = promote_dtypes_inexact(
      points, values, xi, precision)
  d = points.shape[1]

  if xi.shape[1] != d:
    raise ValueError("points and xi must have same trailing dim")
  if precision.shape != (d, d):
    raise ValueError("precision matrix must match data dims")

  whitening = linalg.cholesky(precision, lower=True)
  points = jnp.dot(points, whitening)
  xi = jnp.dot(xi, whitening)
  log_norm = jnp.sum(jnp.log(
      jnp.diag(whitening))) - 0.5 * d * jnp.log(2 * np.pi)

  def kernel(x_test, x_train, y_train):
    arg = log_norm - 0.5 * jnp.sum(jnp.square(x_train - x_test))
    if in_log:
      return jnp.log(y_train) + arg
    else:
      return y_train * jnp.exp(arg)

  reduce = special.logsumexp if in_log else jnp.sum
  reduced_kernel = lambda x: reduce(api.vmap(kernel, in_axes=(None, 0, 0))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Transpose the evaluation points so the trailing dim equals the KDE dimensionality
2. Verify points.shape[1] == kde.d before calling evaluate
3. For custom calls to _gaussian_kernel_eval, ensure points and xi share shape[1]

Example fix

// before
vals = kde.evaluate(X)  # X.shape == (N, k), k != d
// after
vals = kde.evaluate(X[:, :kde.d].T)  # select matching dims and transpose

assert jnp.asarray(xi).shape[1] == jnp.asarray(points).shape[1]

def trailing_dims_match(points, xi) -> bool:
    return jnp.asarray(xi).shape[1] == jnp.asarray(points).shape[1]

When it happens

Trigger: Calling kde.evaluate on points whose trailing dimension differs from kde.d, e.g. passing (m, k) with k != d; mixing row-major data into the (d, m) column-major convention.

Common situations: Same root cause as the (N, d) vs (d, N) transposition issue — this fires on the kernel helper when the reshaping guards were bypassed or arguments were reordered.

Related errors

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

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