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).