{"record":{"id":"652a224da264d12a","repo":"jax-ml/jax","slug":"padding-with-aliasing-not-supported","errorCode":null,"errorMessage":"Padding with aliasing not supported.","messagePattern":"Padding with aliasing not supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/hlo_interpreter.py","lineNumber":362,"sourceCode":"    print(discharged_jaxpr)\n  out = _initialize_output_vals(grid_mapping.block_mappings_output,\n                                args, input_output_aliases)\n  # TODO(b/370563936): Fix correctness issue w/ io aliasing\n  scalars = args[grid_mapping.slice_index_ops]\n  block_args = args[len(scalars):]\n  # invars: [*scalar_prefetch, *consts, *inputs, *outputs, *scratch]\n  # block_args now contains: *consts, *inputs, *outputs\n  scratch_values = tuple(\n      primitives.uninitialized_value(a.shape, a.dtype) for a in scratch_avals\n  )\n\n  carry = []\n  for x, bm in zip(itertools.chain(block_args, out), grid_mapping.block_mappings):\n    padding = [bd.padding if isinstance(bd, pallas_core.Element) else (0, 0)\n               for bd in bm.block_shape]\n    if padding is not None and any(p != (0, 0) for p in padding):\n      if input_output_aliases:\n        raise NotImplementedError(\"Padding with aliasing not supported.\")\n      pad_value = primitives.uninitialized_value(shape=(), dtype=x.dtype)\n      x = lax.pad(x, pad_value, [(*p, 0) for p in padding])\n    carry.append(x)\n\n  block_shapes = [pallas_core._get_block_shape(bm.block_shape)\n                  for bm in grid_mapping.block_mappings]\n  is_squeeze_dim = [\n      tuple(isinstance(bd, pallas_core.Squeezed) for bd in bm.block_shape)\n      for bm in grid_mapping.block_mappings\n  ]\n\n  # Pad values to evenly divide into block dimensions. This matches the\n  # behavior of the non-interpret mode. We pad with NaN, to make it easier\n  # to catch OOB accesses.\n\n  carry = map(_pad_to_block_dimension, carry, block_shapes)\n  carry.extend(scratch_values)\n","sourceCodeStart":344,"sourceCodeEnd":380,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/hlo_interpreter.py#L344-L380","documentation":"During Pallas kernel interpretation, when a block mapping declares non-zero padding, the interpreter materializes padded blocks with lax.pad. If the pallas_call also uses input_output_aliases (inputs aliased to outputs), padding and aliasing cannot be combined, so a NotImplementedError is raised on the input-padding pass.","triggerScenarios":"Calling pallas_call_hlo_interpret (interpret mode) with in_specs that specify padding (block shapes with (lo, hi) padding other than (0,0)) while also passing input_output_aliases to pallas_call.","commonSituations":"Porting GPU-style Pallas kernels that use aliases for in-place updates to a padded/blocked grid; debug-running with the interpreter a kernel that compiles fine but whose spec mixes aliases with padded BlockShapes.","solutions":["Remove input_output_aliases: pass the aliased buffers as separate inputs and copy results explicitly in the kernel","Remove padding from the in_specs (use (0,0) padding and handle boundary checks manually with masked indices)","Verify with the real compiler whether the aliasing+padding combination is even legal there; if not, redesign the kernel","Update JAX — later versions may lift this interpreter restriction"],"exampleFix":"// before\nout = pallas_call(kernel, out_shape, in_specs=..., input_output_aliases=(0,))(x, out)\n// after\nout = pallas_call(kernel, out_shape, in_specs=...)(x)\nout = kernel_result  # write via out_specs without aliasing","handlingStrategy":"validation","validationCode":"def specs_ok_for_interpret(in_specs, aliases):\n    def padded(b):\n        return any(getattr(d, 'padding', (0, 0)) != (0, 0) for d in b.block_shape)\n    return not aliases or not any(padded(s) for s in in_specs)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid mixing input_output_aliases with padded BlockSpecs","Run interpreter smoke tests on every new BlockSpec configuration in CI"],"tags":["jax","pallas","aliasing","padding","interpreter"],"backgroundTag":"unsupported-operation-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}