jax-ml/jax · error · NotImplementedError

JVP with aliasing not supported.

Error message

JVP with aliasing not supported.

What it means

The JVP rule for pallas_call rejects input_output_aliases: if any input is aliased with an output (in-place updated buffer), forward-mode differentiation through the aliased write is not supported and NotImplementedError('JVP with aliasing not supported.') is raised.

Source

Thrown at jax/_src/pallas/pallas_call.py:266

    jaxpr: jax_core.Jaxpr,
    input_output_aliases: tuple[tuple[int, int], ...],
    grid_mapping: GridMapping,
    mesh: pallas_core.Mesh | None,
    debug: bool,
    interpret: Any,
    compiler_params: CompilerParams | None,
    cost_estimate: CostEstimate | None,
    out_avals: tuple[jax_core.AbstractValue, ...],
    metadata: FrozenDict[str, str] | None,
    name: str | None,
):
  debug_info = jaxpr.debug_info
  if grid_mapping.num_dynamic_grid_bounds:
    raise NotImplementedError("interpret with dynamic grid bounds unsupported")
  if grid_mapping.num_index_operands:
    raise NotImplementedError
  if input_output_aliases:
    raise NotImplementedError("JVP with aliasing not supported.")
  if mesh is not None:
    raise NotImplementedError("pallas_call with a mesh does not support JVP")
  nonzero_tangents = [not isinstance(t, ad_util.Zero) for t in tangents]
  tangents = [t for t in tangents if type(t) is not ad_util.Zero]
  nonzero_tangents_with_outputs = nonzero_tangents + [True] * grid_mapping.num_outputs
  closed_jaxpr = jaxpr
  jvp_jaxpr_, _ = ad.jvp_jaxpr(closed_jaxpr, nonzero_tangents_with_outputs, [])
  jvp_jaxpr, () = jvp_jaxpr_, jvp_jaxpr_.consts  # TODO consts
  # `pallas_call` takes in inputs and returns outputs but its jaxpr *does not*.
  # `pallas_call` takes in a stateful jaxpr, meaning the jaxpr accepts input
  # `Ref`s that are read from followed by output `Ref`s that are written to.
  # This means that when we do `jvp_jaxpr` on the `jaxpr`, we get out a new
  # jaxpr that has tangents following primals. In order for this jaxpr to be
  # compatible w/ `pallas_call` (inputs then outputs), we need to shuffle around
  # the jaxpr's invars.
  primal_refs, primal_out_refs, tangent_refs, tangent_out_refs = split_list(
      jvp_jaxpr.invars, [len(primals), grid_mapping.num_outputs, len(tangents)]
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove input_output_aliases so outputs are freshly allocated, making the pallas_call alias-free before differentiating
2. Apply stop_gradient / custom_jvp around the kernel call so the JVP rule is never invoked on the aliased pallas_call
3. Differentiate an alias-free reference implementation (e.g. an equivalent plain-JAX function) instead of the in-place Pallas kernel

Example fix

# before
p = pallas_call(kernel, grid=grid, out_shape=out_shape,
                input_output_aliases=[(0, 0)])
jvp(p, (x, t))
# after
p = pallas_call(kernel, grid=grid, out_shape=out_shape)  # no aliasing
jvp(p, (x, t))

assert not input_output_aliases, 'cannot use input_output_aliases with jvp'

When it happens

Trigger: Creating pallas_call with a non-empty input_output_aliases parameter and then calling jax.jvp (or a transformation that uses the JVP rule) on it.

Common situations: Performance-tuned Pallas kernels that write results in place via input_output_aliases, later run under grad/jvp/checkpointing; refactoring pure kernels to in-place form without considering autodiff constraints.

Related errors

• interpret with dynamic grid bounds unsupported
• pallas_call with a mesh does not support JVP
• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp must have the same t
• primal and tangent arguments to jax.jvp do not match; dtypes

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