{"record":{"id":"34a81e3399ab398a","repo":"jax-ml/jax","slug":"unsupported-block-dimension-type-type-dim-all","errorCode":null,"errorMessage":"Unsupported block dimension type: {type(dim)}. Allowed types: `pl.Squeezed`, `pl.Blocked`, `pl.Element`, `int`, `None`.","messagePattern":"Unsupported block dimension type: (.+?)\\. Allowed types: `pl\\.Squeezed`, `pl\\.Blocked`, `pl\\.Element`, `int`, `None`\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":464,"sourceCode":"\n\ndef _canonicalize_block_dim(dim: BlockDim | int | None) -> BlockDim:\n  match dim:\n    case None:\n      return squeezed\n    case int():\n      return Blocked(int(dim))\n    case Squeezed() | Blocked() | Element() | BoundedSlice() | Indirect():\n      return dim\n    case _:\n      # Handle case where the dim is a symbolic dimension so we assume it is\n      # Blocked.\n      if jax_core.is_symbolic_dim(dim):\n        return Blocked(dim)\n      try:\n        return Blocked(int(dim))\n      except Exception as e:\n        raise ValueError(\n            f\"Unsupported block dimension type: {type(dim)}. Allowed types:\"\n            \" `pl.Squeezed`, `pl.Blocked`, `pl.Element`, `int`, `None`.\"\n        ) from e\n\ndef _canonicalize_block_shape(block_shape: Sequence[BlockDim | int | None]\n                              ) -> tuple[BlockDim, ...]:\n  return tuple(_canonicalize_block_dim(dim) for dim in block_shape)\n\n\ndef _get_block_dim_size(dim: BlockDim) -> int:\n  match dim:\n    case Squeezed():\n      return 1\n    case (\n        Blocked(block_size)\n        | Element(block_size)\n        | BoundedSlice(block_size)\n        | Indirect(block_size)","sourceCodeStart":446,"sourceCodeEnd":482,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L446-L482","documentation":"Block shapes in BlockSpec must consist of ints, None, or the symbolic dims pl.Blocked/pl.Squeezed/pl.Element. _canonicalize_block_dim tries each form and, when the value can't be interpreted (failed int() conversion), raises this ValueError listing the allowed types.","triggerScenarios":"Passing a BlockSpec block_shape containing a float, numpy scalar that fails int conversion, string, or arbitrary object — e.g. block_shape=(102.0, None) instead of (1024, None), or a JAX tracer inside block_shape.","commonSituations":"Computing block sizes with division that yields floats (n / 8 instead of n // 8); passing weakly-typed numpy floats; typos in BlockSpec tuples.","solutions":["Use integer division or int() when computing block dims: n // 8 or cdiv helper","Keep block_shape entries as Python ints, None, or pl.Blocked/Squeezed/Element only","Print block_shape right before BlockSpec construction to spot float/str entries"],"exampleFix":"# before\nblock = (x.shape[0] / tile, None)  # floats\nspec = pl.BlockSpec(block_shape=block, index_map=...)\n# after\nblock = (x.shape[0] // tile, None)\nspec = pl.BlockSpec(block_shape=block, index_map=...)","handlingStrategy":"validation","validationCode":"dims = tuple(int(d) if isinstance(d, (int, float)) and not isinstance(d, bool) else d for d in block_shape)\nblock_shape = tuple(d if isinstance(d, (int, type(None))) or hasattr(d, 'block_size') or hasattr(d, 'size') else int(d) for d in dims)","typeGuard":"def valid_block_dim(d):\n    import jax.experimental.pallas as pl\n    return d is None or isinstance(d, int) or isinstance(d, (pl.Blocked, pl.Squeezed, pl.Element))","tryCatchPattern":null,"preventionTips":["Use // not / when computing tile sizes","Run a small assertion pass over block_shape before BlockSpec"],"tags":["jax","pallas","block-spec","type-validation"],"backgroundTag":"invalid-type-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}