jax-ml/jax · error · TypeError

convolution dimension_numbers

Error message

convolution dimension_numbers[{}] must have len equal to the ndim of lhs and rhs, got {} for lhs and rhs shapes {} and {}.

What it means

Each layout string in dimension_numbers must have length equal to the ndim of lhs and rhs. Since lhs/rhs ranks must match (checked just before), any element longer or shorter than the tensor rank raises this TypeError with the offending index, string length, and both shapes.

Solutions

1. Make each layout string length equal len(lhs_shape) (e.g. 'NHC' for rank-3, 'NHWB C'-style 5 chars for rank-5)
2. Use distinct spatial characters (e.g. 'NHW' vs 'OID' for width/depth) once ranks grow
3. Log shapes and layout strings together when debugging

Example fix

# before
x = jnp.zeros((8, 28, 3)); w = jnp.zeros((3, 3, 8))
dn = lax.conv_dimension_numbers(x.shape, w.shape, ('NHWC', 'HWIO', 'NHWC'))
# after
dn = lax.conv_dimension_numbers(x.shape, w.shape, ('NHC', 'HIO', 'NHC'))

assert all(len(s) == lhs.ndim for s in dimension_numbers), 'layout length must equal tensor rank'

def layouts_match_rank(dn, x) -> bool:
    return all(len(s) == x.ndim for s in dn)

When it happens

Trigger: Passing ('NHWC','HWIO','NHWC') with 3D tensors, or a 5D layout string for 4D tensors, to lax.conv_dimension_numbers / conv_general_dilated.

Common situations: Reusing 2D conv layout strings after switching to 1D or 3D data; adding/removing a channel dim without updating the strings.

Related errors

• conv_general_dilated batch_group_count must divide lhs…
• conv_general_dilated rhs output feature dimension size must…
• conv_general_dilated window and window_strides must have…
• convolution dimension_numbers
• convolution dimension_numbers elements must be strings, got

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

Appendix: source

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

  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)
    return ConvDimensionNumbers(lhs_spec, rhs_spec, out_spec)
  else:
    msg = "convolution dimension_numbers must be tuple/list or None, got {}."
    raise TypeError(msg.format(type(dimension_numbers)))


def conv_general_permutations(dimension_numbers):
  """Utility for convolution dimension permutations relative to Conv HLO."""
  lhs_spec, rhs_spec, out_spec = dimension_numbers
  lhs_char, rhs_char, out_char = charpairs = ("N", "C"), ("O", "I"), ("N", "C")
  for i, (a, b) in enumerate(charpairs):
    if not dimension_numbers[i].count(a) == dimension_numbers[i].count(b) == 1:
      msg = ("convolution dimension_numbers[{}] must contain the characters "
             "'{}' and '{}' exactly once, got {}.")
      raise TypeError(msg.format(i, a, b, dimension_numbers[i]))
    if len(dimension_numbers[i]) != len(set(dimension_numbers[i])):

View on GitHub (pinned to 1e1c6a8fc0)