Prevent unnecessary recompilation due to closures in optimizers

Now that Jax supports dataclasses as PyTrees, would it be possible to switch to using them instead of namedtuple? The benefits are explained here.

The biggest benefit would be preventing unnecessary recompilation. The current Optax code uses closures, which will cause Jax to unnecessarily recompile a jitted function that accepts a GradientTransformation. (The closures are different objects that hash differently, which means that changing the parameters to the GradientTransformation must cause the jitted function to recompile.)

A dataclass version of Optax would look something like this.

I am happy to do submit a pull request if this change is okay.

[rosshemsley]

(Just writing to reassure you that we have not forgotten about this, it’s just a busy week over here 😃 )

[rosshemsley]

Hey @NeilGirdhar,

We have discussed these proposals extensively within the optax maintainers team, and gathered feedback from a number of users and developers. We believe that the approach outlined below would be the best way forward,

The request here is for JAX to treat instances of optax.GradientTransformation as pytrees. In @NeilGirdhar’s proposal, this means combining optimizer hyperparameters with the init / update functions into a single dataclass. In @cgarciae’s proposal the state is additionally also included in the dataclass.

Making gradient transformations into pytree classes provides the ability to pass these objects to JAX jitted functions as dynamic rather than static arguments. On balance, the idea is reasonable and might be easier to understand for users who like OO patterns and the specific dataclass-based strand of OO programming that is being used in some JAX codebases. There are also downsides to this approach though - notably, the current optax API has a very sharp boundary between “functions implementing user behavior” (the init / update methods) and “data containers’’ (i.e. the state). In the above proposals, hyperparameters and/or state are “folded into” the optimizer class, and the pure functions become methods on a mutable class instead. This is not a problem per se (it’s a common pattern in OO!), but it is different from the programming patterns of many of our users.

So what is the way forward? As Optax developers, we do care a lot about users being able to extend and modify the optax components. The whole design focuses on composability and extensibility. We believe the best way forwards here is thus not to use dataclasses in each individual component, but rather to create wrappers that turn optax optimisers into dataclasses with the desired behavior.

This is the approach taken by @n2cholas in his inject_hyperparams (already upstreamed into optax) and is the approach taken by @cgarciae (in Treex). Even when living in libraries outside of optax, these wrappers are lighter to maintain than the current approach in tjax: for instance in @cgarciae’s Treex only one wrapper is maintained, and all optax transformations are exposed after wrapping them programmatically (instead of defining new separate classes for each optax component as in tjax, where new code to be written every time components are added to optax).

Our suggestion is thus to proceed in two stages In an initial stage tjax could define its own wrapper (similarly to Treex). Once a common dataclass implementation is provided by JAX itself we would be happy to see extensions of this kind upstreamed into optax (as we were to see inject_hyperparams).

E.g. we could imagine exposing a triplet of extensions: inject_hyperparams (@n2cholas’s wrapper) wrap_transformation_as_pytree (implementing @NeilGirdhar’s extension) wrap_transformation_as_pytree_with_state (implementing @cgarciae’s extension)

We hope the above is acceptable, and while this might not be everybody’s favorite solution, we want to stress that we really appreciate everybody’s feedback and contributions, and we hope to continue seeing you all engaged in the development of optax and contributing to the library.