jax-ml/jax · error · ValueError

two-arguments must have the same rank ({x.ndim} vs {y.ndim})

Error message

two-arguments must have the same rank ({x.ndim} vs {y.ndim}).

What it means

For cross-spectral density, x and y must have the same number of dimensions so their outer (non-transform) axes can be paired and broadcast. A rank mismatch (e.g. 1-D vs 2-D) makes segment-wise pairing ill-defined, so it is rejected. Note the message has an f-string bug: values are not interpolated in some JAX versions.

Source

Thrown at jax/_src/scipy/signal.py:668

    raise ValueError(
        f"Unknown boundary option '{boundary}', "
        f"must be one of: {list(boundary_funcs.keys())}")

  axis = core.concrete_or_error(operator.index, axis, "axis of windowed-FFT")
  axis = canonicalize_axis(axis, x.ndim)

  if y is None:
    check_arraylike('spectral_helper', x)
    x, = promote_dtypes_inexact(x)
    y_arr = x  # place-holder for type checking
    outershape = tuple_delete(x.shape, axis)
  else:
    if mode != 'psd':
      raise ValueError("two-argument mode is available only when mode=='psd'")
    check_arraylike('spectral_helper', x, y)
    x, y_arr = promote_dtypes_inexact(x, y)
    if x.ndim != y_arr.ndim:
      raise ValueError("two-arguments must have the same rank ({x.ndim} vs {y.ndim}).")
    # Check if we can broadcast the outer axes together
    try:
      outershape = jnp.broadcast_shapes(tuple_delete(x.shape, axis),
                                        tuple_delete(y_arr.shape, axis))
    except ValueError as err:
      raise ValueError('x and y cannot be broadcast together.') from err

  result_dtype = dtypes.to_complex_dtype(x.dtype)
  freq_dtype = np.finfo(result_dtype).dtype

  nperseg_int: int = 0
  nfft_int: int = 0
  noverlap_int: int = 0

  if nperseg is not None:  # if specified by user
    nperseg_int = core.concrete_or_error(
        int, nperseg, "nperseg of windowed-FFT")
    if nperseg_int < 1:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Make ranks equal: add a leading axis with x[None] or remove one with y[0] / squeeze
  2. Broadcast explicitly to a common outer shape before calling csd
  3. Verify x.ndim == y.ndim in a precondition check

Example fix

// before
f, P = jax.scipy.signal.csd(x, Y)  # x:(N,), Y:(C,N)
// after
f, P = jax.scipy.signal.csd(jnp.broadcast_to(x, Y.shape), Y)
Defensive patterns

Strategy: type-guard

Validate before calling

if x.ndim != y.ndim:
    if x.ndim < y.ndim:
        x = x[None]  # promote, adjust to your layout
    else:
        y = y[None]

Type guard

def same_rank(a, b):
    a, b = jnp.asarray(a), jnp.asarray(b)
    if a.ndim == b.ndim:
        return a, b
    raise ValueError(f'rank mismatch: {a.ndim} vs {b.ndim}')

Prevention

When it happens

Trigger: csd(x, y) where x is shape (1000,) and y is (1, 1000) or (32, 1000); mixing a single channel with a multichannel signal.

Common situations: Comparing a reference signal against a batch/channels-first array; forgetting that a leading axis adds a rank.

Related errors


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