jax-ml/jax · error · NotImplementedError

interpret with dynamic grid bounds unsupported

Error message

interpret with dynamic grid bounds unsupported

What it means

The JVP (forward-mode autodiff) rule for pallas_call does not support dynamic grid bounds. If the GridMapping was created with num_dynamic_grid_bounds > 0 (grid sizes computed from runtime values), differentiation is unsupported and NotImplementedError is raised.

Source

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

def _pallas_call_jvp_rule(
    primals,
    tangents,
    *,
    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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass grid bounds as static Python integers/constants instead of traced values
2. Hoist the dynamic dimension out of the differentiated region: compute grid outside jvp and pass it statically
3. Use a different path for differentiation (e.g. custom JVP rule for the kernel, or reverse-mode if supported for your config)

Example fix

# before
pallas_call(kernel, grid=(n,), out_shape=...)(x)  # n is traced; jvp(...)
# after
pallas_call(kernel, grid=(int(n),), out_shape=...)(x)  # static grid

grid_ints = all(isinstance(g, (int,)) and not hasattr(g, 'aval') for g in grid)
assert grid_ints, 'grid bounds must be static ints before jvp'

def is_static_grid(grid) -> bool:
    return all(isinstance(g, int) for g in grid)

When it happens

Trigger: Calling jax.jvp (or jax.grad with forward-mode internals / jax.checkpointed JVP) on a function containing pallas_call whose grid argument depends on traced runtime values, e.g. pallas_call(kernel, grid=(n,), ...) with n a JAX value.

Common situations: Making grid sizes data-dependent (dynamic shapes) and then attempting differentiation through the Pallas kernel; switching a model's loss path containing a Pallas kernel to forward-mode AD or jvp.

Related errors

• JVP with aliasing not supported.
• 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/3221334cd7cc37bc. Report an issue: GitHub.