Julia's AD is compatible with control flow. They have their own issues, but Zygote + ChainRules actually work pretty well