jax-ml/jax · error · ValueError

`weights` input should be of length n

Error message

`weights` input should be of length n

What it means

gaussian_kde weights must have exactly n entries, where n is dataset.shape[1] (the number of samples) after the dataset is made 2-D. Note the length check happens after in-place normalization weights /= jnp.sum(weights), so mismatched weights are rejected here rather than earlier.

Source

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

  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):
      factor = bw_method(self)
    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Make len(weights) == dataset.shape[1]; check dataset = jnp.atleast_2d(dataset).shape first
  2. If dataset is (n, d) user-side, transpose it or slice weights to match the sample axis
  3. Recompute weights after any filtering/subsampling of the data

Example fix

// before
kde = gaussian_kde(X.T, weights=w)  # w computed for X rows, wrong n
// after
kde = gaussian_kde(X.T, weights=w[:X.shape[0]])  # or recompute w for kept rows
Defensive patterns

Strategy: validation

Validate before calling

n = jnp.atleast_2d(jnp.asarray(dataset)).shape[1]
assert len(weights) == n, f'weights must have length {n}'

Prevention

When it happens

Trigger: Passing weights of length m != dataset.shape[1]; passing a scalar weight; passing weights aligned with a transposed dataset layout.

Common situations: Subsampling the dataset but reusing old weights; weights computed against a different split of the data; dataset auto-transposed by atleast_2d so n differs from what the user assumed.

Related errors


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