jax-ml/jax · error · RuntimeError

cannot handle multidimensional fweights

Error message

cannot handle multidimensional fweights

What it means

jnp.cov supports optional frequency weights (fweights) which must be a 1D array with one weight per observation. If np.ndim(fweights) > 1, RuntimeError('cannot handle multidimensional fweights') is raised, mirroring NumPy.

Source

Thrown at jax/_src/numpy/lax_numpy.py:9212

    return array([]).reshape(0, 0)

  if y is not None:
    y_arr = atleast_2d(y)
    if not rowvar and y_arr.shape[0] != 1:
      y_arr = y_arr.T
    X = concatenate((X, y_arr), axis=0)
  if X.shape[1] == 0:
    cov_shape = () if X.shape[0] == 1 else (X.shape[0], X.shape[0])
    return array_creation.full(cov_shape, np.nan, dtype=X.dtype)

  if ddof is None:
    ddof = 1 if bias == 0 else 0

  w: Array | None = None
  if fweights is not None:
    fweights = util.ensure_arraylike("cov", fweights)
    if np.ndim(fweights) > 1:
      raise RuntimeError("cannot handle multidimensional fweights")
    if np.shape(fweights)[0] != X.shape[1]:
      raise RuntimeError("incompatible numbers of samples and fweights")
    if not issubdtype(fweights.dtype, np.integer):
      raise TypeError("fweights must be integer.")
    # Ensure positive fweights; note that numpy raises an error on negative fweights.
    w = abs(fweights)
  if aweights is not None:
    aweights = util.ensure_arraylike("cov", aweights)
    if np.ndim(aweights) > 1:
      raise RuntimeError("cannot handle multidimensional aweights")
    if np.shape(aweights)[0] != X.shape[1]:
      raise RuntimeError("incompatible numbers of samples and aweights")
    # Ensure positive aweights: note that numpy raises an error for negative aweights.
    aweights = abs(aweights)
    w = asarray(aweights if w is None else w * aweights)

  if dtype is not None:
    X = X.astype(dtype)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Flatten weights: fweights=w.ravel() or w.squeeze()
  2. Ensure weights shape is (n_observations,) matching X.shape[1]
  3. Squeeze stray axes introduced by keepdims=True upstream

Example fix

// before
jnp.cov(m, fweights=w[:, None])  # ValueError (RuntimeError)
// after
jnp.cov(m, fweights=w.ravel())
Defensive patterns

Strategy: validation

Validate before calling

if fweights is not None:
    fweights = jnp.asarray(fweights).ravel()
jnp.cov(m, fweights=fweights)

Type guard

def valid_fweights(w, n_obs) -> bool:
    w = jnp.asarray(w)
    return w.ndim == 1 and w.shape[0] == n_obs

Prevention

When it happens

Trigger: jnp.cov(m, fweights=w) where w has 2 or more dimensions, e.g. a (n, 1) column vector instead of (n,) weights.

Common situations: Passing weights that kept a trailing axis after slicing (w[:, None]) or batched weight matrices; forgetting to squeeze a column vector.

Related errors


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