LICENSE
MANIFEST.in
README.md
setup.py
docs/conf.py
docs/ext/coverage_check.py
examples/datasets.py
examples/datasets_test.py
examples/haiku_example.py
examples/lookahead_mnist.py
examples/lookahead_mnist_test.py
optax/__init__.py
optax/optax_test.py
optax.egg-info/PKG-INFO
optax.egg-info/SOURCES.txt
optax.egg-info/dependency_links.txt
optax.egg-info/not-zip-safe
optax.egg-info/requires.txt
optax.egg-info/top_level.txt
optax/_src/alias.py
optax/_src/alias_test.py
optax/_src/combine.py
optax/_src/combine_test.py
optax/_src/control_variates.py
optax/_src/control_variates_test.py
optax/_src/equivalence_test.py
optax/_src/float64_test.py
optax/_src/linear_algebra.py
optax/_src/linear_algebra_test.py
optax/_src/loss.py
optax/_src/loss_test.py
optax/_src/privacy.py
optax/_src/privacy_test.py
optax/_src/schedule.py
optax/_src/schedule_test.py
optax/_src/second_order.py
optax/_src/second_order_test.py
optax/_src/stochastic_gradient_estimators.py
optax/_src/stochastic_gradient_estimators_test.py
optax/_src/test_utils.py
optax/_src/transform.py
optax/_src/transform_test.py
optax/_src/update.py
optax/_src/update_test.py
optax/_src/utils.py
optax/_src/wrappers.py
optax/_src/wrappers_test.py
requirements/requirements-examples.txt
requirements/requirements-test.txt
requirements/requirements.txt