jax-ml/jax · error · TypeError

convolution dimension_numbers[{}] cannot have duplicate char

Error message

convolution dimension_numbers[{}] cannot have duplicate characters, got {}.

What it means

Beyond containing the right special characters, each layout string must not repeat any character (its length must equal the number of distinct characters). Duplicate characters make the permutation ambiguous and raise this TypeError.

Source

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

    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


def _conv_general_vjp_lhs_padding(
    in_shape, window_dimensions, window_strides, out_shape, padding,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Use distinct spatial characters per axis, e.g. 'NCHWD' / 'OIHWD' for 3D convs
  2. Construct layouts programmatically from the known rank to guarantee uniqueness

Example fix

# before
dn = lax.conv_dimension_numbers(x.shape, w.shape, ('NCHHH', 'OIHHH', 'NCHHH'))
# after
dn = lax.conv_dimension_numbers(x.shape, w.shape, ('NCHWD', 'OIHWD', 'NCHWD'))
Defensive patterns

Strategy: validation

Validate before calling

assert all(len(s) == len(set(s)) for s in dimension_numbers), 'duplicate characters in layout'

Type guard

def no_dup_chars(dn) -> bool:
    return all(len(s) == len(set(s)) for s in dn)

Prevention

When it happens

Trigger: Passing a layout like 'NCHH' or 'OIWW' where a spatial character repeats instead of using distinct spatial letters.

Common situations: 3D convs where devs reuse 'H' for both height and depth; auto-generated layout strings from loops that collide on letters.

Related errors


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