jax-ml/jax · error · ValueError

Shape dimension {shape[-2:]} must be divisible by {block_siz

Error message

Shape dimension {shape[-2:]} must be divisible by {block_size}

What it means

The TPU Philox Pallas kernel tiles the two trailing dimensions of the output shape by a fixed block size; philox_4x32_count validates that shape[-2] and shape[-1] are divisible by block_size[-2] and block_size[-1] before computing the grid.

Source

Thrown at jax/experimental/pallas/ops/tpu/random/philox.py:107

      output tensor if it were not padded. This is important for indexing
      calculations within the kernel. If `shape` is not padded, then this
      should be equal to `shape`.
    block_size: The block size of the kernel.
    offset: An optional offset to the counts.
    fuse_output: Whether to fuse the output bits into a single value.

  Returns:
    A tensor of random bits of shape `shape` if fuse_output=True. Otherwise,
    this will return a tensor of shape (2, *shape) with the first channel being
    the high bits and the second channel being the low bits.
  """
  shape = tuple(shape)
  if np.prod(shape) > jnp.iinfo(jnp.uint32).max:
    raise ValueError(
        f"Shape too large: {np.prod(shape)} > {np.iinfo(jnp.uint32).max}")

  if (shape[-2] % block_size[-2] != 0) or (shape[-1] % block_size[-1] != 0):
    raise ValueError(
        f"Shape dimension {shape[-2:]} must be divisible by {block_size}")
  grid_dims = shape[:-2] + (
      shape[-2] // block_size[-2], shape[-1] // block_size[1],)
  offset = jnp.array(offset, dtype=jnp.uint32)
  if offset.ndim != 0:
    raise ValueError(f"Offset must be scalar, got {offset.shape}")
  offset = jnp.reshape(offset, (1,))

  def kernel(offset_ref, key_ref, out_ref):
    counts_idx = tuple(pl.program_id(i) for i in range(len(grid_dims)))
    offset = prng_utils.compute_scalar_offset(
        counts_idx, unpadded_shape, block_shape)
    counts_lo = prng_utils.blocked_iota(block_size, unpadded_shape)
    counts_lo = counts_lo + offset.astype(jnp.uint32) + offset_ref[0]
    counts_lo = counts_lo.astype(jnp.uint32)
    # TODO(justinfu): Support hi bits on count.
    _zeros = jnp.zeros_like(counts_lo)
    k1 = jnp.reshape(key_ref[0, 0], (1, 1))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pad trailing dimensions up to a multiple of the block size and slice after
Use the higher-level JAX random API (jax.random.bits/PRNGKeys) which handles padding
Check the kernel's declared block_size and request shapes aligned to it

Example fix

// before
bits = philox_random_bits(key, 32, (128, 100))
// after
padded = philox_random_bits(key, 32, (128, 128))
bits = padded[:, :100]

Defensive patterns

Strategy: validation

Validate before calling

bs = block_size  # kernel's block size
assert shape[-2] % bs[-2] == 0 and shape[-1] % bs[-1] == 0, 'pad trailing dims to block multiples'
padded = tuple(-(-s % b) * b for s, b in zip(shape[-2:], bs))

Prevention

Prefer public random APIs that pad internally
Keep trailing dims multiples of the kernel block size when calling Pallas PRNG kernels directly

When it happens

Trigger: Requesting random bits with last dimensions not matching the kernel's block size (e.g. shape (128, 100) when blocks are 128x128), i.e. trailing dims that are not multiples of the block.

Common situations: Generating arbitrarily shaped random tensors (e.g. for tests with odd sizes) instead of the padded power-of-two shapes the kernel expects; calling the internal kernel directly instead of a wrapper that pads.

Related errors

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