jax-ml/jax · error · TypeError

convolution requires lhs and rhs ndim to be equal, got {} an

Error message

convolution requires lhs and rhs ndim to be equal, got {} and {}.

What it means

lax.conv_dimension_numbers validates that lhs and rhs have the same rank; a convolution needs matching ndim for its input and kernel. The counts are reported in the message.

Source

Thrown at jax/_src/lax/convolution.py:966

def conv_dimension_numbers(lhs_shape, rhs_shape, dimension_numbers
                           ) -> ConvDimensionNumbers:
  """Converts convolution `dimension_numbers` to a `ConvDimensionNumbers`.

  Args:
    lhs_shape: tuple of nonnegative integers, shape of the convolution input.
    rhs_shape: tuple of nonnegative integers, shape of the convolution kernel.
    dimension_numbers: None or a tuple/list of strings or a ConvDimensionNumbers
      object.

  Returns:
    A `ConvDimensionNumbers` object that represents `dimension_numbers` in the
    canonical form used by lax functions.
  """
  if isinstance(dimension_numbers, ConvDimensionNumbers):
    return dimension_numbers
  if len(lhs_shape) != len(rhs_shape):
    msg = "convolution requires lhs and rhs ndim to be equal, got {} and {}."
    raise TypeError(msg.format(len(lhs_shape), len(rhs_shape)))

  if dimension_numbers is None:
    iota = tuple(range(len(lhs_shape)))
    return ConvDimensionNumbers(iota, iota, iota)
  elif isinstance(dimension_numbers, (list, tuple)):
    if len(dimension_numbers) != 3:
      msg = "convolution dimension_numbers list/tuple must be length 3, got {}."
      raise TypeError(msg.format(len(dimension_numbers)))
    if not all(isinstance(elt, str) for elt in dimension_numbers):
      msg = "convolution dimension_numbers elements must be strings, got {}."
      raise TypeError(msg.format(tuple(map(type, dimension_numbers))))
    msg = ("convolution dimension_numbers[{}] must have len equal to the ndim "
           "of lhs and rhs, got {} for lhs and rhs shapes {} and {}.")
    for i, elt in enumerate(dimension_numbers):
      if len(elt) != len(lhs_shape):
        raise TypeError(msg.format(i, len(elt), lhs_shape, rhs_shape))

    lhs_spec, rhs_spec, out_spec = conv_general_permutations(dimension_numbers)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape/broadcast the kernel so lhs.ndim == rhs.ndim (e.g. kernel shape (k,1,in,out) for 1D conv)
2. Verify your data pipeline preserves the channel dims on both tensors
3. Build dimension_numbers with lax.conv_dimension_numbers, which surfaces the mismatch early with clear shapes

Example fix

# before
x = jnp.zeros((8, 16, 16, 3)); w = jnp.zeros((3, 3, 3))  # missing out-channel dim
out = lax.conv_general_dilated(x, w, (1,1), 'SAME')
# after
w = jnp.zeros((3, 3, 3, 8))  # (H, W, Cin, Cout)
out = lax.conv_general_dilated(x, w, (1,1), 'SAME')

assert lhs.ndim == rhs.ndim, f'rank mismatch: {lhs.ndim} vs {rhs.ndim}'

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

When it happens

Trigger: Calling lax.conv, lax.conv_general_dilated, etc. with a 4D lhs and a 3D rhs (e.g. forgetting the in/out channel dims on the kernel), or rank-mismatched dimension_numbers strings.

Common situations: Using a 1D kernel (out_channels,) instead of (out_channels, in_channels) with lax.conv_general_dilated; mixing 1D and 2D data layouts; bugs in preprocessing that drop a dimension.

Related errors

• convolve2d() only supports 2-dimensional inputs.
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must
• conv_general_dilated window and window_strides must have the
• Wrong number of explicit pads for convolution: expected {},

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