XLA is the limiting factor in a lot of these cases, though maybe saying limiting factor is wrong because it's more of a "trade-off factor". XLA wants to know the static size of a lot of arguments so it can build a mathematical description of the compute graph and fuse linear algebra commands freely. What the Julia libraries like Zygote do is say "there is no good mathematical description of this, so I will generate source code instead" (and some programs like Tapenade are similar). For example, while loops are translated into for loops where a stack of the Boolean choices are stored so they can be ran in reverse during the backpass. The Julia libraries can sometimes have more trouble automatically fusing linear algebra commands though, since then they need to say "my IR lets non-static things occur, so therefore I need to prove it's static before doing transformation X". It's much easier to know you can do such transformations if anything written in the IR obeys the rules required for the transform! So it's a trade-off. In the search for allowing differentiation of any program in the language, the Julia AD tools have gone for extreme flexibility (and can rely on the fact that Julia has a compiler that can JIT compile any generated valid Julia code) and I find it really interesting to try and elucidate what you actually gain from that.