jax-ml/jax · error · ValueError

in1 and in2 should have the same dimensionality

Error message

in1 and in2 should have the same dimensionality

What it means

Validation in jax.scipy.signal.fftconvolve: the FFT-based convolution is computed over matched axes of both inputs, so in1 and in2 must have the same number of dimensions. Mixing dimensionalities (e.g., 1-D with 2-D) fails before any transform is taken.

Solutions

  1. Reshape the smaller operand to match rank, e.g. kernel[None, :] for 2D, or kernel[None] for batched inputs
  2. Use axes= parameter to convolve only specific shared axes
  3. Squeeze irrelevant unit dims deliberately

Example fix

# before
out = signal.fftconvolve(img2d, kernel1d, mode='same')
# after
out = signal.fftconvolve(img2d, kernel1d[None, :], mode='same')
Defensive patterns

Strategy: validation

Validate before calling

assert jnp.asarray(in1).ndim == jnp.asarray(in2).ndim, (in1.shape, in2.shape)

Type guard

def same_rank(a, b) -> bool: return jnp.asarray(a).ndim == jnp.asarray(b).ndim

Prevention

When it happens

Trigger: Convolving a 1D kernel with a 2D image, or a batched (N,H,W) input with an (H,W) filter.

Common situations: Applying 1D smoothing kernels to 2D signals without reshaping; mixing batched and unbatched operands.

Related errors


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

Appendix: source

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

    Specifying ``mode = 'same'`` returns a centered convolution the same size
    as the first input:

    >>> with jax.numpy.printoptions(precision=3):
    ...   print(jax.scipy.signal.fftconvolve(x, y, mode='same'))
    [3. 6. 7. 6. 3.]

    Specifying ``mode = 'valid'`` returns only the portion where the two arrays
    fully overlap:

    >>> 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.

View on GitHub (pinned to 1e1c6a8fc0)