From 4c1845dc49596fd48d373d75465a900fa2cd6c01 Mon Sep 17 00:00:00 2001 From: siddhant Date: Wed, 22 Jul 2026 16:26:39 +0530 Subject: [PATCH] Fix triangle_attention custom_vjp fwd/primal pytree mismatch (tuple vs list) `triangle_attention_custom_vjp` returns `fwd_p.bind(...)`, and for a multiple_results primitive `bind` returns a list. Its custom_vjp fwd rule, however, returned the primal-output element as a tuple `(a, lse, amax)`. JAX versions that enforce exact pytree-structure equality between a custom_vjp fwd rule's primal output and the decorated function's output reject this under reverse-mode AD: TypeError: Custom VJP fwd rule triangle_attention_custom_vjp_fwd ... the fwd rule output's first element had container/pytree structure (tuple ...) while the custom_vjp-decorated function had output container/pytree structure [list ...]. The forward path is unaffected (only bind is called there); only jax.grad / value_and_grad through triangle_attention fails. Return a list so the fwd rule's primal output matches the primal function's output. Signed-off-by: siddhant --- .../cuequivariance_jax/triangle/_triangle_attention.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/cuequivariance_jax/cuequivariance_jax/triangle/_triangle_attention.py b/cuequivariance_jax/cuequivariance_jax/triangle/_triangle_attention.py index 03a88c87..27f7fddd 100644 --- a/cuequivariance_jax/cuequivariance_jax/triangle/_triangle_attention.py +++ b/cuequivariance_jax/cuequivariance_jax/triangle/_triangle_attention.py @@ -263,7 +263,11 @@ def triangle_attention_custom_vjp( def triangle_attention_custom_vjp_fwd(q, k, v, bias, mask, scale, precision=None): a, lse, amax = fwd_p.bind(q, k, v, bias, mask, scale=scale, precision=precision) residuals = (a, lse, q, k, v, bias, mask) - return (a, lse, amax), residuals + # Match the primal output's pytree structure: `triangle_attention_custom_vjp` + # returns `fwd_p.bind(...)`, and a multiple_results primitive's bind returns a + # list. Returning a tuple here trips custom_vjp's fwd/primal structure check + # on JAX versions that enforce exact tuple-vs-list equality (TypeError). + return [a, lse, amax], residuals def triangle_attention_custom_vjp_bwd(scale, precision, residuals, cotangents):