{"record":{"id":"317a27c2a3afc189","repo":"jax-ml/jax","slug":"higher-order-ad-not-supported","errorCode":null,"errorMessage":"Higher-order AD not supported.","messagePattern":"Higher-order AD not supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":251,"sourceCode":"      custom_type=custom_type,\n      attn_logits_soft_cap=attn_logits_soft_cap,\n  )\n\n\ndef _attention_reference_custom_fwd(\n    mask: jax.Array,  # [q_seq_len, kv_seq_len]\n    q: jax.Array,  # [q_seq_len, head_dim]\n    k: jax.Array,  # [kv_seq_len, head_dim]\n    v: jax.Array,  # [kv_seq_len, head_dim]\n    segment_ids: SegmentIds | None,\n    sinks: jax.Array | None,\n    mask_value: float,\n    save_residuals: bool,\n    custom_type: str,\n    attn_logits_soft_cap: float | None,\n):\n  if save_residuals:\n    raise NotImplementedError(\"Higher-order AD not supported.\")\n\n  o, (logsumexp,) = _attention_reference(\n      mask,\n      q,\n      k,\n      v,\n      segment_ids,\n      sinks,\n      mask_value=mask_value,\n      save_residuals=True,\n      custom_type=custom_type,\n      attn_logits_soft_cap=attn_logits_soft_cap,\n  )\n  return o, (mask, q, k, v, segment_ids, sinks, o, logsumexp)\n\n\ndef _attention_reference_custom_bwd(\n    mask_value: float,","sourceCodeStart":233,"sourceCodeEnd":269,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L233-L269","documentation":"Splash attention's reference implementation defines a custom VJP for backward passes; when save_residuals=True is requested inside the custom forward (i.e. differentiating the backward pass), higher-order derivatives are unsupported and it raises NotImplementedError. Second-order AD through this fused TPU attention is not implemented.","triggerScenarios":"Computing Hessian-vector products or any second-order differentiation through splash_attention (e.g. jax.jacrev(jax.grad(loss)) where the loss uses splash attention on TPU).","commonSituations":"Influence-function / curvature estimation research code, score-function methods, or gradient-penalty losses that call grad twice and unknowingly route through the TPU splash attention kernel.","solutions":["Stop-gradient the attention output in the inner derivative so the outer derivative does not traverse the custom VJP: use jax.lax.stop_gradient on attention outputs or on the residual path","Switch to a reference/manual attention implementation (non-Pallas) for the inner gradient so double differentiation is well-defined","Restructure to avoid second-order AD through attention entirely (e.g. finite differences for the outer derivative)"],"exampleFix":"// before\nhvp = jax.grad(lambda eps: jax.grad(loss)(x + eps*vec))  # double AD through splash attention\n// after\ninner = jax.grad(lambda z: loss_with_attn_stopgrad(z))  # wrap attention with lax.stop_gradient for second-order paths\nhvp = jax.grad(lambda eps: inner(x + eps*vec))","handlingStrategy":"fallback","validationCode":"def is_first_order_only(fn) -> bool:\n    return getattr(getattr(fn, 'fn', None) or fn, '__wrapped_first_order__', False) or 'splash' in getattr(fn, '__qualname__', '')","typeGuard":null,"tryCatchPattern":"try:\n    hvp = jax.grad(lambda e: jax.grad(loss)(x + e * v))\nexcept NotImplementedError as e:\n    if 'Higher-order AD' in str(e):\n        hvp = finite_diff_hvp(loss, x, v)  # numerical fallback\n    else:\n        raise","preventionTips":["Wrap splash-attention outputs with jax.lax.stop_gradient in code that may be differentiated twice","Audit grad-of-grad pipelines for fused TPU kernels","Fall back to a reference attention implementation for second-order AD paths"],"tags":["jax","pallas","tpu","splash-attention","autodiff","higher-order-ad"],"backgroundTag":"higher-order-ad-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}