{"record":{"id":"f4a46f09b2052722","repo":"jax-ml/jax","slug":"offset-must-be-scalar-got-offset-shape","errorCode":null,"errorMessage":"Offset must be scalar, got {offset.shape}","messagePattern":"Offset must be scalar, got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/random/philox.py","lineNumber":113,"sourceCode":"\n  Returns:\n    A tensor of random bits of shape `shape` if fuse_output=True. Otherwise,\n    this will return a tensor of shape (2, *shape) with the first channel being\n    the high bits and the second channel being the low bits.\n  \"\"\"\n  shape = tuple(shape)\n  if np.prod(shape) > jnp.iinfo(jnp.uint32).max:\n    raise ValueError(\n        f\"Shape too large: {np.prod(shape)} > {np.iinfo(jnp.uint32).max}\")\n\n  if (shape[-2] % block_size[-2] != 0) or (shape[-1] % block_size[-1] != 0):\n    raise ValueError(\n        f\"Shape dimension {shape[-2:]} must be divisible by {block_size}\")\n  grid_dims = shape[:-2] + (\n      shape[-2] // block_size[-2], shape[-1] // block_size[1],)\n  offset = jnp.array(offset, dtype=jnp.uint32)\n  if offset.ndim != 0:\n    raise ValueError(f\"Offset must be scalar, got {offset.shape}\")\n  offset = jnp.reshape(offset, (1,))\n\n  def kernel(offset_ref, key_ref, out_ref):\n    counts_idx = tuple(pl.program_id(i) for i in range(len(grid_dims)))\n    offset = prng_utils.compute_scalar_offset(\n        counts_idx, unpadded_shape, block_shape)\n    counts_lo = prng_utils.blocked_iota(block_size, unpadded_shape)\n    counts_lo = counts_lo + offset.astype(jnp.uint32) + offset_ref[0]\n    counts_lo = counts_lo.astype(jnp.uint32)\n    # TODO(justinfu): Support hi bits on count.\n    _zeros = jnp.zeros_like(counts_lo)\n    k1 = jnp.reshape(key_ref[0, 0], (1, 1))\n    k2 = jnp.reshape(key_ref[0, 1], (1, 1))\n    o1, o2, _, _ = philox_4x32(_zeros, counts_lo, _zeros, _zeros, k1, k2)\n    if fuse_output:\n      out_bits = o1 ^ o2\n      out_ref[...] = out_bits.reshape(out_ref.shape)\n    else:","sourceCodeStart":95,"sourceCodeEnd":131,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/random/philox.py#L95-L131","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Pass a Python int or 0-d array: offset=int(x) or jnp.squeeze(x)","If you need per-element offsets, call the kernel per offset or restructure using fold_in","Assert data.ndim == 0 before fold_in style usage"],"exampleFix":"// before\nphilox_4x32_count(key, shape, offset=jnp.array([start]))\n// after\nphilox_4x32_count(key, shape, offset=int(start))","handlingStrategy":"validation","validationCode":"offset = int(offset) if np.ndim(offset) == 0 else None\nassert offset is not None or np.ndim(offset := jnp.asarray(offset).squeeze()) == 0","typeGuard":"def is_scalar(x) -> bool:\n    return jnp.asarray(x).ndim == 0","tryCatchPattern":null,"preventionTips":["Always convert offsets to Python ints before passing","Squeeze size-1 arrays coming from vectorized arithmetic"],"tags":["jax","pallas","tpu","random","philox","scalar-validation"],"backgroundTag":"non-scalar-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}