jax-ml/jax · error · TypeError

convolution dimension_numbers list/tuple must be length 3, g

Error message

convolution dimension_numbers list/tuple must be length 3, got {}.

What it means

When dimension_numbers is given as a list/tuple, it must be a 3-element tuple of strings: (lhs_layout, rhs_layout, out_layout) e.g. ('NCHW','OIHW','NCHW'). Anything with a different length raises this TypeError.

Source

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

      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)
    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."""

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Provide exactly three layout strings: input, kernel, output
2. Or pass None to get the default canonical layout
3. Or pass an already-built ConvDimensionNumbers namedtuple

Example fix

# before
dn = lax.conv_dimension_numbers(x.shape, w.shape, ('NCHW', 'OIHW'))
# after
dn = lax.conv_dimension_numbers(x.shape, w.shape, ('NCHW', 'OIHW', 'NCHW'))

assert isinstance(dimension_numbers, (tuple, list)) and len(dimension_numbers) == 3, 'need (lhs, rhs, out) layouts'

def valid_dn_tuple(dn) -> bool:
    return isinstance(dn, (tuple, list)) and len(dn) == 3 and all(isinstance(s, str) for s in dn)

When it happens

Trigger: Passing dimension_numbers=('NCHW','OIHW') (missing output spec) or a 4-tuple to lax.conv_general_dilated or lax.conv_dimension_numbers.

Common situations: Omitting the output layout assuming it's inferred; copy-paste truncation; passing a ConvDimensionNumbers namedtuple unpacked incorrectly.

Related errors

• At most one of batch_group_count and feature_group_count may
• convolution dimension_numbers elements must be strings, got
• convolution dimension_numbers[{}] must have len equal to the
• convolution dimension_numbers must be tuple/list or None, go
• convolution dimension_numbers[{}] must contain the character

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