My bad I didn’t verify what I pasted.
No. I didn’t get any errors when I called gradient of either fast_max or logsumexp_avx.
Surprisingly, if I remove my gradient definition for fast_max the gradients of logsumexp_avx and logsumexp_no_avx matches. I don’t know what’s going on.