jax-ml/jax · error · NotImplementedError

gaussian_kde does not support complex data

Error message

gaussian_kde does not support complex data

What it means

jax.scipy.stats.gaussian_kde explicitly rejects complex-valued datasets because the covariance/Cholesky machinery is implemented for real floating types only. This mirrors an equivalent restriction historically present in scipy.

Source

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

  Parameters:
    dataset: arraylike, real-valued. Data from which to estimate the distribution.
      If 1D, shape is (n_data,). If 2D, shape is (n_dimensions, n_data).
    bw_method: string, scalar, or callable. Either "scott", "silverman", a scalar
      value, or a callable function which takes ``self`` as a parameter.
    weights: arraylike, optional. Weights of the same shape as the dataset.
  """
  neff: Any
  dataset: Any
  weights: Any
  covariance: Any
  inv_cov: Any

  def __init__(self, dataset, bw_method: BwMethod = None, weights=None):
    check_arraylike("gaussian_kde", dataset)
    dataset = jnp.atleast_2d(dataset)
    if dtypes.issubdtype(lax.dtype(dataset), np.complexfloating):
      raise NotImplementedError("gaussian_kde does not support complex data")
    if not dataset.size > 1:
      raise ValueError("`dataset` input should have multiple elements.")

    d, n = dataset.shape
    if weights is not None:
      check_arraylike("gaussian_kde", weights)
      dataset, weights = promote_dtypes_inexact(dataset, weights)
      weights = jnp.atleast_1d(weights)
      weights /= jnp.sum(weights)
      if weights.ndim != 1:
        raise ValueError("`weights` input should be one-dimensional.")
      if len(weights) != n:
        raise ValueError("`weights` input should be of length n")
    else:
      dataset, = promote_dtypes_inexact(dataset)
      weights = jnp.full(n, 1.0 / n, dtype=dataset.dtype)

    self._setattr("dataset", dataset)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Take real (or imaginary/absolute) part of the data before fitting: data = data.real
  2. Convert to real explicitly: dataset = dataset.astype(jnp.float64)
  3. If complex KDE is genuinely needed, fit separate KDEs on real and imaginary parts

Example fix

// before
kde = gaussian_kde(fft_result)  # complex
// after
kde = gaussian_kde(fft_result.real)
Defensive patterns

Strategy: type-guard

Validate before calling

data = jnp.real(jnp.asarray(data)) if jnp.iscomplexobj(data) else data

Type guard

def is_real_dataset(data) -> bool:
    return not jnp.iscomplexobj(jnp.asarray(data))

Prevention

When it happens

Trigger: Constructing gaussian_kde with a complex64/complex128 array, e.g. jax.scipy.stats.gaussian_kde(jnp.array([1+2j, 3-1j, ...])).

Common situations: Applying KDE to FFT output or signal-processing data left in complex dtype after a transform; dtype promotion accidentally producing complex after mixing complex and real arrays.

Related errors


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