{"record":{"id":"56fc37551996ff37","repo":"jax-ml/jax","slug":"shape-dimension-shape-2-must-be-divisible-by","errorCode":null,"errorMessage":"Shape dimension {shape[-2:]} must be divisible by {block_size}","messagePattern":"Shape dimension (.+?) must be divisible by (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/random/philox.py","lineNumber":107,"sourceCode":"      output tensor if it were not padded. This is important for indexing\n      calculations within the kernel. If `shape` is not padded, then this\n      should be equal to `shape`.\n    block_size: The block size of the kernel.\n    offset: An optional offset to the counts.\n    fuse_output: Whether to fuse the output bits into a single value.\n\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))","sourceCodeStart":89,"sourceCodeEnd":125,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/random/philox.py#L89-L125","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\nbits = philox_random_bits(key, 32, (128, 100))\n// after\npadded = philox_random_bits(key, 32, (128, 128))\nbits = padded[:, :100]","handlingStrategy":"validation","validationCode":"bs = block_size  # kernel's block size\nassert shape[-2] % bs[-2] == 0 and shape[-1] % bs[-1] == 0, 'pad trailing dims to block multiples'\npadded = tuple(-(-s % b) * b for s, b in zip(shape[-2:], bs))","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Prefer public random APIs that pad internally","Keep trailing dims multiples of the kernel block size when calling Pallas PRNG kernels directly"],"tags":["jax","pallas","tpu","random","philox","alignment"],"backgroundTag":"shape-block-alignment","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}