jax-ml/jax · error · TypeError

convolution dimension_numbers elements must be strings, got

Error message

convolution dimension_numbers elements must be strings, got {}.

What it means

Each of the three dimension_numbers elements must be a string (like 'NCHW'). If any element is another type (int, list, None), conv_dimension_numbers raises this TypeError showing the actual types of the tuple elements.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use layout strings such as ('NCHW', 'OIHW', 'NCHW')
2. If you have axis permutations, convert them to layout characters first or use the ConvDimensionNumbers namedtuple form via lax.conv_general_permutations output

Example fix

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

assert all(isinstance(elt, str) for elt in dimension_numbers), 'layouts must be strings'

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

When it happens

Trigger: Passing dimension_numbers=(('N','C','H','W'), ('O','I','H','W'), ('N','C','H','W')) (tuples instead of strings), or integer axis permutations.

Common situations: Assuming dimension_numbers takes index tuples like other JAX APIs; converting code from string layouts to permutation lists incorrectly.

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 must be tuple/list or None, go
• 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/dd770c5daa69f8af. Report an issue: GitHub.