{"record":{"id":"3221334cd7cc37bc","repo":"jax-ml/jax","slug":"interpret-with-dynamic-grid-bounds-unsupported","errorCode":null,"errorMessage":"interpret with dynamic grid bounds unsupported","messagePattern":"interpret with dynamic grid bounds unsupported","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/pallas_call.py","lineNumber":262,"sourceCode":"def _pallas_call_jvp_rule(\n    primals,\n    tangents,\n    *,\n    jaxpr: jax_core.Jaxpr,\n    input_output_aliases: tuple[tuple[int, int], ...],\n    grid_mapping: GridMapping,\n    mesh: pallas_core.Mesh | None,\n    debug: bool,\n    interpret: Any,\n    compiler_params: CompilerParams | None,\n    cost_estimate: CostEstimate | None,\n    out_avals: tuple[jax_core.AbstractValue, ...],\n    metadata: FrozenDict[str, str] | None,\n    name: str | None,\n):\n  debug_info = jaxpr.debug_info\n  if grid_mapping.num_dynamic_grid_bounds:\n    raise NotImplementedError(\"interpret with dynamic grid bounds unsupported\")\n  if grid_mapping.num_index_operands:\n    raise NotImplementedError\n  if input_output_aliases:\n    raise NotImplementedError(\"JVP with aliasing not supported.\")\n  if mesh is not None:\n    raise NotImplementedError(\"pallas_call with a mesh does not support JVP\")\n  nonzero_tangents = [not isinstance(t, ad_util.Zero) for t in tangents]\n  tangents = [t for t in tangents if type(t) is not ad_util.Zero]\n  nonzero_tangents_with_outputs = nonzero_tangents + [True] * grid_mapping.num_outputs\n  closed_jaxpr = jaxpr\n  jvp_jaxpr_, _ = ad.jvp_jaxpr(closed_jaxpr, nonzero_tangents_with_outputs, [])\n  jvp_jaxpr, () = jvp_jaxpr_, jvp_jaxpr_.consts  # TODO consts\n  # `pallas_call` takes in inputs and returns outputs but its jaxpr *does not*.\n  # `pallas_call` takes in a stateful jaxpr, meaning the jaxpr accepts input\n  # `Ref`s that are read from followed by output `Ref`s that are written to.\n  # This means that when we do `jvp_jaxpr` on the `jaxpr`, we get out a new\n  # jaxpr that has tangents following primals. In order for this jaxpr to be\n  # compatible w/ `pallas_call` (inputs then outputs), we need to shuffle around","sourceCodeStart":244,"sourceCodeEnd":280,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/pallas_call.py#L244-L280","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Pass grid bounds as static Python integers/constants instead of traced values","Hoist the dynamic dimension out of the differentiated region: compute grid outside jvp and pass it statically","Use a different path for differentiation (e.g. custom JVP rule for the kernel, or reverse-mode if supported for your config)"],"exampleFix":"# before\npallas_call(kernel, grid=(n,), out_shape=...)(x)  # n is traced; jvp(...)\n# after\npallas_call(kernel, grid=(int(n),), out_shape=...)(x)  # static grid","handlingStrategy":"validation","validationCode":"grid_ints = all(isinstance(g, (int,)) and not hasattr(g, 'aval') for g in grid)\nassert grid_ints, 'grid bounds must be static ints before jvp'","typeGuard":"def is_static_grid(grid) -> bool:\n    return all(isinstance(g, int) for g in grid)","tryCatchPattern":null,"preventionTips":["Always pass static integer grids to pallas_call that will be differentiated","Compute dynamic sizes outside the differentiated function and close over them as constants"],"tags":["jax","pallas","jvp","autodiff","dynamic-grid","notimplementederror"],"backgroundTag":"autodiff-unsupported-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}