{"record":{"id":"369a53190f9b72d9","repo":"jax-ml/jax","slug":"unsupported-block-dimension-type-type-bd","errorCode":null,"errorMessage":"Unsupported block dimension type: {type(bd)}","messagePattern":"Unsupported block dimension type: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/pipeline.py","lineNumber":838,"sourceCode":"    assert len(self.block_shape) == len(indices)\n    indexer = []\n    for bd, idx in zip(self.block_shape, indices, strict=True):\n      match bd:\n        case None | Squeezed():\n          # Dimension is squeezed out so we don't do anything.\n          indexer.append(idx)\n        case Element(block_size, padding=padding):\n          if padding != (0, 0):\n            raise ValueError(f\"Element with {padding=} is not supported.\")\n          indexer.append(ds(idx, block_size))\n        case BoundedSlice(block_size):\n          indexer.append(ds(idx.start, block_size))\n        case Blocked(block_size):\n          indexer.append(ds(idx * block_size, block_size))\n        case int():\n          indexer.append(ds(idx * bd, bd))\n        case _:\n          raise ValueError(f\"Unsupported block dimension type: {type(bd)}\")\n    return tuple(indexer)\n\n  def initialize_slots(self) -> BufferedRef:\n    if self.window_ref is None and self.prefetched_count > 0:\n      raise ValueError(\n          \"Expected external window buffer to be bound for prefetched input \"\n          f\"(prefetched_count={self.prefetched_count}), but window_ref is None. \"\n          \"Ensure .with_window_ref(...) is called on the BufferedRef in allocations.\"\n      )\n    return dataclasses.replace(\n        self,\n        copy_in_slot=jnp.uint32(0) if self.buffer_type.is_input else None,\n        wait_in_slot=jnp.uint32(0) if self.buffer_type.is_input else None,\n        copy_out_slot=jnp.uint32(0) if self.buffer_type.is_output else None,\n        wait_out_slot=jnp.uint32(0) if self.buffer_type.is_output else None,\n        next_fetch=(\n            tuple(jnp.int32(0) for _ in range(self._grid_rank))\n            if self._grid_rank is not None","sourceCodeStart":820,"sourceCodeEnd":856,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/pipeline.py#L820-L856","documentation":"computeSlice converts each BlockDimension of a BlockSpec into a dynamic_slice indexer via Python match. Only None, Squeezed, Element, BoundedSlice, Blocked, and int are supported; any other type falls through to this ValueError.","triggerScenarios":"Passing a custom or invalid object as a block dimension in a BlockSpec, e.g. a numpy integer (np.int64) instead of a Python int, or a future/unsupported BlockDimension subclass, into a pipelined Mosaic kernel.","commonSituations":"Using np.int64/np.intp values (from numpy arrays or computed sizes) as block_shape entries; defining a custom BlockDimension subclass expecting it to be honored.","solutions":["Ensure every block dimension is a Python int or an instance of a supported type (Blocked, BoundedSlice, Element, Squeezed, None)","Convert numpy scalars with int(...) before constructing the BlockSpec","If you subclassed a block dimension type, replace it with the closest supported type"],"exampleFix":"# before\nblock_shape=(np.int64(128),)\n# after\nblock_shape=(int(np.int64(128)),)","handlingStrategy":"type-guard","validationCode":"import numpy as np\nblock_shape = tuple(int(d) if isinstance(d, np.integer) else d for d in block_shape)","typeGuard":"SUPPORTED = (int, Element, BoundedSlice, Blocked, Squeezed, type(None))\ndef dims_supported(block_shape) -> bool:\n    return all(d is None or isinstance(d, SUPPORTED) and not isinstance(d, bool) or d is None for d in block_shape)","tryCatchPattern":null,"preventionTips":["Never put numpy scalars or custom objects in block_shape","Lint spec construction with isinstance checks in a helper"],"tags":["jax","pallas","blockspec","type-validation"],"backgroundTag":"type-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}