jax-ml/jax · error · ValueError

convolve2d() only supports 2-dimensional inputs.

Error message

convolve2d() only supports 2-dimensional inputs.

What it means

convolve2d requires both inputs to be exactly 2-D; it is a thin wrapper over the N-D _convolve_nd but validates rank to match the SciPy 2-D API contract. Higher- or lower-rank arrays raise this.

Source

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

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

    >>> jax.scipy.signal.convolve2d(x, y, mode='same')
    Array([[22., 17.],
           [30., 32.]], dtype=float32)

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

    >>> jax.scipy.signal.convolve2d(x, y, mode='valid')
    Array([[22., 17.],
           [30., 32.]], dtype=float32)
  """
  if boundary != 'fill' or fillvalue != 0:
    raise NotImplementedError("convolve2d() only supports boundary='fill', fillvalue=0")
  if np.ndim(in1) != 2 or np.ndim(in2) != 2:
    raise ValueError("convolve2d() only supports 2-dimensional inputs.")
  return _convolve_nd(in1, in2, mode, precision=precision)


def correlate(in1: Array, in2: Array, mode: ModeString = 'full', method: str = 'auto',
              precision: PrecisionLike = None) -> Array:
  """Cross-correlation of two N-dimensional arrays.

  JAX implementation of :func:`scipy.signal.correlate`.

  Args:
    in1: left-hand input to the cross-correlation.
    in2: right-hand input to the cross-correlation. Must have ``in1.ndim == in2.ndim``.
    mode: controls the size of the output. Available operations are:

      * ``"full"``: (default) output the full cross-correlation of the inputs.
      * ``"same"``: return a centered portion of the ``"full"`` output which
        is the same size as ``in1``.
      * ``"valid"``: return the portion of the ``"full"`` output which do not

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use jax.scipy.signal.convolve for N-D inputs (it supports any matching rank)
2. squeeze extra axes: convolve2d(x[0], k) for batched data, or vmap over the batch
3. Reshape 1-D inputs to (1, n) or (n, 1) if 2-D semantics are wanted

Example fix

// before
y = jax.scipy.signal.convolve2d(batched_x, k)  # batched_x: (8,H,W)
// after
y = jax.vmap(lambda im: jax.scipy.signal.convolve2d(im, k))(batched_x)

assert np.ndim(in1) == 2 and np.ndim(in2) == 2, 'convolve2d needs 2-D inputs'

def is_2d(a) -> bool:
    return jnp.asarray(a).ndim == 2

When it happens

Trigger: Passing a batched (N,H,W) tensor, a 1-D vector, or a scalar alongside a 2-D kernel to jax.scipy.signal.convolve2d.

Common situations: Adding a batch dimension for vmap/jit pipelines then forgetting to drop it before calling convolve2d; using convolve2d where convolve (N-D) was intended.

Related errors

• convolution requires lhs and rhs ndim to be equal, got {} an
• zero-size arrays not supported in convolutions, got shapes {
• correlate2d() only supports 2-dimensional inputs.
• two-arguments must have the same rank ({x.ndim} vs {y.ndim})
• conv_general_dilated batch_group_count must divide lhs batch

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