jax-ml/jax · error · TypeError

convolution dimension_numbers

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.

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'))
Defensive patterns

Strategy: validation

Validate before calling

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

Type guard

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

Prevention

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

Understand the failure class

Background: "Invalid ... format", "must be in format X", "does not look like a ..." — invalid argument format errors across CLI tools and libraries — this error's family across 17 libraries.

Related errors

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

Appendix: 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)