{"record":{"id":"aa660c13b361a9f6","repo":"jax-ml/jax","slug":"element-with-padding-is-not-supported","errorCode":null,"errorMessage":"Element with {padding=} is not supported.","messagePattern":"Element with (.+?) is not supported\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/pipeline.py","lineNumber":829,"sourceCode":"  def unbind_refs(self):\n    if not self.is_buffered and not self.has_allocated_buffer:\n      return dataclasses.replace(self, window_ref=None)\n    return self\n\n  def compute_slice(self, grid_indices):\n    \"\"\"Compute the indexers for the window at given grid indices.\"\"\"\n    indices = self.compute_index(*grid_indices)\n    assert self.block_shape is not None\n    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      )","sourceCodeStart":811,"sourceCodeEnd":847,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/pipeline.py#L811-L847","documentation":"In JAX's Mosaic GPU pipeline (jax.experimental.pallas / tpu), each dimension of a BufferedRef's windowing spec is described by a BlockDimension. When a dimension uses Element(block_size, padding=...) with non-zero padding, compute_slice cannot translate it into a dynamic_slice and raises this ValueError. Zero padding is fine; any other padding tuple is rejected.","triggerScenarios":"Passing a BlockSpec whose block dimensions include an Element with padding != (0, 0) to a kernel using pipelined/buffered execution (e.g. emit_pipeline with buffering, fetch_with_lookahead, or make_output_bref paths that build indexers via compute_slice).","commonSituations":"Converting a Pallas kernel from CPU/TPU interp mode to GPU Mosaic pipeline mode where padded BlockSpecs were previously tolerated; using halos/borders expressed as Element padding instead of explicit slice bounds.","solutions":["Change padding=(left, right) on Element block dimensions to (0, 0) and handle out-of-bounds edges by adjusting the index map / block shape","Use a BoundedSlice or explicit start indices to express the padded window instead of Element padding","Run in interpreter mode (jax.platform or interpret=True) if padding semantics are required and pipeline support isn't needed"],"exampleFix":"# before\nBlockSpec(index_map=lambda i: (i * 128,), block_shape=(128,), padding=((2, 2),))\n# after\nBlockSpec(index_map=lambda i: (i * 128 - 2,), block_shape=(132,), padding=((0, 0),))","handlingStrategy":"validation","validationCode":"bad = [d for d in spec.block_shape if isinstance(d, Element) and d.padding != (0, 0)]\nassert not bad, f'Unsupported Element padding: {bad}'","typeGuard":"from jax._src.pallas.mosaic.pipeline import Element\ndef has_zero_padding(spec) -> bool:\n    return all(not isinstance(d, Element) or d.padding == (0, 0)\n               for d in spec.block_shape)","tryCatchPattern":null,"preventionTips":["Avoid Element padding in BlockSpecs destined for Mosaic pipelines; express halos via index maps","Test kernels in interpret mode first, then under the real pipeline backend","Keep block_shape entries as plain ints or Blocked/BoundedSlice"],"tags":["jax","pallas","mosaic","pipeline","blockspec","padding"],"backgroundTag":"unsupported-argument-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}