[mattjjatgoogle]

You're right that downstream libraries have often tended to introduce magic (some more than others), and moreover one library's magic is typically incompatible with other libraries'. It's something that we're working on but we don't have much to show for it yet. Two avenues are:

1. as you say, exposing patterns and tools for library authors to implement transformations/higher-order primitives using JAX's machinery rather than requiring each library to introduce bespoke magic to do the same;

2. adding JAX core infrastructure which directly solves the common problems that libraries tend to solve independently (and with bespoke magic).

[6gvONxR4sf7o]

Thanks for the info! And I want to be clear that I make all these comments and questions out of a love for the tool. It's one of my favorite tools that I wish I could use at work. The design and community engagement are both fantastic.