AdamW implementation, without correction

[jbarrow]

Implementation of AdamW, but without correction of the first and second moments. Following the convention of the Adam implementation in the repo.

If you want the paper convention of correction, that implementation is here:

https://github.com/jbarrow/mlxllama/blob/cb2a808eefafc340db91e7d8f97ca334f6afff2a/mlxllama/optim.py

[angeloskath]

That looks great.

Do you mind instead of repeating the Adam implementation to extend Adam and do sth like calling

return super().apply_single(gradient, parameter * (1 - lr * wd))

This way if we want to change something down the line we only need to change it in one place.

[jbarrow]

Alright, made that change and tested it locally in mlxllama

[angeloskath]

Looking good!