mlx
mlx copied to clipboard
Add AdaBelief optimizer
- Implement AdaBelief optimizer following the NeurIPS 2020 paper
- Add unit tests following MLX standard patterns
- Update optimizer documentation to include AdaBelief
Fixes #2479
Proposed changes
- Added
AdaBeliefclass topython/mlx/optimizers/optimizers.py - Implemented with bias correction enabled by default (matching paper recommendations)
- Supports decoupled weight decay
- Added unit tests in
python/tests/test_optimizers.py - Updated documentation in
docs/src/python/optimizers/common_optimizers.rst
Checklist
- [x] I have read the CONTRIBUTING document
- [x] I have run
pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes - [x] I have added tests that prove my fix is effective or that my feature works
- [x] I have updated the necessary documentation (if needed)