jax-ml/jax · error · ValueError

in1 and in2 must have the same number of dimensions

Error message

in1 and in2 must have the same number of dimensions

What it means

Guard at the entry of _convolve_nd: the convolution is evaluated elementwise over paired axes of the two inputs, so in1 and in2 must have the same number of dimensions. Different dimensionalities cannot be paired and fail.

Solutions

Reshape to match ranks: kernel[None, :] for 2D, kernel[None, None] for batched 3D
Use jnp.squeeze on stray unit dims deliberately
Prefer fftconvolve with axes= for selective convolution

Example fix

# before
signal.convolve2d(img, kernel1d, mode='same')
# after
signal.convolve2d(img, kernel1d[None, :], mode='same')

Defensive patterns

Strategy: validation

Validate before calling

assert in1.ndim == in2.ndim, (in1.shape, in2.shape)

Type guard

def same_rank(a, b) -> bool: return jnp.asarray(a).ndim == jnp.asarray(b).ndim

Prevention

Use a helper that reshapes kernels to input rank automatically

When it happens

Trigger: convolve2d expects 2D inputs; passing a (H,W) image with a (K,) kernel, or a batched (N,H,W) with a 2D filter.

Common situations: 1D-vs-2D kernel confusion, batched inputs forgetting the leading dim on the filter.

Related errors

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

Appendix: source

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

    out_shape = in1.shape
  elif mode == "valid":
    out_shape = tuple(s1 - s2 + 1 for s1, s2 in zip(in1.shape, in2.shape))
  else:
    raise ValueError(f"Unrecognized {mode=}")

  start_indices = tuple((full_size - out_size) // 2
                        for full_size, out_size in zip(full_shape, out_shape))
  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':

View on GitHub (pinned to 1e1c6a8fc0)