jax-ml/jax · error · TypeError

expected a 1-d array for weights

Error message

expected a 1-d array for weights

What it means

In polyfit, sample weights w must be a 1-D array with one weight per observation, because they are applied row-wise to the design matrix and responses (lhs *= w[:, None]). A multi-dimensional or scalar-broadcast weight array is rejected with a TypeError.

Source

Thrown at jax/_src/numpy/polynomial.py:258

  if x_arr.size == 0:
    raise TypeError("expected non-empty vector for x")
  if y_arr.ndim < 1 or y_arr.ndim > 2:
    raise TypeError("expected 1D or 2D array for y")
  if x_arr.shape[0] != y_arr.shape[0]:
    raise TypeError("expected x and y to have same length")

  if rcond is None:
    rcond = len(x_arr) * float(finfo(x_arr.dtype).eps)
  rcond = core.concrete_or_error(float, rcond, "rcond must be float")
  # set up least squares equation for powers of x
  lhs = vander(x_arr, order)
  rhs = y_arr

  # apply weighting
  if w is not None:
    w_arr, = promote_dtypes_inexact(w)
    if w_arr.ndim != 1:
      raise TypeError("expected a 1-d array for weights")
    if w_arr.shape[0] != y_arr.shape[0]:
      raise TypeError("expected w and y to have the same length")
    lhs *= w_arr[:, np.newaxis]
    if rhs.ndim == 2:
      rhs *= w_arr[:, np.newaxis]
    else:
      rhs *= w_arr

  # scale lhs to improve condition number and solve
  scale = sqrt((lhs*lhs).sum(axis=0))
  lhs /= scale[np.newaxis, :]
  c, resids, rank, s = linalg.lstsq(lhs, rhs, rcond)

  # Broadcasting scale coefficients
  if c.ndim > 1:
    # For multi-dimensional output, make scale (1, order) to divide
    # across the c.T of shape (num_rhs, order)
    c = (c.T / scale[np.newaxis, :]).T

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Flatten weights: w=W.ravel().
  2. Use None or jnp.ones(len(x)) for unweighted fits instead of a scalar.
  3. Check w.ndim == 1 and w.shape[0] == y.shape[0] before the call.

Example fix

// before
c = jnp.polyfit(x, y, 3, w=W)  # W.shape == (n, 1)
// after
c = jnp.polyfit(x, y, 3, w=W.ravel())
Defensive patterns

Strategy: validation

Validate before calling

if w is not None:
    w = jnp.ravel(jnp.asarray(w))
    assert w.ndim == 1
c = jnp.polyfit(x, y, deg, w=w)

Type guard

def weights_ok(w) -> bool:
    return w is None or (getattr(w, 'ndim', 0) == 1)

Prevention

When it happens

Trigger: jnp.polyfit(x, y, deg, w=W) where W.ndim != 1 (e.g. shape (n, 1) column matrix or a per-target (n, k) weight grid); passing a scalar weight instead of a vector of ones.

Common situations: Weights loaded from data pipelines as column vectors; reusing per-channel weight matrices from other libraries (sklearn sample_weight is 1-D, but intermediate processing may add axes).

Related errors


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