What's Changed
- Beginning of 0.2.4 development by @copybara-service in #1003
- Fix gallery display, add images by @vroulet in #1005
- Fix docs for dog by @jungtaekkim in #1008
- docs: remove
multi_normalfrom utilities.rst by @Abhinavcode13 in #1009 - Fix docs for dowg by @jungtaekkim in #1013
- feat: add a mathematical description of AdaGrad optimizer by @Abhinavcode13 in #1011
- fix: refactor AdaGrad optimizer recent changes by @Abhinavcode13 in #1016
- Make masking compatible with callable pytrees a la Equinox by @copybara-service in #1015
- Enable zero lrate for schedule-free optimization by @copybara-service in #1018
- keep a local .pylintrc file by @fabianp in #1024
- Add bias_correction and eps_in_sqrt options to rmsprop and associated transforms by @copybara-service in #1019
- Replace adam(b1=0) by rmsprop for schedule_free by @copybara-service in #1025
- Update init value for zakharov problem from 1e4 to 1e3 by @copybara-service in #1027
- Fix typo by @gil2rok in #1030
- Updated docs cosine_decay_schedule by @bhargavyagnik in #1032
- feat: add mathematical notation docs of SM3 optimizer by @Abhinavcode13 in #1012
- DOC: misc improvements in docstring of softmax_cross_entropy* by @copybara-service in #1033
- add doctest to constant_schedule by @fabianp in #1034
- Add axis and where arguments to loss functions. by @carlosgmartin in #912
- Fix doctest error in make_fenchel_young_loss by @copybara-service in #1035
- add doctest for polynomial_schedule by @fabianp in #1037
- add missing schedule_free_* methods by @fabianp in #1043
- fix error in softmax_cross_entropy formula by @fabianp in #1041
- Fix typo in formula of cosine_decay_schedule by @fabianp in #1044
- schedule_free: fix broadcasting of scalar arrays to 1d arrays by @n-gao in #1042
- Update polynomial_schedule doctest per @vroulet's feedback by @fabianp in #1045
- Fix linting schedule_free_test by @copybara-service in #1048
- more robust tests by @copybara-service in #1050
- Generalizes safe_int32_increment to safe_increment by @copybara-service in #1054
- Add dtype option to tree_random_like by @copybara-service in #1056
- Add double precision tests for safe_increment and fix warnings on float64_test.py by @copybara-service in #1055
- Add optax.tree_utils.tree_random_split. by @carlosgmartin in #1063
- Fix test.sh, which uses set -u, so that it works when JAX_VERSION is unset. by @carlosgmartin in #1070
- Migrate from jax.tree_util legacy APIs to new jax.tree API. by @carlosgmartin in #1066
- Ensure optimizers return updates of same dtype as params. by @copybara-service in #1060
- Fix test.sh to not modify .pylintrc. by @carlosgmartin in #1071
- Replace deprecated typing.Hashable with collections.abc.Hashable. by @carlosgmartin in #1068
- Relax absolute tolerance for failing tests involving chex.assert_trees_all_close. by @carlosgmartin in #1069
- Fix doctests by @copybara-service in #1073
- Tidy up test.sh and make it clean up properly. by @carlosgmartin in #1074
- Add missing initializer argument of 0 to tree_reduce in tree_vdot and tree_sum. by @carlosgmartin in #1065
- bump chex version for #1076 by @fabianp in #1078
- correct RST references by @fabianp in #1079
- deprecate methods in
optax/monte_carlo. by @copybara-service in #1076 - schedule-free optimier: ensure it's possible to donate both the state and the params by @enolan in #1059
- add link to examples from docstring by @fabianp in #1085
- Adopt US spelling for documentation and fix typos by @miguelcsx in #1087
- Update docs: Note RMSprop usage instead of Adam for memory savings in… by @nasyxx in #1086
- adding a perturbations module. Can take pytrees as inputs by @copybara-service in #827
- Fix initial step of scale_by_optimistic_gradient. by @carlosgmartin in #1088
- Add Hungarian algorithm for the linear assignment problem. by @carlosgmartin in #1083
- Allow safe_increment to handle unsigned integers. by @carlosgmartin in #1092
- Fix formatting issues with gallery entry for linear assignment problem. by @carlosgmartin in #1095
- use sphinx references instead of hardcoded links. by @fabianp in #1096
- Remove dtype safeguards by @copybara-service in #1099
- cosmetic improvements perturbations module by @fabianp in #1097
- update jax.tree.map to comply with jax 0.4.34 by @a1302z in #1094
- Add Adan optimizer. by @carlosgmartin in #1090
- Fix typo in projection_simplex docstring. by @copybara-service in #1105
- add config for link checking, and mark 429 (too many requests) as fine by @fabianp in #1103
- Fix docstring for hungarian_algorithm. by @carlosgmartin in #1102
- Add optax.optimistic_adam. by @carlosgmartin in #1089
- Add projection_l1_sphere and projection_l1_ball. by @copybara-service in #1106
- Add projection_l2_sphere and projection_l2_ball. by @copybara-service in #1114
- Add tree_max. by @copybara-service in #1115
- Add projection_linf_ball. by @copybara-service in #1117
- remove test that leaked jax tracers by @copybara-service in #1123
- Add a mathematical description for Lion by @aman2304 in #1121
- Fix the sign for the update in the math equation for nadam in the docs. by @carlosgmartin in #1128
- ntxent fix by @GrantMcConachie in #946
- Add Nesterov momentum to AdaBelief optimizer. by @carlosgmartin in #1127
- fix: Coherent dtypes of updates with and without MultiSteps by @hlzl in #1122
- Fix AdaBelief implementation. by @carlosgmartin in #1130
- Revisiting linesearches and LBFGS. by @copybara-service in #1133
New Contributors
- @vroulet made their first contribution in #1005
- @jungtaekkim made their first contribution in #1008
- @Abhinavcode13 made their first contribution in #1009
- @gil2rok made their first contribution in #1030
- @bhargavyagnik made their first contribution in #1032
- @n-gao made their first contribution in #1042
- @enolan made their first contribution in #1059
- @miguelcsx made their first contribution in #1087
- @a1302z made their first contribution in #1094
- @aman2304 made their first contribution in #1121
- @hlzl made their first contribution in #1122
Full Changelog: v0.2.3...v0.2.4