jax-ml/jax · error · ValueError

One input must be smaller than the other in every dimension.

Error message

One input must be smaller than the other in every dimension.

What it means

Raised when neither input is uniformly >= the other across all dimensions, i.e. the shapes cross (in1 bigger in one dim, smaller in another). The implementation must order inputs as larger-then-smaller before handing off to lax.conv_general_dilated, which requires an unambiguous ordering; SciPy supports mixed cases but JAX does not.

Source

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

  return lax.dynamic_slice(conv, start_indices, out_shape)


# Note: we do not reuse the code from jax.numpy.convolve here, because the handling
# of padding differs slightly between the two implementations (particularly for
# mode='same').
def _convolve_nd(in1: Array, in2: Array, mode: ModeString, *, precision: PrecisionLike) -> Array:
  if mode not in ["full", "same", "valid"]:
    raise ValueError("mode must be one of ['full', 'same', 'valid']")
  if in1.ndim != in2.ndim:
    raise ValueError("in1 and in2 must have the same number of dimensions")
  if in1.size == 0 or in2.size == 0:
    raise ValueError(f"zero-size arrays not supported in convolutions, got shapes {in1.shape} and {in2.shape}.")
  in1, in2 = promote_dtypes_inexact(in1, in2)

  no_swap = all(s1 >= s2 for s1, s2 in zip(in1.shape, in2.shape))
  swap = all(s1 <= s2 for s1, s2 in zip(in1.shape, in2.shape))
  if not (no_swap or swap):
    raise ValueError("One input must be smaller than the other in every dimension.")

  shape_o = in2.shape
  if swap:
    in1, in2 = in2, in1
  shape = in2.shape
  in2 = jnp.flip(in2)

  if mode == 'valid':
    padding = [(0, 0) for s in shape]
  elif mode == 'same':
    padding = [(s - 1 - (s_o - 1) // 2, s - s_o + (s_o - 1) // 2)
               for (s, s_o) in zip(shape, shape_o)]
  elif mode == 'full':
    padding = [(s - 1, s - 1) for s in shape]
  else:
    raise ValueError(f'unsupported mode: {mode}')

  strides = tuple(1 for s in shape)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pad the smaller array so one input dominates in every dimension, then crop the result
  2. Redesign kernel dimensions so the kernel is <= the signal in all axes
  3. Compute per-dimension separable convolutions if the kernel factorizes

Example fix

// before
y = jax.scipy.signal.convolve(x, k)  # x:(5,3), k:(4,4)
// after
k_p = jnp.pad(k, ((0,1),(0,0)))  # (5,4)
y = jax.scipy.signal.convolve(x, k_p)[: x.shape[0]+k.shape[0]-1, : x.shape[1]+k.shape[1]-1]
Defensive patterns

Strategy: validation

Validate before calling

def shapes_nested(a, b):
    return all(s1 >= s2 for s1, s2 in zip(a.shape, b.shape)) or \
           all(s1 <= s2 for s1, s2 in zip(a.shape, b.shape))
assert shapes_nested(x, k), 'one input must dominate in every dim'

Prevention

When it happens

Trigger: convolve(jnp.ones((5,3)), jnp.ones((4,4))) — 5>=4 in dim 0 but 3<4 in dim 1; any N-D convolution where shapes are not nested.

Common situations: Porting SciPy signal code that mixed kernel/data sizes; kernels sized per-dimension from unrelated config values.

Related errors


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