jax-ml/jax · error · NotImplementedError

correlate2d() only supports boundary='fill', fillvalue=0

Error message

correlate2d() only supports boundary='fill', fillvalue=0

What it means

correlate2d in JAX only supports the default zero-fill boundary; 'wrap', 'symm', and nonzero fill values are not implemented (NotImplementedError, not ValueError — the feature is absent rather than invalid).

Source

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

           [12.,  7.,  7.,  2.]], dtype=float32)

    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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pre-pad inputs with jnp.pad using the desired mode and use mode='valid'
  2. Use jax.lax.conv_general_dilated directly for full padding control
  3. Fall back to scipy on CPU for that op if exact boundary semantics are required

Example fix

// before
r = jax.scipy.signal.correlate2d(img, tpl, mode='same', boundary='symm')
// after
ph, pw = tpl.shape[0]//2, tpl.shape[1]//2
img_p = jnp.pad(img, ((ph,ph),(pw,pw)), mode='symmetric')
r = jax.scipy.signal.correlate2d(img_p, tpl, mode='valid')
Defensive patterns

Strategy: fallback

Validate before calling

def corr2d_symm(x, tpl):
    ph, pw = tpl.shape[0] // 2, tpl.shape[1] // 2
    xp = jnp.pad(x, ((ph, ph), (pw, pw)), mode='symmetric')
    return jax.scipy.signal.correlate2d(xp, tpl, mode='valid')

Prevention

When it happens

Trigger: jax.scipy.signal.correlate2d(x, template, boundary='wrap') or fillvalue != 0, typical in template matching with periodic images.

Common situations: Porting scipy.signal.correlate2d normalized cross-correlation pipelines that use reflecting boundaries.

Related errors


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