jax-ml/jax · error · TypeError

convolution dimension_numbers must be tuple/list or None, go

Error message

convolution dimension_numbers must be tuple/list or None, got {}.

What it means

conv_dimension_numbers accepts only None, a list/tuple of three strings, or a ConvDimensionNumbers instance. Any other type (dict, int, keyword string) hits the else branch and raises this TypeError with the actual type.

Source

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

    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])):
      msg = ("convolution dimension_numbers[{}] cannot have duplicate "
             "characters, got {}.")
      raise TypeError(msg.format(i, dimension_numbers[i]))
  if not (set(lhs_spec) - set(lhs_char) == set(rhs_spec) - set(rhs_char) ==
          set(out_spec) - set(out_char)):
    msg = ("convolution dimension_numbers elements must each have the same "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass None for the canonical layout, a 3-tuple of strings, or a ConvDimensionNumbers namedtuple
2. Build the namedtuple via lax.conv_dimension_numbers(None...) or lax.conv_general_permutations once and reuse it

Example fix

# before
out = lax.conv_general_dilated(x, w, (1,1), 'SAME', dimension_numbers='NCHW')
# after
out = lax.conv_general_dilated(x, w, (1,1), 'SAME', dimension_numbers=('NCHW', 'OIHW', 'NCHW'))

assert dimension_numbers is None or isinstance(dimension_numbers, (tuple, list, jax.lax.ConvDimensionNumbers)), type(dimension_numbers)

from jax._src.lax.convolution import ConvDimensionNumbers
def valid_dn(dn) -> bool:
    return dn is None or isinstance(dn, (tuple, list, ConvDimensionNumbers))

When it happens

Trigger: Passing dimension_numbers='NCHW' (a single string), a dict like {'lhs':'NCHW',...}, or an integer enum to lax.conv_general_dilated.

Common situations: Guessing the API shape from other libraries; passing an object that used to work with an older custom wrapper.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• convolution dimension_numbers elements must be strings, got
• convolution dimension_numbers list/tuple must be length 3, g
• convolution dimension_numbers[{}] must have len equal to the
• convolution dimension_numbers[{}] must contain the character
• convolution dimension_numbers[{}] cannot have duplicate char

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