Stateless Transformation

Right now, there’s some boilerplate for defining simple gradient transformations. If a user wants to implement custom weight decay, clipping, constraints, etc without any state, they still have to define a nested function with an init that does nothing, and handle the empty state.

I think it’d be convenient to provide a stateless transformation that accepts a function to apply to the updates and params. We can also do the jax.tree_multimap for the user by default.

weight_decay = optax.stateless(lambda g, p: g + 0.1 * p)

If the user wants to define a function that does the tree_multimap themselves:

def my_function(updates, params):
    return jax.tree_multimap(lambda g, p: ..., updates, params)

optim = optax.stateless(my_function, on_leaves=False)

In my view, this is a very clean way to implement simple stateless transformations. I think the on_leaves argument could use a better name, though. I’d be happy to implement this if it sounds reasonable.

[n2cholas]

No more action needed on your side @mtthss, I’ll work on this over the weekend!

[mtthss]

Great, thanks @n2cholas! And apologies for the slow turnaround on this issue.