jax-ml/jax · error · TypeError

convolution dimension_numbers elements must each have the sa

Error message

convolution dimension_numbers elements must each have the same set of spatial characters, got {}.

What it means

After removing the batch/channel characters, the remaining (spatial) character sets of the three layout strings must be identical — lhs, rhs, and out must agree on which spatial axes exist. Mismatched spatial sets raise this TypeError showing the whole dimension_numbers.

Source

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

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,
    lhs_dilation, rhs_dilation) -> list[tuple[int, int]]:
  lhs_dilated_shape = lax._dilate_shape(in_shape, lhs_dilation)
  rhs_dilated_shape = lax._dilate_shape(window_dimensions, rhs_dilation)
  out_dilated_shape = lax._dilate_shape(out_shape, window_strides)
  pad_before = np.subtract(rhs_dilated_shape, [lo for lo, _ in padding]) - 1

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use the same spatial characters (same set) in all three layout strings, only N/C vs O/I may differ
2. Regenerate all three strings together from one rank/layout decision
3. For transposed convs remember out layout spatial chars must still match lhs/rhs

Example fix

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

spatial = [set(s) - set(sc) for s, sc in zip(dn, [('NC','C'),('OI','I'),('NC','C')])]
assert spatial[0] == spatial[1] == spatial[2], 'spatial chars must match'

def matching_spatial_sets(dn) -> bool:
    sets = [set(dn[0]) - {'N','C'}, set(dn[1]) - {'O','I'}, set(dn[2]) - {'N','C'}]
    return sets[0] == sets[1] == sets[2]

When it happens

Trigger: Passing ('NCHW','OIHD','NCHW') — rhs uses 'D' where others use 'W'; or 1D lhs/out layouts with a 2D kernel layout.

Common situations: Editing only one of the three strings during a layout refactor (NHWC→NCHW); using different spatial letters per string in 3D convs.

Related errors

• convolution dimension_numbers[{}] must contain the character
• convolution dimension_numbers[{}] cannot have duplicate char
• 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/b33512cec5882e4a. Report an issue: GitHub.