{"record":{"id":"cb1b68ef02baacf8","repo":"jax-ml/jax","slug":"block-size-len-block-size-and-tile-size-len","errorCode":null,"errorMessage":"block_size ({len(block_size)}) and tile_size ({len(tile_size)}) must have the same length.","messagePattern":"block_size \\((.+?)\\) and tile_size \\((.+?)\\) must have the same length\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/random.py","lineNumber":203,"sourceCode":"  Args:\n    sampler_fn: A sampling function that consumes a key and returns\n      random samples.\n    global_key: The global key to use for sampling.\n    block_size: The shape of an individual block.\n    tile_size: The shape of a ``tile``, which is the smallest unit at\n      which samples are generated. This should be selected to be a divisor\n      of all block sizes one needs to be invariant to.\n    total_size: The total size of the array to sample.\n    block_index: The index denoting which block to generate keys for. Defaults\n      to the program_id for each block axis.\n    **kwargs: Additional arguments to pass to the sampler_fn.\n\n  Returns:\n    A ``block_size`` shaped array of samples for the current block corresponding\n    to ``block_index``.\n  \"\"\"\n  if len(block_size) != len(tile_size):\n    raise ValueError(f\"block_size ({len(block_size)}) and tile_size \"\n                     f\"({len(tile_size)}) must have the same length.\")\n\n  if block_index is None:\n    num_axes = len(block_size)\n    block_index = tuple(\n      primitives.program_id(axis) for axis in range(num_axes))\n\n  keys = blocked_sampler.blocked_fold_in(\n      global_key, total_size, block_size, tile_size, block_index)\n  return blocked_sampler.sample_block(\n      sampler_fn, keys, block_size, tile_size, **kwargs)\n","sourceCodeStart":185,"sourceCodeEnd":215,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/random.py#L185-L215","documentation":"sample_block generates random samples for one tile of a larger block, so block_size and tile_size must have the same rank (number of axes). A mismatch raises ValueError before any sampling happens.","triggerScenarios":"Calling random.sample_block(key, block_size=(128,), tile_size=(128, 128)) or any combination where len(block_size) != len(tile_size), e.g. 1D block with 2D tiles.","commonSituations":"Using 2D tiles with a flattened 1D block shape; editing one of the two size tuples during refactoring and forgetting the other.","solutions":["Give block_size and tile_size the same number of axes (append 1s or reshape accordingly)","For 2D sampling pass block_size=(rows, cols) matching tile rank"],"exampleFix":"# before\nx = random.sample_block(key, block_size=(8192,), tile_size=(128, 128))\n# after\nx = random.sample_block(key, block_size=(64, 128), tile_size=(128, 128))","handlingStrategy":"validation","validationCode":"assert len(block_size) == len(tile_size), 'block/tile rank mismatch'","typeGuard":"def ranks_match(block_size, tile_size) -> bool:\n    return len(tuple(block_size)) == len(tuple(tile_size))","tryCatchPattern":null,"preventionTips":["Keep block_size and tile_size in one nested tuple structure so ranks stay in sync"],"tags":["jax","pallas","tpu","random","shape-mismatch"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}