Zygote debugging

In the end, Zygote could not handle the code and I had to write my own adjoint. Maybe I will take a look at Zygote in the future.