{"record":{"id":"edb55246d0f75bf5","repo":"jax-ml/jax","slug":"points-and-xi-must-have-same-trailing-dim","errorCode":null,"errorMessage":"points and xi must have same trailing dim","messagePattern":"points and xi must have same trailing dim","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/stats/kde.py","lineNumber":262,"sourceCode":"                d, self.d))\n    return points\n\n\ndef _gaussian_kernel_convolve(chol, norm, target, weights, mean):\n  diff = target - mean[:, None]\n  alpha = linalg.cho_solve(chol, diff)\n  arg = 0.5 * jnp.sum(diff * alpha, axis=0)\n  return norm * jnp.sum(jnp.exp(-arg) * weights)\n\n\n@api.jit(static_argnums=0)\ndef _gaussian_kernel_eval(in_log, points, values, xi, precision):\n  points, values, xi, precision = promote_dtypes_inexact(\n      points, values, xi, precision)\n  d = points.shape[1]\n\n  if xi.shape[1] != d:\n    raise ValueError(\"points and xi must have same trailing dim\")\n  if precision.shape != (d, d):\n    raise ValueError(\"precision matrix must match data dims\")\n\n  whitening = linalg.cholesky(precision, lower=True)\n  points = jnp.dot(points, whitening)\n  xi = jnp.dot(xi, whitening)\n  log_norm = jnp.sum(jnp.log(\n      jnp.diag(whitening))) - 0.5 * d * jnp.log(2 * np.pi)\n\n  def kernel(x_test, x_train, y_train):\n    arg = log_norm - 0.5 * jnp.sum(jnp.square(x_train - x_test))\n    if in_log:\n      return jnp.log(y_train) + arg\n    else:\n      return y_train * jnp.exp(arg)\n\n  reduce = special.logsumexp if in_log else jnp.sum\n  reduced_kernel = lambda x: reduce(api.vmap(kernel, in_axes=(None, 0, 0))","sourceCodeStart":244,"sourceCodeEnd":280,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/stats/kde.py#L244-L280","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Transpose the evaluation points so the trailing dim equals the KDE dimensionality","Verify points.shape[1] == kde.d before calling evaluate","For custom calls to _gaussian_kernel_eval, ensure points and xi share shape[1]"],"exampleFix":"// before\nvals = kde.evaluate(X)  # X.shape == (N, k), k != d\n// after\nvals = kde.evaluate(X[:, :kde.d].T)  # select matching dims and transpose","handlingStrategy":"validation","validationCode":"assert jnp.asarray(xi).shape[1] == jnp.asarray(points).shape[1]","typeGuard":"def trailing_dims_match(points, xi) -> bool:\n    return jnp.asarray(xi).shape[1] == jnp.asarray(points).shape[1]","tryCatchPattern":null,"preventionTips":["Route queries through kde.evaluate so reshaping guards apply","Keep the (d, N) convention end-to-end","Assert trailing dims before low-level kernel calls"],"tags":["jax","scipy","kde","shape-validation"],"backgroundTag":"invalid-array-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}