{"record":{"id":"0cf7612941a367b5","repo":"jax-ml/jax","slug":"jvp-with-aliasing-not-supported","errorCode":null,"errorMessage":"JVP with aliasing not supported.","messagePattern":"JVP with aliasing not supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/pallas_call.py","lineNumber":266,"sourceCode":"    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\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  )","sourceCodeStart":248,"sourceCodeEnd":284,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/pallas_call.py#L248-L284","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Remove input_output_aliases so outputs are freshly allocated, making the pallas_call alias-free before differentiating","Apply stop_gradient / custom_jvp around the kernel call so the JVP rule is never invoked on the aliased pallas_call","Differentiate an alias-free reference implementation (e.g. an equivalent plain-JAX function) instead of the in-place Pallas kernel"],"exampleFix":"# before\np = pallas_call(kernel, grid=grid, out_shape=out_shape,\n                input_output_aliases=[(0, 0)])\njvp(p, (x, t))\n# after\np = pallas_call(kernel, grid=grid, out_shape=out_shape)  # no aliasing\njvp(p, (x, t))","handlingStrategy":"validation","validationCode":"assert not input_output_aliases, 'cannot use input_output_aliases with jvp'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep a non-aliased variant of each Pallas kernel for autodiff paths","Only enable input_output_aliases on inference/performance-critical call sites"],"tags":["jax","pallas","jvp","autodiff","input-output-alias","notimplementederror"],"backgroundTag":"autodiff-unsupported-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}