This is just speculative discussion from my own experiences, but I appreciate that the metaprogramming constructs are restricted to operate on pure functions, v

This is just speculative discussion from my own experiences, but I appreciate that the metaprogramming constructs are restricted to operate on pure functions, vs attempting to handle things like (e.g.) Python data declarations, or syntactical transformations. For one, it's not clear to me that you should attempt to have 1 system which handles both syntax-level expansion (macros) and function transformation (what this system, and JAX, does) -- unless you're writing a Lisp. That's not Python -- so I'm sort of glad that these systems don't try to handle the former, and instead make a nice go at the latter. The restriction to code which is "pure" (where pure is restricted to a definition which makes sense for these array languages) is also nice from a program transformation perspective. Otherwise, authors of transformations must reason about effects, etc. JAX has started to include utilities to help with this sort of reasoning for custom (read: user-defined) primitives.