jax-ml/jax · error · ValueError

Shape {shape} and strides {strides} must have the same lengt

Error message

Shape {shape} and strides {strides} must have the same length

What it means

Tiling.tile_nested_shape_strides requires parallel shape and strides sequences of the same length; a length mismatch means the caller paired a shape with strides from a different rank/layout.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:208

      tiles.append(t)
    return Tiling(tuple(tiles))

  def tile_nested_shape_strides(
      self,
      shape: tuple[tuple[int, ...], ...],
      strides: tuple[tuple[int, ...], ...],
  ) -> tuple[tuple[tuple[int, ...], ...], tuple[tuple[int, ...], ...]]:
    """A fused version of `tile_shape` and `tile_strides` for nested shapes.

    By nested shape we mean that each logical dimension (i.e. each element of
    shape/strides) is actually composed out of multiple physical dimensions.
    For example, a row-major array of logical shape (128, 128) that is tiled
    into (64, 64) tiles would have a nested shape ((2, 64), (2, 64)) (i.e. each
    dim is split into two sub-dims) and nested strides of
    ((2 * 64 * 64, 64), (64 * 64, 1)).
    """
    if len(shape) != len(strides):
      raise ValueError(
          f"Shape {shape} and strides {strides} must have the same length"
      )
    def fail_if(cond, shape=shape):  # Capture shape now.
      if cond:
        raise ValueError(f"Tiling {self.tiles} does not apply to shape {shape}")
    for tile in self.tiles:
      fail_if(len(tile) > len(shape))
      untiled_shape, tiled_shape = shape[:-len(tile)], shape[-len(tile):]
      untiled_strides, tiled_strides = strides[:-len(tile)], strides[-len(tile):]
      major_dim_shapes, major_dim_strides = [], []
      minor_dim_shapes, minor_dim_strides = [], []
      for t, dim_shape, dim_strides in zip(tile, tiled_shape, tiled_strides):
        major_dim_shape_rev, major_dim_stride_rev = [], []
        minor_dim_shape_rev, minor_dim_stride_rev = [], []
        for d, s in zip(reversed(dim_shape), reversed(dim_strides), strict=True):
          if d < t:  # We will need to tile more dims
            fail_if(t % d != 0)
            t //= d

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Recompute strides from the same shape/rank you pass in (keep them derived together)
2. Assert len equality before calling
3. Use the library helpers that produce shape+strides jointly

Example fix

// before
shape, strides = (2, 64), compute_strides((2, 64, 64))
tiling.tile_nested_shape_strides(shape, strides)
// after
assert len(shape) == len(strides)
tiling.tile_nested_shape_strides(shape, strides)

assert len(shape) == len(strides), 'shape/strides rank mismatch'

def matching_rank(shape, strides) -> bool:
    return len(shape) == len(strides)

When it happens

Trigger: Calling tile_nested_shape_strides(shape, strides) where len(shape) != len(strides) — e.g. passing strides computed from a different (pre- or post-reshape) layout, or forgetting the stride for one dim.

Common situations: Hand-assembling tiled transfer descriptors (transfer_tiled / plan_tiled_transfer) with mismatched metadata; shape changes after strides were computed.

Related errors

• Tiles must have a decreasing rank
• Tiles must not be empty
• Tile shape must only have positive sizes, got: {self.tiles}
• Tiling {self.tiles} does not apply to shape {orig_shape}
• shape {orig_shape} is not a valid result of applying tiling

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