{"record":{"id":"232273abf582607b","repo":"jax-ml/jax","slug":"unexpected-dim-block-size-dim-block-size","errorCode":null,"errorMessage":"Unexpected dim_block_size: {dim_block_size}","messagePattern":"Unexpected dim_block_size: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":1923,"sourceCode":"  full_size = math.prod(full_shape) * array_dtype.itemsize\n  # Use 64-bit indexing when offset might be >= 2**32 bytes.\n  offset_eltype = ir.IntegerType.get_signless(64 if full_size > 2**32 else 32)\n  if indexer_shape:\n    offsets = _zeros(ir.RankedTensorType.get(indexer_shape, offset_eltype))\n  else:\n    offsets = _ir_constant(0, offset_eltype)\n\n  indexer_iter = iter(indices)\n  for dim_stride, dim_block_size, start_offset in zip(\n      strides, block_info.block_shape, block_info.start_indices\n  ):\n    match dim_block_size:\n      case pallas_core.Squeezed():\n        index = _ir_constant(0, offset_eltype)\n      case int():\n        index = next(indexer_iter)\n      case _:\n        raise ValueError(f\"Unexpected dim_block_size: {dim_block_size}\")\n\n    if isinstance(index, slice):\n      index = primitives.Slice.from_slice(\n          index, pallas_core.get_block_size(dim_block_size)\n      )\n\n    if isinstance(index, primitives.Slice):\n      if index.is_dynamic_start or (index.stride != 1):\n        if not index.is_dynamic_start:\n          start = _ir_constant(index.start, offset_eltype)\n        else:\n          assert isinstance(index.start, ir.Value)\n          start = index.start\n        start = _ir_cast(start, offset_eltype, signed=False)\n\n        iota = _ir_cast(\n            _make_range(0, int(index.size)), offset_eltype, signed=False\n        )","sourceCodeStart":1905,"sourceCodeEnd":1941,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L1905-L1941","documentation":"When computing pointer offsets from Pallas block indices, each dimension's block specification must be either a Squeezed() dimension (lowered to constant 0) or a plain int block size (consumed from the index iterator). Any other object in the match triggers this internal ValueError, indicating an internal inconsistency or a new/unsupported BlockMapping form.","triggerScenarios":"Passing a BlockMapping whose dimension block size is neither pallas_core.Squeezed nor an int (e.g. None, a custom sentinel, or an object from a version-mismatched pallas_core) when launching a Triton Pallas kernel; reached from load/store pointer computation.","commonSituations":"Mixing JAX versions (jax + jaxlib or a vendored pallas_core with mismatched types); using experimental BlockMapping features not yet supported by the Triton backend; internal API changes across JAX releases.","solutions":["Ensure jax/jaxlib versions match exactly (pip install -U jax jaxlib) and that only one pallas install is present","Avoid custom or experimental BlockMapping dimension specs; use standard int block sizes or squeezed dims","Reproduce with a minimal kernel and report the internal error to the JAX repository","Pin to a JAX version where your GridSpec/BlockMapping usage was known to work"],"exampleFix":null,"handlingStrategy":"retry","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    compiled = kernel.compile(...)\nexcept ValueError as e:\n    if 'Unexpected dim_block_size' in str(e):\n        raise RuntimeError('JAX internal version mismatch; align jax/jaxlib/pallas versions') from e\n    raise","preventionTips":["Install matching jax and jaxlib versions in one command","Avoid mixing pip JAX with source-built pallas/triton plugins","Pin versions in requirements to avoid surprise internal API drift"],"tags":["jax","triton","pallas","internal","block-mapping","version-mismatch"],"backgroundTag":"library-version-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}