Are there higher-level routines for gradient-based optimizations?

[idnm]

My primary use-case for JAX is to perform gradient-based optimization. I am aware of optax which implements many of the popular optimizers. However, to optimize a given loss function using, say, ADAM optimizer, multiple further lines of code need to be written such as initialization, a complete jit-compilible update step, and some kind of lax loop wrapper for the latter. I find myself recreating this setup often in different scenarios and wonder if maybe there already exists a higher level API implementing the whole pipeline?

[yoeldr]

I think JAXline may suit your needs: https://github.com/deepmind/jaxline

[hawkinsp]

We didn't want to be prescriptive about training loops in JAX, and thought it would be better for us to provide no abstraction at all rather than a complicated abstraction that no-one likes. There are libraries built on top of JAX that are more opinionated. jaxline is one example. flax also has many examples of training loops, but flax mostly subscribes to a "you should fork the training loop and modify" philosophy, rather than building a library abstraction for it.