Thanks! One last thing, since I have your ear. The function transformation aspects of jax seem to make their way into downstream libraries like haiku, resulting in a lot of "magic" that can be difficult to examine and debug. Are there any utils you made to make jax's own transformations more transparent, which you think might be helpful to third party transformations?
Higher order functions are difficult in general, and it would be fantastic to have core patterns or tools for breaking them open.
Higher order functions are difficult in general, and it would be fantastic to have core patterns or tools for breaking them open.