{"record":{"id":"a851068db0cf98cb","repo":"jax-ml/jax","slug":"unsupported-block-dimension-type-type-bd-for-b","errorCode":null,"errorMessage":"Unsupported block dimension type: {type(bd)} for block shape: {bm.block_shape}","messagePattern":"Unsupported block dimension type: (.+?) for block shape: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":1241,"sourceCode":"          tpu_memory_space == tpu_core.MemorySpace.VMEM\n          and bm.has_trivial_window()\n      ):\n        pipeline_mode = pallas_core.Buffered(1)\n\n      # If we have an extended dtype, we need to add the block shape for the\n      # remaining physical dtype.\n      block_shape += list(_get_aval_physical_dtype_shape(bm.block_aval.inner_aval))\n      block_shape = dynamic_shape_replacement_fn(block_shape)\n      window_shape = ir.DenseI64ArrayAttr.get(block_shape)\n      block_params: dict[str, ir.Attribute] = dict(\n          window_bounds=window_shape,\n          transform_indices=ir.FlatSymbolRefAttr.get(func_name),\n      )\n      for bd in bm.block_shape:\n        if not isinstance(\n            bd, (pallas_core.Element, pallas_core.Squeezed, pallas_core.Blocked)\n        ):\n          raise NotImplementedError(\n              \"Unsupported block dimension type: \"\n              f\"{type(bd)} for block shape: {bm.block_shape}\"\n          )\n      is_element_block = [isinstance(bd, pallas_core.Element)\n                          for bd in bm.block_shape]\n      if any(is_element_block):\n        is_element_or_squeezed_block = [\n            isinstance(bd, (pallas_core.Element, pallas_core.Squeezed))\n            for bd in bm.block_shape\n        ]\n        if not all(is_element_or_squeezed_block):\n          raise NotImplementedError(\n              \"All block dimensions must be Elements or none of them can be\"\n              \" Elements.\"\n          )\n        padding = [\n            bd.padding if isinstance(bd, pallas_core.Element) else (0, 0)\n            for bd in bm.block_shape","sourceCodeStart":1223,"sourceCodeEnd":1259,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L1223-L1259","documentation":"Block shapes in Pallas grid mappings must be built from pallas_core.Element, pallas_core.Squeezed, or pallas_core.Blocked dimension descriptors. If any dimension of bm.block_shape is another type (e.g. a plain int from manual BlockSpec construction), the TPU lowering raises NotImplementedError listing the offending type.","triggerScenarios":"Constructing BlockSpec/BlockMapping objects manually (or via internal APIs) with block_shape given as raw ints/strings instead of pallas_core.Blocked(...) / Element(...); intercepting or transforming block mappings with custom code that replaces descriptors with ints.","commonSituations":"Writing custom interpreters/partials over Pallas internals; monkey-patching or serializing/deserializing block mappings; upgrading JAX versions where block_shape representation changed from ints to descriptor objects, breaking custom glue code.","solutions":["Build block shapes with the public API (BlockSpec with plain tuples) so JAX constructs Blocked/Element descriptors for you","If manipulating internals, convert each dim via pallas_core.Blocked(...) or keep the original descriptor objects","Update custom code that assumed integer block dims after a JAX upgrade; check the BlockSpec constructor signature in your JAX version","Avoid hand-assembling GridMapping; use pallas.create_grid / pallas_call's automatic mapping instead"],"exampleFix":"# before\nGridMapping(block_mappings=[BlockMapping(block_shape=[16, 16], ...)])\n\n# after\n# let pallas_call build mappings from BlockSpec:\npallas_call(kernel, BlockSpec(block_shape=(16, 16), index_map=...), grid=(...))","handlingStrategy":"type-guard","validationCode":"from jax._src.pallas import pallas_core\ndef check_block_shape_types(bm):\n    for bd in bm.block_shape:\n        if not isinstance(bd, (pallas_core.Element, pallas_core.Squeezed, pallas_core.Blocked)):\n            raise TypeError(f'bad block dim {bd!r} in {bm.block_shape}')","typeGuard":"def is_valid_block_shape(bm) -> bool:\n    return all(isinstance(bd, (pallas_core.Element, pallas_core.Squeezed, pallas_core.Blocked))\n               for bd in bm.block_shape)","tryCatchPattern":"try:\n    lower_jaxpr_into_pipelined_module(...)\nexcept NotImplementedError as e:\n    if 'Unsupported block dimension type' in str(e):\n        rebuild mappings via public BlockSpec and retry","preventionTips":["Never hand-construct GridMapping/BlockMapping; use BlockSpec + pallas_call","After JAX upgrades, re-verify internal APIs your code touches","Keep block shapes as plain tuples at the user level"],"tags":["jax","pallas","tpu","block-shape","api-misuse"],"backgroundTag":"pallas-invalid-block-shape-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}