jax-ml/jax · error · NotImplementedError

Padding with aliasing not supported.

Error message

Padding with aliasing not supported.

What it means

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.

Source

Thrown at jax/_src/pallas/hlo_interpreter.py:362

    print(discharged_jaxpr)
  out = _initialize_output_vals(grid_mapping.block_mappings_output,
                                args, input_output_aliases)
  # TODO(b/370563936): Fix correctness issue w/ io aliasing
  scalars = args[grid_mapping.slice_index_ops]
  block_args = args[len(scalars):]
  # invars: [*scalar_prefetch, *consts, *inputs, *outputs, *scratch]
  # block_args now contains: *consts, *inputs, *outputs
  scratch_values = tuple(
      primitives.uninitialized_value(a.shape, a.dtype) for a in scratch_avals
  )

  carry = []
  for x, bm in zip(itertools.chain(block_args, out), grid_mapping.block_mappings):
    padding = [bd.padding if isinstance(bd, pallas_core.Element) else (0, 0)
               for bd in bm.block_shape]
    if padding is not None and any(p != (0, 0) for p in padding):
      if input_output_aliases:
        raise NotImplementedError("Padding with aliasing not supported.")
      pad_value = primitives.uninitialized_value(shape=(), dtype=x.dtype)
      x = lax.pad(x, pad_value, [(*p, 0) for p in padding])
    carry.append(x)

  block_shapes = [pallas_core._get_block_shape(bm.block_shape)
                  for bm in grid_mapping.block_mappings]
  is_squeeze_dim = [
      tuple(isinstance(bd, pallas_core.Squeezed) for bd in bm.block_shape)
      for bm in grid_mapping.block_mappings
  ]

  # Pad values to evenly divide into block dimensions. This matches the
  # behavior of the non-interpret mode. We pad with NaN, to make it easier
  # to catch OOB accesses.

  carry = map(_pad_to_block_dimension, carry, block_shapes)
  carry.extend(scratch_values)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove input_output_aliases: pass the aliased buffers as separate inputs and copy results explicitly in the kernel
2. Remove padding from the in_specs (use (0,0) padding and handle boundary checks manually with masked indices)
3. Verify with the real compiler whether the aliasing+padding combination is even legal there; if not, redesign the kernel
4. Update JAX — later versions may lift this interpreter restriction

Example fix

// before
out = pallas_call(kernel, out_shape, in_specs=..., input_output_aliases=(0,))(x, out)
// after
out = pallas_call(kernel, out_shape, in_specs=...)(x)
out = kernel_result  # write via out_specs without aliasing

def specs_ok_for_interpret(in_specs, aliases):
    def padded(b):
        return any(getattr(d, 'padding', (0, 0)) != (0, 0) for d in b.block_shape)
    return not aliases or not any(padded(s) for s in in_specs)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Parameter {key} is not a Jaxpr or sequence of Jaxprs: {value
• run_scoped interpret rule does not support collective axes
• Advanced indexers are not supported on TPU
• Revisited block {output_ranges[i]} of output {i} in iteratio
• Aliasing of scalar prefetch arguments is not currently suppo

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/652a224da264d12a. Report an issue: GitHub.