jax-ml/jax · error · NotImplementedError

pallas_call with a mesh does not support JVP

Error message

pallas_call with a mesh does not support JVP

What it means

The JVP rule for pallas_call raises NotImplementedError when the primitive was created with a mesh (multi-device/SPMD Pallas, e.g. via the mesh= parameter). Forward-mode differentiation of mesh-partitioned pallas_call kernels is not implemented.

Source

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

    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)]
  )
  invars = (*primal_refs, *tangent_refs, *primal_out_refs, *tangent_out_refs)
  jvp_jaxpr = jvp_jaxpr.replace(invars=invars)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop the mesh parameter and run the kernel per-device (manual replication) if forward-mode differentiation is required
2. Use jax.custom_jvp to define the forward-mode rule manually for the mesh-based kernel
3. Differentiate outside the mesh context with a non-Pallas implementation and only run the meshed pallas_call in the forward pass with stop_gradient

Example fix

# before
p = pallas_call(kernel, grid=grid, out_shape=out_shape, mesh=my_mesh)
jvp(p, (x, t))
# after
p = pallas_call(kernel, grid=grid, out_shape=out_shape)  # no mesh
jvp(p, (x, t))

assert mesh is None or not use_jvp, 'meshed pallas_call does not support jvp'

try:
    jax.jvp(f, (x,), (t,))
except NotImplementedError as e:
    if 'mesh' in str(e):
        return non_meshed_jvp_path(x, t)  # per-device or custom_jvp fallback
    raise

When it happens

Trigger: Passing mesh=... to pallas_call (or pallas_call setup APIs that attach a mesh) and subsequently calling jax.jvp on the resulting function.

Common situations: Running multi-device Pallas kernels (TPU pods, multi-GPU) inside a training loop that uses forward-mode AD or jax.grad paths triggering JVP; upgrading single-device kernels to mesh-based SPMD kernels while keeping autodiff.

Related errors

• interpret with dynamic grid bounds unsupported
• JVP with aliasing not supported.
• 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/893b10de5624d9e1. Report an issue: GitHub.