jax-ml/jax · error · ValueError

mapped axes must have same shape; got {in1.shape=} {in2.shap

Error message

mapped axes must have same shape; got {in1.shape=} {in2.shape=} {axes=}

What it means

When fftconvolve is called with explicit axes, the remaining (mapped/batched) axes are vmap-ed over and must match in size between in1 and in2. Otherwise the batched convolution is ill-defined.

Source

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

    >>> with jax.numpy.printoptions(precision=3):
    ...   print(jax.scipy.signal.fftconvolve(x, y, mode='valid'))
    [6. 7. 6.]
  """
  check_arraylike('fftconvolve', in1, in2)
  in1, in2 = promote_dtypes_inexact(in1, in2)
  if in1.ndim != in2.ndim:
    raise ValueError("in1 and in2 should have the same dimensionality")
  if mode not in ["same", "full", "valid"]:
    raise ValueError("mode must be one of ['same', 'full', 'valid']")
  _fftconvolve = partial(_fftconvolve_unbatched, mode=mode)
  if axes is None:
    return _fftconvolve(in1, in2)
  axes = _ensure_index_tuple(axes)
  axes = tuple(canonicalize_axis(ax, in1.ndim) for ax in axes)
  mapped_axes = set(range(in1.ndim)) - set(axes)
  if any(in1.shape[i] != in2.shape[i] for i in mapped_axes):
    raise ValueError(f"mapped axes must have same shape; got {in1.shape=} {in2.shape=} {axes=}")
  for ax in sorted(mapped_axes):
    _fftconvolve = api.vmap(_fftconvolve, in_axes=ax, out_axes=ax)
  return _fftconvolve(in1, in2)

def _fftconvolve_unbatched(in1: Array, in2: Array, mode: str) -> Array:
  full_shape = tuple(s1 + s2 - 1 for s1, s2 in zip(in1.shape, in2.shape))

  # TODO(jakevdp): potentially use next_fast_len to evaluate with a more efficient shape.
  fft_shape = full_shape  # tuple(next_fast_len(s) for s in full_shape)

  if mode == 'valid':
    no_swap = all(s1 >= s2 for s1, s2 in zip(in1.shape, in2.shape))
    swap = all(s1 <= s2 for s1, s2 in zip(in1.shape, in2.shape))
    if not (no_swap or swap):
      raise ValueError("For 'valid' mode, One input must be at least as "
                       "large as the other in every dimension.")
    if swap:
      in1, in2 = in2, in1

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make batch (non-axes) dims equal: slice, pad, or broadcast one input
2. Alternatively vmap manually if you actually want per-sample behavior
3. Double-check shapes after preprocessing pipelines

Example fix

# before
signal.fftconvolve(a, b, axes=1)  # a:(8,100), b:(16,100)
# after
b = b[:8]  # align batch dims
signal.fftconvolve(a, b, axes=1)

mapped = set(range(in1.ndim)) - set(axes)
assert all(in1.shape[i] == in2.shape[i] for i in mapped)

def batch_axes_match(in1, in2, axes) -> bool:
    m = set(range(in1.ndim)) - set(axes)
    return all(in1.shape[i] == in2.shape[i] for i in m)

When it happens

Trigger: Convolving (8, 100) and (16, 100) along axes=1 — batch dims 8 vs 16 differ.

Common situations: Batched signals where one side has a different batch size after slicing or padding.

Related errors

• in1 and in2 should have the same dimensionality
• For 'valid' mode, One input must be at least as large as the
• in1 and in2 must have the same number of dimensions
• Incompatible shapes for Sylvester equation: A: {A.shape} B:
• coordinates must be a sequence of length input.ndim, but {}

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