{"record":{"id":"cfd2e1be0c873de0","repo":"jax-ml/jax","slug":"shape-dimension-shape-2-must-be-divisible-by-cfd2e1","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/threefry.py","lineNumber":57,"sourceCode":"  Args:\n    key: A threefry key of shape (2,).\n    shape: The shape of the output. Must be divisible by `block_size`.\n    unpadded_shape: If `shape` is padded, then this is the shape of the\n      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\n  Returns:\n    A tensor of random bits of shape `shape`.\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\n  def kernel(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)\n    counts_lo = counts_lo.astype(jnp.uint32)\n    # TODO(justinfu): Support hi bits on count.\n    counts_hi = 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 = threefry2x32.threefry2x32_p.bind(\n        k1, k2, counts_hi, counts_lo)\n    out_bits = o1 ^ o2","sourceCodeStart":39,"sourceCodeEnd":75,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/random/threefry.py#L39-L75","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\nbits = threefry_2x32_count(key, (10, 10))\n// after\nbits = threefry_2x32_count(key, (128, 128))[:10, :10]","handlingStrategy":"validation","validationCode":"assert shape[-2] % block_size[-2] == 0 and shape[-1] % block_size[-1] == 0, 'pad to block multiples and slice after'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","pallas","tpu","random","threefry","alignment"],"backgroundTag":"shape-block-alignment","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}