jax-ml/jax · error · ValueError

Offset must be scalar, got {offset.shape}

Error message

Offset must be scalar, got {offset.shape}

What it means

philox_4x32_count accepts an offset (used to advance the Philox counter stream) that must be a scalar; after conversion to jnp.uint32, offset.ndim must be 0. Passing an array offset makes the counter arithmetic ill-defined and is rejected.

Source

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

  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))
    k2 = jnp.reshape(key_ref[0, 1], (1, 1))
    o1, o2, _, _ = philox_4x32(_zeros, counts_lo, _zeros, _zeros, k1, k2)
    if fuse_output:
      out_bits = o1 ^ o2
      out_ref[...] = out_bits.reshape(out_ref.shape)
    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a Python int or 0-d array: offset=int(x) or jnp.squeeze(x)
2. If you need per-element offsets, call the kernel per offset or restructure using fold_in
3. Assert data.ndim == 0 before fold_in style usage

Example fix

// before
philox_4x32_count(key, shape, offset=jnp.array([start]))
// after
philox_4x32_count(key, shape, offset=int(start))

offset = int(offset) if np.ndim(offset) == 0 else None
assert offset is not None or np.ndim(offset := jnp.asarray(offset).squeeze()) == 0

def is_scalar(x) -> bool:
    return jnp.asarray(x).ndim == 0

When it happens

Trigger: Passing offset as a 1-element array (e.g. jnp.array([0])), or broadcasting a per-element offset vector, e.g. philox_fold_in style calls where data kept an axis.

Common situations: Computing offsets with arithmetic that accidentally keeps a shape, e.g. offset = idx * stride where idx is an array; adapting fold_in code that asserts scalar but receives size-1 tensors.

Related errors

• Shape too large: {np.prod(shape)} > {np.iinfo(jnp.uint32).ma
• Shape dimension {shape[-2:]} must be divisible by {block_siz
• Only 32-bit PRNG supported.
• PRNG keys must be loaded from SMEM. Did you set the memory s
• Bit width must be 32

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