Optax 0.1.4.

What's Changed

New features

• Expose adamax and adamaxw by @nasyxx in #433
• Support feeding in extra dictionary data in optax/experimental in #373
• Add NovoGrad optimizer by @DT6A in #385
• added optimistic gradient descent in #387
• Add types to utils.py by @atgctg in #367
• Add hyperparam_dtype override to hyperparam injection in #414
• added typing to linear_algebra.py by @shivance in #413
• Add amsgrad optimizer by @merajhashemi in #382
• Add Hinge Loss by @heytanay in #409

Bug fixes

• [optax] Increase chex version requirement to fix issue #456. in #457
• Start inject_hyperparams step count at 0. in #416
• Add noise before multiplying by the learning rate (noisy_sgd ) by @atgctg in #369

Miscellaneous

• Remove flags from lookahead_mnist example. in #395
• Bring alias docstrings more in line with style guide. in #406
• Bugfix in alias_test: Adamax named but not tested. in #419
• Test that all optimizers can be wrapped in inject_hyperparams. in #420
• Add an example for gradient accumulation. in #425
• [docs] Start adding numerical definitions to key parts of optax. in #430
• [optax] Add basic mnist example based on the lookahead_mnist example. in #436
• Install dependencies for dp-accounting manually. in #375
• Update documentation for AdamW in #376
• [Docs] softmax_cross_entropy_with_integer_label by @lkhphuc in #351
• Update Returns section in gradient transformations' docstrings. in #388
• Update logo and theme for the documentation in #391
• Set the test version of flax used in the equivalence test. in #398
• Add holounic to contributors list. by in #400
• Bring transform docstrings more in line with style guide. by in #405
• Update citation. in #407
• Refine the doc of sigmoid_binary_cross_entropy to not assume the meaning of last dimension. in #418
• fix integer_pow recompilation in _bias_correction by @epignatelli in #329
• Use auto instead of /proc/cpuinfo. in #454

New Contributors

• @atgctg made their first contribution in #367
• @nasyxx made their first contribution in #433
• @shivance made their first contribution in #413
• @epignatelli made their first contribution in #329
• @merajhashemi made their first contribution in #382
• @heytanay made their first contribution in #409

Full Changelog: v0.1.3...v0.1.4

Optax 0.1.3.

What's Changed

• Do not use None in optax.masked. in #338
• Implement Adamax and AdamaxW optimizers. in #342
• Add Kullback-Leibler Divergence Loss by @holounic in #309
• Add optax.softmax_cross_entropy_with_integer_labels. in #343
• Publicize private methods in transform in #348
• Update .pylintrc in #354
• Support mixture of dtypes for parameters when clipping. in #355
• Update "jax.tree_util" functions in #370

New Contributors

• @holounic made their first contribution in #309

Full Changelog: v0.1.2...v0.1.3

Optax 0.1.2

What's Changed

• Improve the documentation and discoverability of set_to_zero. by @copybara-service in #299
• Add pages to documentation for contributors. by @copybara-service in #301
• Replace _build_simple_adam with _build_stateful_sgd. by @n2cholas in #302
• Make masked preserve param structure by @n2cholas in #300
• Allow to set custom schedule for the second momentum in Adafactor. One useful use case is setting the upper bound for the second momentum, as otherwise it converges to 1.0 and effectively freezes the second momentum. It was used in https://arxiv.org/abs/2106.04560. by @copybara-service in #303
• exporting CTC loss for public use by @copybara-service in #305
• Specify that mask must be static in masked wrapper docstring. by @copybara-service in #306
• Clarify the docstring of cosine_decay_schedule. by @copybara-service in #310
• Typo in LARS docstring. by @Olamon in #311
• Fix adam's mu_dtype casting. by @copybara-service in #313
• Fix docs rendering of loss functions. by @grahamgower in #318
• Clarifies optax.adamw(weight_decay) parameter. by @copybara-service in #322
• Enhance ctc_loss to return forward probs. by @yotarok in #321
• Replace jax.tree_multimap, which is now deprecated, with jax.tree_map. by @copybara-service in #330
• Rename argument names for CTC loss functions. by @yotarok in #331

New Contributors

• @Olamon made their first contribution in #311
• @grahamgower made their first contribution in #318
• @yotarok made their first contribution in #321

Full Changelog: v0.1.1...v0.1.2

Optax 0.1.1

What's Changed

