jax-ml/jax · error · TypeError

convolution dimension_numbers[{}] must contain the character

Error message

convolution dimension_numbers[{}] must contain the characters '{}' and '{}' exactly once, got {}.

What it means

In conv_general_permutations, each layout string must contain its two designated non-spatial characters exactly once: lhs and out must each contain 'N' and 'C' once, rhs must contain 'O' and 'I' once. Violations (missing, duplicated, or swapped) raise this TypeError naming the element index and required characters.

Source

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

      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 "
           "set of spatial characters, got {}.")
    raise TypeError(msg.format(dimension_numbers))

  def getperm(spec, charpair):
    spatial = (i for i, c in enumerate(spec) if c not in charpair)
    if spec is not rhs_spec:
      spatial = sorted(spatial, key=lambda i: rhs_spec.index(spec[i]))
    return (spec.index(charpair[0]), spec.index(charpair[1])) + tuple(spatial)

  lhs_perm, rhs_perm, out_perm = map(getperm, dimension_numbers, charpairs)
  return lhs_perm, rhs_perm, out_perm

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use 'OIHW'/'HWIO' style layouts for the kernel: O and I exactly once
2. Ensure lhs/out layouts contain N and C exactly once
3. Prefer lax.conv_dimension_numbers and validated constants over hand-typed strings

Example fix

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

required = [('N','C'), ('O','I'), ('N','C')]
assert all(s.count(a) == s.count(b) == 1 for s, (a, b) in zip(dimension_numbers, required))

def valid_layout_chars(dn) -> bool:
    req = [('N','C'), ('O','I'), ('N','C')]
    return all(s.count(a) == 1 and s.count(b) == 1 for s, (a, b) in zip(dn, req))

When it happens

Trigger: Passing rhs layout 'OOHW' (two O's, no I) or lhs layout 'NHHW' (no C); using 'NCHW' for the kernel instead of 'OIHW'.

Common situations: Copy-pasting the input layout into the kernel position; typos in layout strings; misunderstanding that kernel layouts use O (out channels) and I (in channels).

Related errors

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

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