jax-ml/jax · error · NotImplementedError

gaussian_kde does not support complex coordinates

Error message

gaussian_kde does not support complex coordinates

What it means

All evaluation methods of gaussian_kde (evaluate, logpdf, pdf) reject complex-valued query points via _reshape_points, mirroring the complex-dataset restriction in __init__. The kernel math is real-only.

Source

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

    result = _gaussian_kernel_eval(True, self.dataset.T, self.weights[:, None],
                                   x.T, self.inv_cov)
    return result[:, 0]

  def integrate_box(self, low_bounds, high_bounds, maxpts=None):
    """This method is not implemented in the JAX interface."""
    del low_bounds, high_bounds, maxpts
    raise NotImplementedError(
        "only 1D box integrations are supported; use `integrate_box_1d`")

  def set_bandwidth(self, bw_method=None):
    """This method is not implemented in the JAX interface."""
    del bw_method
    raise NotImplementedError(
        "dynamically changing the bandwidth method is not supported")

  def _reshape_points(self, points):
    if dtypes.issubdtype(lax.dtype(points), np.complexfloating):
      raise NotImplementedError(
          "gaussian_kde does not support complex coordinates")
    points = jnp.atleast_2d(points)
    d, m = points.shape
    if d != self.d:
      if d == 1 and m == self.d:
        points = jnp.reshape(points, (self.d, 1))
      else:
        raise ValueError(
            "points have dimension {}, dataset has dimension {}".format(
                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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Evaluate on real coordinates: kde.evaluate(points.real)
2. Cast the query array: points = points.astype(jnp.float64)
3. Split into real/imaginary grids and evaluate separately if both are needed

Example fix

// before
vals = kde.evaluate(z)  # z complex
// after
vals = kde.evaluate(z.real)

if jnp.iscomplexobj(points):
    points = points.real

def points_are_real(points) -> bool:
    return not jnp.iscomplexobj(jnp.asarray(points))

When it happens

Trigger: kde.evaluate(jnp.array([1+1j, 2-1j])) or kde.logpdf(complex_grid).

Common situations: Evaluating a KDE on the output of an FFT/analytic-signal pipeline left in complex dtype; complex creeping in via promotion with a complex constant.

Related errors

• gaussian_kde does not support complex data
• only 1D box integrations are supported; use `integrate_box_1
• dynamically changing the bandwidth method is not supported
• jnp.partition for complex dtype is not implemented.
• jnp.argpartition for complex dtype is not implemented.

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