jax-ml/jax · error · ValueError

two-argument mode is available only when mode=='psd'

Error message

two-argument mode is available only when mode=='psd'

What it means

Cross-spectral computation (two arrays x and y) is only meaningful for power spectral density estimation; STFT mode processes a single signal. Passing both x and y while mode is 'stft' violates that contract and raises this internal ValueError.

Source

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

  }

  # Check/ normalize inputs
  if boundary not in boundary_funcs:
    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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use jax.scipy.signal.csd(x, y) for two-signal spectral analysis
2. For STFT of two signals, call stft on each and combine the complex outputs

Example fix

// before
_ = _spectral_helper(x, y, fs, mode='stft', ...)  # internal misuse
// after
Zx = jax.scipy.signal.stft(x)[2]
Zy = jax.scipy.signal.stft(y)[2]
cross = Zx * jnp.conj(Zy)

assert mode == 'psd' or y is None  # two-argument form only for psd

When it happens

Trigger: Calling _spectral_helper(x, y, mode='stft', ...) with a non-None y. Public stft() never passes y; only csd (mode='psd') does.

Common situations: Using or copying jax internals to build custom transforms; monkeypatching stft to accept two signals.

Related errors

• Unknown value for mode {mode}, must be one of: ('psd', 'stft
• Unrecognized {mode=}
• Accumulator aval mismatch: expected {aval}, got {acc.aval}
• {name} cannot accept args with unreduced_kind={mat.unreduced
• Cannot determine the ``__name__`` of the caller.

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