{"record":{"id":"1c21398b6a0961e0","repo":"jax-ml/jax","slug":"all-block-dimensions-must-be-elements-or-none-of-t","errorCode":null,"errorMessage":"All block dimensions must be Elements or none of them can be Elements.","messagePattern":"All block dimensions must be Elements or none of them can be Elements\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":1253,"sourceCode":"          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\n        ]\n        pad_low, pad_high = map(list, zip(*padding))\n        block_params[\"window_kind\"] = ir.Attribute.parse(\n            f\"#tpu.element_window<{pad_low},{pad_high}>\"\n        )\n      if pipeline_mode is not None:\n        if not isinstance(pipeline_mode, pallas_core.Buffered):\n          raise LoweringException(\n              f\"Unsupported pipeline mode: {pipeline_mode}.\"\n          )\n        if pipeline_mode.use_lookahead:\n          raise NotImplementedError(","sourceCodeStart":1235,"sourceCodeEnd":1271,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L1235-L1271","documentation":"When any dimension of a block shape is a pallas_core.Element (an element-level windowed block), every other dimension must be an Element or Squeezed. Mixing Element dims with Blocked dims is unsupported because the TPU lowering can only emit an element_window attribute for uniformly element-wise blocks.","triggerScenarios":"A block_shape like [Element(...), Blocked(16)] — some dims element-level, others tiled — reaching lower_jaxpr_into_pipelined_module; typically from custom BlockSpecs that use pallas_core.Element for one dimension only (e.g. to get per-element padding on the last dim while tiling the first).","commonSituations":"Hand-built block mappings using Element descriptors for fine-grained padding; experimental kernels combining element windows with tiling; internal code that wraps some dims in Element for dynamic shapes while leaving others Blocked.","solutions":["Make all dimensions uniformly Blocked (tile the whole shape) and handle padding inside the kernel body","If element semantics are needed, use Element/Squeezed for every dimension","Redesign the kernel: split the array so the element-windowed access is a separate rank-appropriate operand"],"exampleFix":"# before\nblock_shape=[Blocked(128), Element(padding=(1, 1))]\n\n# after\nblock_shape=[Blocked(128), Blocked(1)]  # tile uniformly; pad in-kernel","handlingStrategy":"type-guard","validationCode":"from jax._src.pallas import pallas_core\ndef check_uniform_element_blocks(bm):\n    if any(isinstance(bd, pallas_core.Element) for bd in bm.block_shape):\n        if not all(isinstance(bd, (pallas_core.Element, pallas_core.Squeezed)) for bd in bm.block_shape):\n            raise ValueError('Element dims cannot be mixed with Blocked dims')","typeGuard":"def has_uniform_element_blocks(bm) -> bool:\n    flags = [isinstance(bd, (pallas_core.Element, pallas_core.Squeezed))\n             for bd in bm.block_shape]\n    return all(flags) if any(isinstance(bd, pallas_core.Element) for bd in bm.block_shape) else True","tryCatchPattern":"try:\n    pallas_call(...)\nexcept NotImplementedError as e:\n    if 'Elements or none' in str(e):\n        switch all dims to Blocked and handle padding in-kernel","preventionTips":["Avoid Element descriptors unless the entire block is element-wise","Do padding/slicing inside the kernel body rather than via element windows"],"tags":["jax","pallas","tpu","block-shape","element-window"],"backgroundTag":"pallas-mixed-element-block-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}