{"record":{"id":"893b10de5624d9e1","repo":"jax-ml/jax","slug":"pallas-call-with-a-mesh-does-not-support-jvp","errorCode":null,"errorMessage":"pallas_call with a mesh does not support JVP","messagePattern":"pallas_call with a mesh does not support JVP","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/pallas_call.py","lineNumber":268,"sourceCode":"    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\n  # the jaxpr's invars.\n  primal_refs, primal_out_refs, tangent_refs, tangent_out_refs = split_list(\n      jvp_jaxpr.invars, [len(primals), grid_mapping.num_outputs, len(tangents)]\n  )\n  invars = (*primal_refs, *tangent_refs, *primal_out_refs, *tangent_out_refs)\n  jvp_jaxpr = jvp_jaxpr.replace(invars=invars)","sourceCodeStart":250,"sourceCodeEnd":286,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/pallas_call.py#L250-L286","documentation":"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.","triggerScenarios":"Passing mesh=... to pallas_call (or pallas_call setup APIs that attach a mesh) and subsequently calling jax.jvp on the resulting function.","commonSituations":"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.","solutions":["Drop the mesh parameter and run the kernel per-device (manual replication) if forward-mode differentiation is required","Use jax.custom_jvp to define the forward-mode rule manually for the mesh-based kernel","Differentiate outside the mesh context with a non-Pallas implementation and only run the meshed pallas_call in the forward pass with stop_gradient"],"exampleFix":"# before\np = pallas_call(kernel, grid=grid, out_shape=out_shape, mesh=my_mesh)\njvp(p, (x, t))\n# after\np = pallas_call(kernel, grid=grid, out_shape=out_shape)  # no mesh\njvp(p, (x, t))","handlingStrategy":"fallback","validationCode":"assert mesh is None or not use_jvp, 'meshed pallas_call does not support jvp'","typeGuard":null,"tryCatchPattern":"try:\n    jax.jvp(f, (x,), (t,))\nexcept NotImplementedError as e:\n    if 'mesh' in str(e):\n        return non_meshed_jvp_path(x, t)  # per-device or custom_jvp fallback\n    raise","preventionTips":["Gate mesh usage behind a flag so AD paths can run without mesh","Predefine custom_jvp rules for SPMD Pallas kernels used in training"],"tags":["jax","pallas","jvp","autodiff","mesh","spmd","notimplementederror"],"backgroundTag":"autodiff-unsupported-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}