• Tweak the meta-learning example from the docs by @copybara-service in #233
• Fix small bugs in metalearning example. by @copybara-service in #236
• Do not reuse mini-batches between epochs in DPSGD example. by @copybara-service in #230
• Make the version of typing_extensions less constrained. by @copybara-service in #238
• [JAX] move example libraries from jax.experimental to jax.example_libraries by @copybara-service in #200
• Export ScaleByBeliefState by @NeilGirdhar in #239
• MultiStep optimizer: align parameter naming and type annotations of update function with signature of GradientTransform.update. by @copybara-service in #243
• Fix imports of datasets in examples folder. by @copybara-service in #242
• Enable example tests on github. Fix the bugs that were uncovered. by @copybara-service in #244
• Formatting. by @copybara-service in #249
• Add test for multi steps wrapper, verifying that the aggregated gradient is the mean of the input gradients. by @copybara-service in #255
• MultiStep optimizer wrapper: replace naive streaming average gradient implementation with numerically stabler one. by @copybara-service in #254
• Added ord, axis, and keepdims args to safe_norm by @copybara-service in #252
• Add badges and include RTD build into CI tests. by @copybara-service in #256
• Write a clearer doc-string for GradientTransformation by @copybara-service in #257
• Refactor clipping.py by @copybara-service in #260
• Implement split real norm by @wdphy16 in #241
• Monkey-patch sphinx to output correct type annotations for the most common cases (e.g. params, opt state) in the documentation. by @copybara-service in #266
• Improve docs by @copybara-service in #268
• Implement stateless wrapper. by @n2cholas in #246
• Replace _ with params to ensure you can always call init with named args. by @copybara-service in #270
• Improve docs. by @copybara-service in #269
• Add missing ` in two places. by @copybara-service in #273
• Add option to cache examples datasets after pre-processing. by @copybara-service in #272
• Fix an error in README.md rendering. by @copybara-service in #275
• Remove the old venv directory before testing the package. by @copybara-service in #289
• Fix Yogi optimizer by @wdphy16 in #288
• Bump ipython from 7.16.1 to 7.16.3 in /requirements by @dependabot in #286
• Clarifies optax.adamw(mask) parameter. by @copybara-service in #284
• Fix the link to the complex-valued optim proposal in RTD. by @copybara-service in #295
• Implement complex norm in optimizers by @wdphy16 in #279
• Change add_noise to match the target variance by scaling by its sqrt. by @Rupt in #294
• Minor tweaks to the optax documentation. by @copybara-service in #297
• Bump version to 0.1.1 from 0.1.0 by @copybara-service in #298

New Contributors

• @wdphy16 made their first contribution in #241
• @dependabot made their first contribution in #286
• @Rupt made their first contribution in #294

Full Changelog: v0.1.0...v0.1.1

Optax 0.1.0

Support for Python 3.6 has been dropped following the JAX deprecation policy. Please upgrade to a supported Python version.

Full Changelog: v0.0.91...v0.1.0

Optax 0.0.91

It is the latest version compatible with Python 3.6. See #222 for more details.

Main changes since 0.0.9

• Use an explicit type for base classes of optimizer states. #211
• Replace optax.OptState with NamedTuple, when it is used as a base class. #218
• Set optax.OptState to chex.ArrayTree. #219
• Add quick start notebook to Optax. #143
• Add Optax regression flax_example.py. by @8bitmp3 in #138
• Add scale_gradient. #198
• Add a set_to_zero transformation. #202
• Add utilities for numerical stability in the public API. #227
• Change int64 to safely incremented int32 in wrappers. by @rdaems in #187
• Add an example of meta-optimization to the docs. #231

Full Changelog: v0.0.9...v0.0.91

Optax 0.0.9

Added:

• multi_transform
• LARS optimiser
• AdaFactor optimiser

Fixed:

• masked wrapper for empty params nest
• label type in loss.py
• checks in loss.py
• MultiSteps arguments handling
• Chex asserts in constrain_test.py

Optax 0.0.8

• Added clipping_by_block_rms
• Added sgdr_schedule
• Fixed inconsistency in ema's outputs
• Added linear_schedule

Optax 0.0.6

Changelog

Full Changelog

Closed issues:

• Documentation #3

Merged pull requests:

• Fix section headers in docs #79 (KristianHolsheimer)
• Define MANIFEST.in to include requirements.txt in pypi bundle. Increment the version. #76 (copybara-service[bot])
• Fix RTD build error. #74 (copybara-service[bot])
• Improve doc-strings and formatting of popular optimiser instances - Part 1 #71 (copybara-service[bot])
• Set up sphinix-built documentation for RTD to build and serve documentation from the public Github. #61 (copybara-service[bot])

* This Changelog was automatically generated by github_changelog_generator

Optax 0.0.5

Changelog

Note: this is a first GitHub release of Optax. It includes all changes since the repo was created.

Full Changelog

Implemented enhancements:

• Implement lookahead optimiser #17
• Implement support for Yogi optimiser #9
• Implement rectified Adam #8
• Implement gradient centralisation #7
• Implement scaling by AdaBelief #6

Closed issues:

• Multiple optimizers using optax #59
• Change masked wrapper to use mask_fn instead of mask #57
• Prevent creating unnecessary momentum variables #52
• Implement Differentially Private Stochastic Gradient Descent #50
• RMSProp does not match original Tensorflow impl #49
• JITted Adam results in NaN when setting decay to integer 0 #46
• Option to not decay bias with additive_weight_decay #25
• Support specifying end_value for exponential_decay #21
• Schedules for Non-Learning Rate Hyper-parameters #20
• Implement OneCycle Learning Rate Schedule #19
• adam does not learn? #18
• Which JAX-based libraries is optax compatible with? #14
• Manually setting the learning_rate? #4

Merged pull requests:

• Fix pylint errors. #73 (copybara-service[bot])
• Add PyPI release workflow and increment the version. #70 (copybara-service[bot])
• Add flax to requirements for tests. #69 (copybara-service[bot])
• Add first flax equivalence test. #68 (copybara-service[bot])
• Targets optional in l2loss and huberloss. #67 (copybara-service[bot])
• Add .pylintrc and run pylint checks in CI workflow. #66 (copybara-service[bot])
• Increase optax version #63 (copybara-service[bot])
• Add utilities for eigenvector and matrix inverse pth root computation. #62 (copybara-service[bot])
• Add Callable option to optax.masked. #60 (n2cholas)
• Increase optax version for PyPi release. #58 (copybara-service[bot])
• Add momentum and initial_scale to RMSProp #55 (rwightman)
• Prevent creating unnecessary momentum variables. #54 (n2cholas)
• Implement DPSGD #53 (n2cholas)
• Add inject_hyperparams wrapper #48 (n2cholas)
• Format tests and parallelize pytest runs. #47 (copybara-service[bot])
• Provide a canonical implementation of canonical losses used in gradient based optimisation. #45 (copybara-service[bot])
• Expose optax transform's init and update function signatures to facilitate type annotation in user code. #44 (copybara-service[bot])
• Add a transformation and a transformation wrapper. #43 (copybara-service[bot])
• Update reference arxiv link. #41 (copybara-service[bot])
• Move equivalence tests to a separate file, as we will be adding more. #40 (copybara-service[bot])
• Optax: Add MNIST example with Adam optimizer and lookahead wrapper. #39 (copybara-service[bot])
• Optax: gradient transformation for non-negative parameters. #38 (copybara-service[bot])
• Aliases support LR schedules in addition to constant scalar LRs. #37 (copybara-service[bot])
• Optax: add datasets module for image classifier example. #36 (copybara-service[bot])
• Ensure the number of update functions and states is the same in chain. #34 (copybara-service[bot])
• Rename additive_weight_decay to add_decayed_weights. #33 (copybara-service[bot])
• Remove scale_by_fromage. #32 (copybara-service[bot])
• Add AGC to optax __init__ and add comment noting regarding 1D conv weights. #30 (copybara-service[bot])
• Clean nits to make loss and hk.transform() slightly more clear. #29 (copybara-service[bot])
• Disable macos-latest tests (to speed up CI) and add CI status badge. #28 (copybara-service[bot])
• Add a mask wrapper. #27 (n2cholas)
• Support end_value for exponential_decay #26 (n2cholas)
• Add piecewise_interpolate_schedule, linear_onecycle, and cos_onecycle. #22 (n2cholas)
• Yogi #16 (joaogui1)
• Radam #15 (joaogui1)
• gradient centralization #13 (joaogui1)
• Fix haiku_example.py #5 (asmith26)

* This Changelog was automatically generated by github_changelog_generator