There is definitely no free lunch, it's good to really delineate the engineering trade-offs you're making! A lot of this work actually comes from the fact that

There is definitely no free lunch, it's good to really delineate the engineering trade-offs you're making! A lot of this work actually comes from the fact that some people I work with were building tools that could efficiently handle dynamic control flow without requiring tracing (see the description of Zygote.jl https://arxiv.org/abs/1810.07951 https://arxiv.org/abs/1810.07951). I had to bring up the question: why? It's much harder to build, needs more machinery, and in some cases can make less assumptions/less fusions (a general form of vmap is much harder for example if you cannot trace, see KernelAbstractions.jl for details). This line of inquiry led an example of why you might want to support such dynamic behaviors, so I'll leave it up to someone else to declare whether the maintenance or complexity cost is worth it to them. I wouldn't say that this means Jax or Tensorflow are doomed (far from it: simple ML architectures are quasi-static, so it's building for the correct audience), but it's good to know what exactly you're leaving out when you make a simplifying assumption.