jax-ml/jax · error · ValueError

Length of major_to_minor and the rank of the value should ma

Error message

Length of major_to_minor and the rank of the value should match. Got major_to_minor={self.major_to_minor} and shape={aval_shape}

What it means

A jax Layout's major_to_minor permutation must have exactly as many entries as the array's rank. check_compatible_aval enforces this before using the layout with a value of a given shape.

Source

Thrown at jax/_src/layout.py:122

                  kwargs['sub_byte_element_size_in_bits'])

  def _to_xla_layout(self, dtype) -> xc.Layout:
    if self.tiling is None:
      xla_layout = xc.Layout(self.major_to_minor[::-1])
    else:
      if self.sub_byte_element_size_in_bits != 0:
        sub_byte_size = self.sub_byte_element_size_in_bits
      elif issubdtype(dtype, np.integer):
        sub_byte_size = iinfo(dtype).bits if iinfo(dtype).bits < 8 else 0
      else:
        sub_byte_size = 0
      xla_layout = xc.Layout(self.major_to_minor[::-1], self.tiling,
                              sub_byte_size)
    return xla_layout

  def check_compatible_aval(self, aval_shape: Shape):
    if len(self.major_to_minor) != len(aval_shape):
      raise ValueError(
          f'Length of major_to_minor and the rank of the value should match.'
          f' Got major_to_minor={self.major_to_minor} and shape={aval_shape}')


LayoutOptions = Layout | None | AutoLayoutSingleton
ShardingOptions = Sharding | None


class Format:
  __slots__ = ['layout', 'sharding']

  def __init__(self, layout: LayoutOptions = None,
               sharding: ShardingOptions = None):
    # If layout is concrete and sharding is not, error.
    if isinstance(layout, Layout) and sharding is None:
      raise ValueError(
          'Sharding has to be concrete when layout is of type'
          f' {type(layout)}. Please pass a'

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Build major_to_minor with length equal to the target rank, e.g. list(range(ndim))[::-1]
  2. Derive the layout from the array: use ndim of the actual value when constructing

Example fix

# before
layout = Layout(major_to_minor=(1,0))
layout.check_compatible_aval(x.shape)  # x is 3-D
# after
layout = Layout(major_to_minor=tuple(range(x.ndim)[::-1]))
layout.check_compatible_aval(x.shape)
Defensive patterns

Strategy: validation

Validate before calling

assert len(layout.major_to_minor) == len(shape), f'{len(layout.major_to_minor)} vs rank {len(shape)}'

Type guard

def layout_fits(layout, shape): return len(layout.major_to_minor) == len(shape)

Prevention

When it happens

Trigger: Creating a Layout(major_to_minor=[0,1]) and applying it to a 3-D array's abstract value.

Common situations: Reusing a layout object across tensors of different rank in sharding/layout APIs; constructing layouts from hardcoded permutations.

Related errors


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