jax-ml/jax · error · NotImplementedError

Higher-order AD not supported

Error message

Higher-order AD not supported

What it means

The forward pass wrapper for Splash Attention supports only first-order autodiff. If save_residuals=True is passed (which happens when differentiating a function that itself is being differentiated), higher-order AD is explicitly unsupported and raises NotImplementedError.

Source

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

    q: jax.Array,
    k: jax.Array,
    v: jax.Array,
    segment_ids: SegmentIds | None,
    sinks: jax.Array | None,
    save_residuals: bool,
    mask_value: float,
    is_mqa: bool,
    block_sizes: BlockSizes,
    residual_checkpoint_name: str | None,
    mask_function: MaskFunctionType | None,
    attn_logits_soft_cap: float | None = None,
    interpret: bool = False,
) -> tuple[
    jax.Array,
    SplashResidualsType,
]:
  if save_residuals:
    raise NotImplementedError("Higher-order AD not supported")

  out, (logsumexp,) = _splash_attention_forward(
      fwd_mask_info,
      q,
      k,
      v,
      segment_ids,
      sinks,
      mask_value=mask_value,
      is_mqa=is_mqa,
      block_sizes=block_sizes,
      residual_checkpoint_name=residual_checkpoint_name,
      save_residuals=True,
      mask_function=mask_function,
      attn_logits_soft_cap=attn_logits_soft_cap,
      interpret=interpret,
  )
  return out, (

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restructure to first-order AD only, e.g. use jax.vjp/jax.jvp once, or stop_gradient on the inner gradient
2. Use a different attention implementation (e.g. reference jax.nn.dot_product_attention or flash attention on GPU) inside the twice-differentiated path
3. For Hessian-free methods, use implicit differentiation or finite differences of the gradient

Example fix

// before
h = jax.hessian(loss_with_splash_attn)(params)
// after
attn_loss = remake_loss_with(attn=jax.nn.dot_product_attention)
h = jax.hessian(attn_loss)(params)

def safe_grad(fn, x):
    try:
        return jax.grad(fn)(x)
    except NotImplementedError:
        return jax.grad(lambda p: jax.grad(fn, allow_int=False)(p).sum())(x)  # not 2nd-order; see tip

try:
    h = jax.hessian(loss)(p)
except NotImplementedError as e:
    if 'Higher-order AD' in str(e):
        h = jax.hessian(replace_splash_with_reference(loss))(p)
    else:
        raise

When it happens

Trigger: Calling jax.grad twice on a loss that uses splash attention, or jax.jacfwd/hessian over a function containing splash attention; explicitly passing save_residuals=True to the internal fwd function.

Common situations: Meta-learning (MAML) second derivatives; computing Hessian-vector products of an attention-based loss; score-function estimators that differentiate through gradients.

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/d459150be344b94d. Report an issue: GitHub.