The problem is that you were implicitly doing a struct of array to array of structs transformation by splatting the a. That’s going to be really slow anyways. I would suggest using the recommended Zygote, since that will use the adjoint overloads. See this page for details:
1 Like