> Some of this complexity may be necessary for achieving optimal performance in Jax. E.g. extra indirection to avoid the compiler making some bad fusion decision, or multiple calls so something can be marked as static for the jit in the outer call

certainly some of it is but not the lion's share - I have a much simpler (private) codebase which scales pretty similarly afaict.

the complexity of Maxtext feels more Serious Engineering ™ flavored, following Best Practices.