jax-ml/jax · error · ValueError

correlate2d() only supports 2-dimensional inputs.

Error message

correlate2d() only supports 2-dimensional inputs.

What it means

correlate2d validates that both inputs are exactly 2-D before delegating to the N-D machinery; arrays of any other rank are rejected with this ValueError, mirroring the SciPy 2-D-only API.

Source

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

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

    >>> jax.scipy.signal.correlate2d(x, y, mode='same')
    Array([[15., 24.,  7.],
           [28., 14.,  9.],
           [ 7.,  7.,  2.]], dtype=float32)

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

    >>> jax.scipy.signal.correlate2d(x, y, mode='valid')
    Array([[15., 24.],
           [28., 14.]], dtype=float32)
  """
  if boundary != 'fill' or fillvalue != 0:
    raise NotImplementedError("correlate2d() only supports boundary='fill', fillvalue=0")
  if np.ndim(in1) != 2 or np.ndim(in2) != 2:
    raise ValueError("correlate2d() only supports 2-dimensional inputs.")

  swap = all(s1 <= s2 for s1, s2 in zip(in1.shape, in2.shape))
  same_shape =  all(s1 == s2 for s1, s2 in zip(in1.shape, in2.shape))

  if mode == "same":
    in1, in2 = jnp.flip(in1), in2.conj()
    result = jnp.flip(_convolve_nd(in1, in2, mode, precision=precision))
  elif mode == "valid":
    if swap and not same_shape:
      in1, in2 = jnp.flip(in2), in1.conj()
      result = _convolve_nd(in1, in2, mode, precision=precision)
    else:
      in1, in2 = jnp.flip(in1), in2.conj()
      result = jnp.flip(_convolve_nd(in1, in2, mode, precision=precision))
  else:
    if swap:
      in1, in2 = jnp.flip(in2), in1.conj()
      result = _convolve_nd(in1, in2, mode, precision=precision).conj()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Use jax.scipy.signal.correlate (N-D) for non-2-D inputs
  2. vmap over the batch dimension or index a single image
  3. Check np.ndim(in1) == np.ndim(in2) == 2 before calling

Example fix

// before
r = jax.scipy.signal.correlate2d(imgs, tpl)  # imgs: (B,H,W)
// after
r = jax.vmap(lambda im: jax.scipy.signal.correlate2d(im, tpl))(imgs)
Defensive patterns

Strategy: validation

Validate before calling

assert np.ndim(in1) == 2 and np.ndim(in2) == 2

Type guard

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

Prevention

When it happens

Trigger: Passing (N,H,W) batched feature maps, 1-D signals, or scalar templates to jax.scipy.signal.correlate2d.

Common situations: Template matching inside batched neural-net code without vmap; accidentally passing a list-of-lists-of-lists that becomes 3-D.

Related errors


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