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 Pallas Threefry kernel tiles the output's two trailing dimensions by a fixed block size; threefry_2x32_count validates shape[-2] % block_size[-2] == 0 and shape[-1] % block_size[-1] == 0 before launching.

Source

Thrown at jax/experimental/pallas/ops/tpu/random/threefry.py:57

  Args:
    key: A threefry key of shape (2,).
    shape: The shape of the output. Must be divisible by `block_size`.
    unpadded_shape: If `shape` is padded, then this is the shape of the
      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.

  Returns:
    A tensor of random bits of shape `shape`.
  """
  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],)

  def kernel(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)
    counts_lo = counts_lo.astype(jnp.uint32)
    # TODO(justinfu): Support hi bits on count.
    counts_hi = jnp.zeros_like(counts_lo)
    k1 = jnp.reshape(key_ref[0, 0], (1, 1))
    k2 = jnp.reshape(key_ref[0, 1], (1, 1))
    o1, o2 = threefry2x32.threefry2x32_p.bind(
        k1, k2, counts_hi, counts_lo)
    out_bits = o1 ^ o2

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pad the trailing dims to block multiples and slice the result
Use the wrapper plthreefry_random_bits / jax.random.bits, which handles padding
Align requested shapes to the kernel's block_size

Example fix

// before
bits = threefry_2x32_count(key, (10, 10))
// after
bits = threefry_2x32_count(key, (128, 128))[:10, :10]

Defensive patterns

Strategy: validation

Validate before calling

assert shape[-2] % block_size[-2] == 0 and shape[-1] % block_size[-1] == 0, 'pad to block multiples and slice after'

Prevention

Use plthreefry_random_bits (the wrapper) rather than the raw kernel; it pads for you
Pad trailing dims up to block multiples when calling kernels directly

When it happens

Trigger: Requesting random bits whose last two dims are not multiples of the kernel block size, e.g. shape (10, 10) against 128x128 blocks, especially small or odd-shaped tensors.

Common situations: Calling the Pallas threefry path directly with arbitrary user shapes; the public wrapper pads internally, so this mostly bites when invoking the kernel function directly.

Related errors

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