jax-ml/jax · error · ValueError

`weights` input should be one-dimensional.

Error message

`weights` input should be one-dimensional.

What it means

When weights are passed to gaussian_kde they must be 1-D after jnp.atleast_1d, i.e. one weight per sample point along the last axis of the (d, n) dataset. A 2-D weights array raises this error. Note weights are normalized in place before the check, and a 2-D input survives normalization only to be rejected here.

Source

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

  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)
    self._setattr("weights", weights)
    neff = self._setattr("neff", 1 / jnp.sum(weights**2))

    bw_method = "scott" if bw_method is None else bw_method
    if bw_method == "scott":
      factor = jnp.power(neff, -1. / (d + 4))
    elif bw_method == "silverman":
      factor = jnp.power(neff * (d + 2) / 4.0, -1. / (d + 4))
    elif jnp.isscalar(bw_method) and not isinstance(bw_method, str):
      factor = cast(Array, bw_method)
    elif callable(bw_method):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Flatten/squeeze weights to shape (n,): weights = weights.squeeze() or weights.ravel()
2. If you truly need per-dimension importance, resample or transform the data instead of the weights
3. Verify weights length matches dataset.shape[1]

Example fix

// before
kde = gaussian_kde(data, weights=w)  # w.shape == (n, 1)
// after
kde = gaussian_kde(data, weights=w.ravel())

weights = jnp.ravel(jnp.asarray(weights))
n = jnp.atleast_2d(jnp.asarray(dataset)).shape[1]
assert weights.ndim == 1 and weights.shape[0] == n

def weights_valid(weights, dataset) -> bool:
    w = jnp.atleast_1d(jnp.asarray(weights))
    n = jnp.atleast_2d(jnp.asarray(dataset)).shape[1]
    return w.ndim == 1 and w.shape[0] == n

When it happens

Trigger: Passing weights shaped (d, n) or (n, 1) instead of (n,) when dataset has shape (d, n).

Common situations: Using per-dimension weights (unsupported — KDE weights are per-sample); forgetting to squeeze weights produced by broadcasting or from a column-vector-like array.

Related errors

• `weights` input should be of length n
• 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/b6d36831f877e5ee. Report an issue: GitHub.