jax-ml/jax · error · NotImplementedError

Higher-order AD not supported.

Error message

Higher-order AD not supported.

What it means

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.

Source

Thrown at jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py:251

      custom_type=custom_type,
      attn_logits_soft_cap=attn_logits_soft_cap,
  )


def _attention_reference_custom_fwd(
    mask: jax.Array,  # [q_seq_len, kv_seq_len]
    q: jax.Array,  # [q_seq_len, head_dim]
    k: jax.Array,  # [kv_seq_len, head_dim]
    v: jax.Array,  # [kv_seq_len, head_dim]
    segment_ids: SegmentIds | None,
    sinks: jax.Array | None,
    mask_value: float,
    save_residuals: bool,
    custom_type: str,
    attn_logits_soft_cap: float | None,
):
  if save_residuals:
    raise NotImplementedError("Higher-order AD not supported.")

  o, (logsumexp,) = _attention_reference(
      mask,
      q,
      k,
      v,
      segment_ids,
      sinks,
      mask_value=mask_value,
      save_residuals=True,
      custom_type=custom_type,
      attn_logits_soft_cap=attn_logits_soft_cap,
  )
  return o, (mask, q, k, v, segment_ids, sinks, o, logsumexp)


def _attention_reference_custom_bwd(
    mask_value: float,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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
2. Switch to a reference/manual attention implementation (non-Pallas) for the inner gradient so double differentiation is well-defined
3. Restructure to avoid second-order AD through attention entirely (e.g. finite differences for the outer derivative)

Example fix

// before
hvp = jax.grad(lambda eps: jax.grad(loss)(x + eps*vec))  # double AD through splash attention
// after
inner = jax.grad(lambda z: loss_with_attn_stopgrad(z))  # wrap attention with lax.stop_gradient for second-order paths
hvp = jax.grad(lambda eps: inner(x + eps*vec))

def is_first_order_only(fn) -> bool:
    return getattr(getattr(fn, 'fn', None) or fn, '__wrapped_first_order__', False) or 'splash' in getattr(fn, '__qualname__', '')

try:
    hvp = jax.grad(lambda e: jax.grad(loss)(x + e * v))
except NotImplementedError as e:
    if 'Higher-order AD' in str(e):
        hvp = finite_diff_hvp(loss, x, v)  # numerical fallback
    else:
        raise

When it happens

Trigger: 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).

Common situations: 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.

Related errors

• Higher-order AD not supported
• Block sizes for dq kernel are not needed with a fused kernel
• Mask function must return a boolean-valued array, but got: {
• block_kv must be a multiple of {NUM_LANES}
• block_q must be a multiple of {NUM_LANES}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/317a27c2a3afc189. Report an issue: GitHub.