Fix triangle_attention custom_vjp fwd/primal pytree mismatch (tuple vs list)

[ssiddhantsharma]

Summary

triangle_attention raises a TypeError under reverse-mode AD when the call sits inside a jax.checkpoint'd body run under lax.scan (the usual AlphaFold3-style Pairformer stack), on JAX versions that enforce exact pytree-structure equality between a custom_vjp fwd rule's primal output and the decorated function's output. Observed on jax==0.10.1:

TypeError: Custom VJP fwd rule triangle_attention_custom_vjp_fwd for function
triangle_attention_custom_vjp must produce a pair (list or tuple of length two)
where the first element represents the primal output ... the fwd rule output's
first element had container/pytree structure:
    (float32[...], float32[...,1], float32[...,1])
while the custom_vjp-decorated function had output container/pytree structure:
    [float32[...], float32[...,1], float32[...,1]].

Root cause

triangle_attention_custom_vjp returns fwd_p.bind(...), and bind on a multiple_results primitive returns a list. Its fwd rule (triangle_attention_custom_vjp_fwd), however, returns the primal-output element as a tuple (a, lse, amax). custom_vjp compares the fwd rule's primal-output structure to the decorated function's output; tuple != list, so JAX rejects it under AD.

The forward-only path is unaffected (it only calls bind); only jax.grad / value_and_grad fails.

Fix

One line in cuequivariance_jax/triangle/_triangle_attention.py — return a list so the fwd rule matches the primal:

-    return (a, lse, amax), residuals
+    return [a, lse, amax], residuals

Reproduction (jax 0.10.1, CUDA 12)

import jax, jax.numpy as jnp
from cuequivariance_jax import triangle_attention
B, N, H, S, D = 1, 2, 4, 16, 32
q0 = jax.random.normal(jax.random.PRNGKey(0), (B, N, H, S, D))
k  = jax.random.normal(jax.random.PRNGKey(1), (B, N, H, S, D))
v  = jax.random.normal(jax.random.PRNGKey(2), (B, N, H, S, D))
bias = jnp.zeros((B, 1, H, S, S)); mask = jnp.ones((B, N, 1, 1, S), dtype=bool)

def body(carry, _):
    return triangle_attention(carry, k, v, bias, mask, scale=1.0)[0], None

def loss(q):
    qf, _ = jax.lax.scan(jax.checkpoint(body), q, None, length=2)
    return jnp.sum(qf)

jax.grad(loss)(q0)   # TypeError before this PR; works after

A bare jax.grad of a single triangle_attention call does not trip it — the strict fwd/primal structure check fires under checkpoint/scan, which is exactly how the kernel is used inside a Pairformer trunk.

Validation

On jax==0.10.1 + cuequivariance-jax==0.10.0 + cuequivariance-ops-jax-cu12==0.10.0 (L40S):

• Before: the reproduction raises the TypeError above.
• After: jax.grad succeeds, and the fused kernel's gradient matches the module's own reference implementation (triangle_attention_jax_fwd) at cosine = 0.999997 (max|Δ| = 1.0e-1) — the fix changes nothing numerically.

Note: newer JAX (0.11.0) appears to have relaxed this fwd/primal structure check, so the crash is not observed there — but the fwd rule should still structurally match the primal, and this restores reverse-mode AD on JAX 0.10.x.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.