jax-ml/jax · error · ValueError

block_size ({len(block_size)}) and tile_size ({len(tile_size

Error message

block_size ({len(block_size)}) and tile_size ({len(tile_size)}) must have the same length.

What it means

sample_block generates random samples for one tile of a larger block, so block_size and tile_size must have the same rank (number of axes). A mismatch raises ValueError before any sampling happens.

Source

Thrown at jax/_src/pallas/mosaic/random.py:203

  Args:
    sampler_fn: A sampling function that consumes a key and returns
      random samples.
    global_key: The global key to use for sampling.
    block_size: The shape of an individual block.
    tile_size: The shape of a ``tile``, which is the smallest unit at
      which samples are generated. This should be selected to be a divisor
      of all block sizes one needs to be invariant to.
    total_size: The total size of the array to sample.
    block_index: The index denoting which block to generate keys for. Defaults
      to the program_id for each block axis.
    **kwargs: Additional arguments to pass to the sampler_fn.

  Returns:
    A ``block_size`` shaped array of samples for the current block corresponding
    to ``block_index``.
  """
  if len(block_size) != len(tile_size):
    raise ValueError(f"block_size ({len(block_size)}) and tile_size "
                     f"({len(tile_size)}) must have the same length.")

  if block_index is None:
    num_axes = len(block_size)
    block_index = tuple(
      primitives.program_id(axis) for axis in range(num_axes))

  keys = blocked_sampler.blocked_fold_in(
      global_key, total_size, block_size, tile_size, block_index)
  return blocked_sampler.sample_block(
      sampler_fn, keys, block_size, tile_size, **kwargs)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Give block_size and tile_size the same number of axes (append 1s or reshape accordingly)
2. For 2D sampling pass block_size=(rows, cols) matching tile rank

Example fix

# before
x = random.sample_block(key, block_size=(8192,), tile_size=(128, 128))
# after
x = random.sample_block(key, block_size=(64, 128), tile_size=(128, 128))

assert len(block_size) == len(tile_size), 'block/tile rank mismatch'

def ranks_match(block_size, tile_size) -> bool:
    return len(tuple(block_size)) == len(tuple(tile_size))

When it happens

Trigger: Calling random.sample_block(key, block_size=(128,), tile_size=(128, 128)) or any combination where len(block_size) != len(tile_size), e.g. 1D block with 2D tiles.

Common situations: Using 2D tiles with a flattened 1D block shape; editing one of the two size tuples during refactoring and forgetting the other.

Related errors

• PRNG keys must be loaded from SMEM. Did you set the memory s
• Can only load scalars from SMEM
• Expected value and mask to have the same shape, but got valu
• Can only store scalars to SMEM
• The number of sources must match the packing factor ({packin

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