jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

convolve2d in JAX only implements zero-padding (boundary='fill', fillvalue=0); unlike SciPy it does not support 'wrap', 'reflect', or nonzero fill values, because the underlying lax convolution only exposes explicit padding. Passing anything else raises NotImplementedError.

Source

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

           [13., 30., 32., 20.],
           [ 3., 13., 18.,  8.]], dtype=float32)

    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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use boundary='fill', fillvalue=0 and manually pad the input with jnp.pad(mode='wrap'/'symmetric') to the equivalent extent, then use mode='valid'
2. Use jax.lax.conv_general_dilated with a custom padding config for asymmetric padding
3. Keep that operation in scipy/numpy on CPU if boundary semantics are essential

Example fix

// before
y = jax.scipy.signal.convolve2d(x, k, boundary='wrap', mode='same')
// after
xp = jnp.pad(x, ((k.shape[0]//2,)*2, (k.shape[1]//2,)*2), mode='wrap')
y = jax.scipy.signal.convolve2d(xp, k, mode='valid')

def conv2d_wrap(x, k, mode='same'):
    # emulate wrap boundary via explicit padding
    ph, pw = k.shape[0] // 2, k.shape[1] // 2
    xp = jnp.pad(x, ((ph, ph), (pw, pw)), mode='wrap')
    return jax.scipy.signal.convolve2d(xp, k, mode='valid')

When it happens

Trigger: jax.scipy.signal.convolve2d(x, k, boundary='wrap'); boundary='symm'; or fillvalue=1.0 — common in image-processing code ported from scipy.signal.

Common situations: Migrating scipy.signal.convolve2d image filters (e.g. Laplacian with symmetric boundaries) to JAX.

Related errors

• correlate2d() only supports boundary='fill', fillvalue=0
• For 'valid' mode, One input must be at least as large as the
• mode must be one of ['full', 'same', 'valid']
• in1 and in2 must have the same number of dimensions
• zero-size arrays not supported in convolutions, got shapes {

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