├── .github ├── check_license_headers.py └── workflows │ ├── mlc_config.json │ ├── pypi-publish.yml │ └── tests.yml ├── .gitignore ├── .pre-commit-config.yaml ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── docs ├── .gitignore ├── Makefile ├── _static │ └── css │ │ └── custom.css ├── api │ ├── apply_updates.rst │ ├── assignment.rst │ ├── combining_optimizers.rst │ ├── contrib.rst │ ├── experimental.rst │ ├── losses.rst │ ├── optimizer_schedules.rst │ ├── optimizer_wrappers.rst │ ├── optimizers.rst │ ├── perturbations.rst │ ├── projections.rst │ ├── stochastic_gradient_estimators.rst │ ├── transformations.rst │ └── utilities.rst ├── conf.py ├── development.md ├── ext │ └── coverage_check.py ├── gallery.rst ├── getting_started.ipynb ├── images │ ├── examples │ │ ├── adversarial.png │ │ ├── cifar10_resnet.png │ │ ├── contrib │ │ │ ├── ademamix_rosenbrock.png │ │ │ ├── reduce_on_plateau.png │ │ │ └── sam.png │ │ ├── flax_optax.png │ │ ├── lbfgs.png │ │ ├── linear_assignment_problem.png │ │ ├── lookahead.png │ │ ├── mnist.png │ │ ├── ogda.png │ │ ├── perturbations.png │ │ └── tiny_shakespeare.png │ ├── favicon.svg │ └── logo.svg ├── index.rst └── pyproject.toml ├── examples ├── README.md ├── adversarial_training.ipynb ├── cifar10_resnet.ipynb ├── contrib │ ├── README.md │ ├── differentially_private_sgd.ipynb │ ├── muon.ipynb │ ├── reduce_on_plateau.ipynb │ ├── rosenbrock_ademamix.ipynb │ └── sam.ipynb ├── flax_example.ipynb ├── freezing_parameters.ipynb ├── gradient_accumulation.ipynb ├── lbfgs.ipynb ├── linear_assignment_problem.ipynb ├── lookahead_mnist.ipynb ├── meta_learning.ipynb ├── mlp_mnist.ipynb ├── nanolm.ipynb ├── ogda_example.ipynb ├── perturbations.ipynb └── pyproject.toml ├── optax ├── CHANGELOG.md ├── __init__.py ├── _src │ ├── alias.py │ ├── alias_test.py │ ├── base.py │ ├── base_test.py │ ├── combine.py │ ├── combine_test.py │ ├── constrain.py │ ├── deprecations.py │ ├── factorized.py │ ├── factorized_test.py │ ├── float64_test.py │ ├── linear_algebra.py │ ├── linear_algebra_test.py │ ├── linesearch.py │ ├── linesearch_test.py │ ├── lookahead.py │ ├── lookahead_test.py │ ├── numerics.py │ ├── numerics_test.py │ ├── schedule.py │ ├── sharding_test.py │ ├── test_utils.py │ ├── transform.py │ ├── transform_test.py │ ├── update.py │ ├── update_test.py │ ├── utils.py │ ├── utils_test.py │ └── wrappers.py ├── assignment │ ├── __init__.py │ ├── _hungarian_algorithm.py │ └── _hungarian_algorithm_test.py ├── contrib │ ├── __init__.py │ ├── _acprop.py │ ├── _ademamix.py │ ├── _adopt.py │ ├── _cocob.py │ ├── _common_test.py │ ├── _complex_valued.py │ ├── _complex_valued_test.py │ ├── _dadapt_adamw.py │ ├── _dog.py │ ├── _mechanic.py │ ├── _mechanic_test.py │ ├── _momo.py │ ├── _muon.py │ ├── _muon_test.py │ ├── _privacy.py │ ├── _privacy_test.py │ ├── _prodigy.py │ ├── _reduce_on_plateau.py │ ├── _reduce_on_plateau_test.py │ ├── _sam.py │ ├── _sam_test.py │ ├── _schedule_free.py │ ├── _schedule_free_test.py │ └── _sophia.py ├── experimental │ ├── __init__.py │ ├── _aggregating.py │ ├── _aggregating_test.py │ ├── _microbatching.py │ └── _microbatching_test.py ├── losses │ ├── __init__.py │ ├── _classification.py │ ├── _classification_test.py │ ├── _fenchel_young.py │ ├── _fenchel_young_test.py │ ├── _ranking.py │ ├── _ranking_test.py │ ├── _regression.py │ ├── _regression_test.py │ ├── _segmentation.py │ ├── _segmentation_test.py │ ├── _self_supervised.py │ ├── _self_supervised_test.py │ ├── _smoothing.py │ └── _smoothing_test.py ├── optax_test.py ├── perturbations │ ├── __init__.py │ ├── _make_pert.py │ └── _make_pert_test.py ├── projections │ ├── __init__.py │ ├── _projections.py │ └── _projections_test.py ├── schedules │ ├── __init__.py │ ├── _inject.py │ ├── _inject_test.py │ ├── _join.py │ ├── _join_test.py │ ├── _schedule.py │ ├── _schedule_test.py │ └── inject.py ├── second_order │ ├── __init__.py │ └── _deprecated.py ├── transforms │ ├── __init__.py │ ├── _accumulation.py │ ├── _accumulation_test.py │ ├── _adding.py │ ├── _adding_test.py │ ├── _clipping.py │ ├── _clipping_test.py │ ├── _combining.py │ ├── _combining_test.py │ ├── _conditionality.py │ ├── _conditionality_test.py │ ├── _constraining.py │ ├── _constraining_test.py │ ├── _freezing.py │ ├── _freezing_test.py │ ├── _layouts.py │ ├── _layouts_test.py │ ├── _masking.py │ ├── _masking_test.py │ ├── _monitoring.py │ └── _monitoring_test.py ├── tree │ └── __init__.py └── tree_utils │ ├── __init__.py │ ├── _casting.py │ ├── _casting_test.py │ ├── _random.py │ ├── _random_test.py │ ├── _state_utils.py │ ├── _state_utils_test.py │ ├── _tree_math.py │ └── _tree_math_test.py ├── pyproject.toml ├── readthedocs.yml └── test.sh /.github/check_license_headers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/.github/check_license_headers.py -------------------------------------------------------------------------------- /.github/workflows/mlc_config.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/.github/workflows/mlc_config.json -------------------------------------------------------------------------------- /.github/workflows/pypi-publish.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/.github/workflows/pypi-publish.yml -------------------------------------------------------------------------------- /.github/workflows/tests.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/.github/workflows/tests.yml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/.gitignore -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/.pre-commit-config.yaml -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/CONTRIBUTING.md -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/README.md -------------------------------------------------------------------------------- /docs/.gitignore: -------------------------------------------------------------------------------- 1 | _build 2 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/Makefile -------------------------------------------------------------------------------- /docs/_static/css/custom.css: -------------------------------------------------------------------------------- 1 | div.math { 2 | flex-direction: row; 3 | } 4 | -------------------------------------------------------------------------------- /docs/api/apply_updates.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/apply_updates.rst -------------------------------------------------------------------------------- /docs/api/assignment.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/assignment.rst -------------------------------------------------------------------------------- /docs/api/combining_optimizers.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/combining_optimizers.rst -------------------------------------------------------------------------------- /docs/api/contrib.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/contrib.rst -------------------------------------------------------------------------------- /docs/api/experimental.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/experimental.rst -------------------------------------------------------------------------------- /docs/api/losses.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/losses.rst -------------------------------------------------------------------------------- /docs/api/optimizer_schedules.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/optimizer_schedules.rst -------------------------------------------------------------------------------- /docs/api/optimizer_wrappers.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/optimizer_wrappers.rst -------------------------------------------------------------------------------- /docs/api/optimizers.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/optimizers.rst -------------------------------------------------------------------------------- /docs/api/perturbations.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/perturbations.rst -------------------------------------------------------------------------------- /docs/api/projections.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/projections.rst -------------------------------------------------------------------------------- /docs/api/stochastic_gradient_estimators.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/stochastic_gradient_estimators.rst -------------------------------------------------------------------------------- /docs/api/transformations.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/transformations.rst -------------------------------------------------------------------------------- /docs/api/utilities.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/api/utilities.rst -------------------------------------------------------------------------------- /docs/conf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/conf.py -------------------------------------------------------------------------------- /docs/development.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/development.md -------------------------------------------------------------------------------- /docs/ext/coverage_check.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/ext/coverage_check.py -------------------------------------------------------------------------------- /docs/gallery.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/gallery.rst -------------------------------------------------------------------------------- /docs/getting_started.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/getting_started.ipynb -------------------------------------------------------------------------------- /docs/images/examples/adversarial.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/examples/adversarial.png -------------------------------------------------------------------------------- /docs/images/examples/cifar10_resnet.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/examples/cifar10_resnet.png -------------------------------------------------------------------------------- /docs/images/examples/contrib/ademamix_rosenbrock.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/examples/contrib/ademamix_rosenbrock.png -------------------------------------------------------------------------------- /docs/images/examples/contrib/reduce_on_plateau.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/examples/contrib/reduce_on_plateau.png -------------------------------------------------------------------------------- /docs/images/examples/contrib/sam.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/examples/contrib/sam.png -------------------------------------------------------------------------------- /docs/images/examples/flax_optax.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/examples/flax_optax.png -------------------------------------------------------------------------------- /docs/images/examples/lbfgs.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/examples/lbfgs.png -------------------------------------------------------------------------------- /docs/images/examples/linear_assignment_problem.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/examples/linear_assignment_problem.png -------------------------------------------------------------------------------- /docs/images/examples/lookahead.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/examples/lookahead.png -------------------------------------------------------------------------------- /docs/images/examples/mnist.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/examples/mnist.png -------------------------------------------------------------------------------- /docs/images/examples/ogda.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/examples/ogda.png -------------------------------------------------------------------------------- /docs/images/examples/perturbations.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/examples/perturbations.png -------------------------------------------------------------------------------- /docs/images/examples/tiny_shakespeare.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/examples/tiny_shakespeare.png -------------------------------------------------------------------------------- /docs/images/favicon.svg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/favicon.svg -------------------------------------------------------------------------------- /docs/images/logo.svg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/images/logo.svg -------------------------------------------------------------------------------- /docs/index.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/index.rst -------------------------------------------------------------------------------- /docs/pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/docs/pyproject.toml -------------------------------------------------------------------------------- /examples/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/README.md -------------------------------------------------------------------------------- /examples/adversarial_training.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/adversarial_training.ipynb -------------------------------------------------------------------------------- /examples/cifar10_resnet.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/cifar10_resnet.ipynb -------------------------------------------------------------------------------- /examples/contrib/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/contrib/README.md -------------------------------------------------------------------------------- /examples/contrib/differentially_private_sgd.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/contrib/differentially_private_sgd.ipynb -------------------------------------------------------------------------------- /examples/contrib/muon.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/contrib/muon.ipynb -------------------------------------------------------------------------------- /examples/contrib/reduce_on_plateau.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/contrib/reduce_on_plateau.ipynb -------------------------------------------------------------------------------- /examples/contrib/rosenbrock_ademamix.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/contrib/rosenbrock_ademamix.ipynb -------------------------------------------------------------------------------- /examples/contrib/sam.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/contrib/sam.ipynb -------------------------------------------------------------------------------- /examples/flax_example.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/flax_example.ipynb -------------------------------------------------------------------------------- /examples/freezing_parameters.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/freezing_parameters.ipynb -------------------------------------------------------------------------------- /examples/gradient_accumulation.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/gradient_accumulation.ipynb -------------------------------------------------------------------------------- /examples/lbfgs.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/lbfgs.ipynb -------------------------------------------------------------------------------- /examples/linear_assignment_problem.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/linear_assignment_problem.ipynb -------------------------------------------------------------------------------- /examples/lookahead_mnist.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/lookahead_mnist.ipynb -------------------------------------------------------------------------------- /examples/meta_learning.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/meta_learning.ipynb -------------------------------------------------------------------------------- /examples/mlp_mnist.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/mlp_mnist.ipynb -------------------------------------------------------------------------------- /examples/nanolm.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/nanolm.ipynb -------------------------------------------------------------------------------- /examples/ogda_example.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/ogda_example.ipynb -------------------------------------------------------------------------------- /examples/perturbations.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/perturbations.ipynb -------------------------------------------------------------------------------- /examples/pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/examples/pyproject.toml -------------------------------------------------------------------------------- /optax/CHANGELOG.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/CHANGELOG.md -------------------------------------------------------------------------------- /optax/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/__init__.py -------------------------------------------------------------------------------- /optax/_src/alias.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/alias.py -------------------------------------------------------------------------------- /optax/_src/alias_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/alias_test.py -------------------------------------------------------------------------------- /optax/_src/base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/base.py -------------------------------------------------------------------------------- /optax/_src/base_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/base_test.py -------------------------------------------------------------------------------- /optax/_src/combine.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/combine.py -------------------------------------------------------------------------------- /optax/_src/combine_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/combine_test.py -------------------------------------------------------------------------------- /optax/_src/constrain.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/constrain.py -------------------------------------------------------------------------------- /optax/_src/deprecations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/deprecations.py -------------------------------------------------------------------------------- /optax/_src/factorized.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/factorized.py -------------------------------------------------------------------------------- /optax/_src/factorized_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/factorized_test.py -------------------------------------------------------------------------------- /optax/_src/float64_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/float64_test.py -------------------------------------------------------------------------------- /optax/_src/linear_algebra.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/linear_algebra.py -------------------------------------------------------------------------------- /optax/_src/linear_algebra_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/linear_algebra_test.py -------------------------------------------------------------------------------- /optax/_src/linesearch.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/linesearch.py -------------------------------------------------------------------------------- /optax/_src/linesearch_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/linesearch_test.py -------------------------------------------------------------------------------- /optax/_src/lookahead.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/lookahead.py -------------------------------------------------------------------------------- /optax/_src/lookahead_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/lookahead_test.py -------------------------------------------------------------------------------- /optax/_src/numerics.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/numerics.py -------------------------------------------------------------------------------- /optax/_src/numerics_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/numerics_test.py -------------------------------------------------------------------------------- /optax/_src/schedule.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/schedule.py -------------------------------------------------------------------------------- /optax/_src/sharding_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/sharding_test.py -------------------------------------------------------------------------------- /optax/_src/test_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/test_utils.py -------------------------------------------------------------------------------- /optax/_src/transform.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/transform.py -------------------------------------------------------------------------------- /optax/_src/transform_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/transform_test.py -------------------------------------------------------------------------------- /optax/_src/update.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/update.py -------------------------------------------------------------------------------- /optax/_src/update_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/update_test.py -------------------------------------------------------------------------------- /optax/_src/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/utils.py -------------------------------------------------------------------------------- /optax/_src/utils_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/utils_test.py -------------------------------------------------------------------------------- /optax/_src/wrappers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/_src/wrappers.py -------------------------------------------------------------------------------- /optax/assignment/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/assignment/__init__.py -------------------------------------------------------------------------------- /optax/assignment/_hungarian_algorithm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/assignment/_hungarian_algorithm.py -------------------------------------------------------------------------------- /optax/assignment/_hungarian_algorithm_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/assignment/_hungarian_algorithm_test.py -------------------------------------------------------------------------------- /optax/contrib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/__init__.py -------------------------------------------------------------------------------- /optax/contrib/_acprop.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_acprop.py -------------------------------------------------------------------------------- /optax/contrib/_ademamix.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_ademamix.py -------------------------------------------------------------------------------- /optax/contrib/_adopt.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_adopt.py -------------------------------------------------------------------------------- /optax/contrib/_cocob.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_cocob.py -------------------------------------------------------------------------------- /optax/contrib/_common_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_common_test.py -------------------------------------------------------------------------------- /optax/contrib/_complex_valued.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_complex_valued.py -------------------------------------------------------------------------------- /optax/contrib/_complex_valued_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_complex_valued_test.py -------------------------------------------------------------------------------- /optax/contrib/_dadapt_adamw.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_dadapt_adamw.py -------------------------------------------------------------------------------- /optax/contrib/_dog.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_dog.py -------------------------------------------------------------------------------- /optax/contrib/_mechanic.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_mechanic.py -------------------------------------------------------------------------------- /optax/contrib/_mechanic_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_mechanic_test.py -------------------------------------------------------------------------------- /optax/contrib/_momo.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_momo.py -------------------------------------------------------------------------------- /optax/contrib/_muon.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_muon.py -------------------------------------------------------------------------------- /optax/contrib/_muon_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_muon_test.py -------------------------------------------------------------------------------- /optax/contrib/_privacy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_privacy.py -------------------------------------------------------------------------------- /optax/contrib/_privacy_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_privacy_test.py -------------------------------------------------------------------------------- /optax/contrib/_prodigy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_prodigy.py -------------------------------------------------------------------------------- /optax/contrib/_reduce_on_plateau.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_reduce_on_plateau.py -------------------------------------------------------------------------------- /optax/contrib/_reduce_on_plateau_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_reduce_on_plateau_test.py -------------------------------------------------------------------------------- /optax/contrib/_sam.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_sam.py -------------------------------------------------------------------------------- /optax/contrib/_sam_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_sam_test.py -------------------------------------------------------------------------------- /optax/contrib/_schedule_free.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_schedule_free.py -------------------------------------------------------------------------------- /optax/contrib/_schedule_free_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_schedule_free_test.py -------------------------------------------------------------------------------- /optax/contrib/_sophia.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/contrib/_sophia.py -------------------------------------------------------------------------------- /optax/experimental/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/experimental/__init__.py -------------------------------------------------------------------------------- /optax/experimental/_aggregating.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/experimental/_aggregating.py -------------------------------------------------------------------------------- /optax/experimental/_aggregating_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/experimental/_aggregating_test.py -------------------------------------------------------------------------------- /optax/experimental/_microbatching.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/experimental/_microbatching.py -------------------------------------------------------------------------------- /optax/experimental/_microbatching_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/experimental/_microbatching_test.py -------------------------------------------------------------------------------- /optax/losses/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/__init__.py -------------------------------------------------------------------------------- /optax/losses/_classification.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_classification.py -------------------------------------------------------------------------------- /optax/losses/_classification_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_classification_test.py -------------------------------------------------------------------------------- /optax/losses/_fenchel_young.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_fenchel_young.py -------------------------------------------------------------------------------- /optax/losses/_fenchel_young_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_fenchel_young_test.py -------------------------------------------------------------------------------- /optax/losses/_ranking.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_ranking.py -------------------------------------------------------------------------------- /optax/losses/_ranking_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_ranking_test.py -------------------------------------------------------------------------------- /optax/losses/_regression.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_regression.py -------------------------------------------------------------------------------- /optax/losses/_regression_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_regression_test.py -------------------------------------------------------------------------------- /optax/losses/_segmentation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_segmentation.py -------------------------------------------------------------------------------- /optax/losses/_segmentation_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_segmentation_test.py -------------------------------------------------------------------------------- /optax/losses/_self_supervised.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_self_supervised.py -------------------------------------------------------------------------------- /optax/losses/_self_supervised_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_self_supervised_test.py -------------------------------------------------------------------------------- /optax/losses/_smoothing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_smoothing.py -------------------------------------------------------------------------------- /optax/losses/_smoothing_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/losses/_smoothing_test.py -------------------------------------------------------------------------------- /optax/optax_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/optax_test.py -------------------------------------------------------------------------------- /optax/perturbations/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/perturbations/__init__.py -------------------------------------------------------------------------------- /optax/perturbations/_make_pert.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/perturbations/_make_pert.py -------------------------------------------------------------------------------- /optax/perturbations/_make_pert_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/perturbations/_make_pert_test.py -------------------------------------------------------------------------------- /optax/projections/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/projections/__init__.py -------------------------------------------------------------------------------- /optax/projections/_projections.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/projections/_projections.py -------------------------------------------------------------------------------- /optax/projections/_projections_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/projections/_projections_test.py -------------------------------------------------------------------------------- /optax/schedules/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/schedules/__init__.py -------------------------------------------------------------------------------- /optax/schedules/_inject.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/schedules/_inject.py -------------------------------------------------------------------------------- /optax/schedules/_inject_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/schedules/_inject_test.py -------------------------------------------------------------------------------- /optax/schedules/_join.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/schedules/_join.py -------------------------------------------------------------------------------- /optax/schedules/_join_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/schedules/_join_test.py -------------------------------------------------------------------------------- /optax/schedules/_schedule.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/schedules/_schedule.py -------------------------------------------------------------------------------- /optax/schedules/_schedule_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/schedules/_schedule_test.py -------------------------------------------------------------------------------- /optax/schedules/inject.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/schedules/inject.py -------------------------------------------------------------------------------- /optax/second_order/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/second_order/__init__.py -------------------------------------------------------------------------------- /optax/second_order/_deprecated.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/second_order/_deprecated.py -------------------------------------------------------------------------------- /optax/transforms/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/__init__.py -------------------------------------------------------------------------------- /optax/transforms/_accumulation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_accumulation.py -------------------------------------------------------------------------------- /optax/transforms/_accumulation_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_accumulation_test.py -------------------------------------------------------------------------------- /optax/transforms/_adding.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_adding.py -------------------------------------------------------------------------------- /optax/transforms/_adding_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_adding_test.py -------------------------------------------------------------------------------- /optax/transforms/_clipping.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_clipping.py -------------------------------------------------------------------------------- /optax/transforms/_clipping_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_clipping_test.py -------------------------------------------------------------------------------- /optax/transforms/_combining.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_combining.py -------------------------------------------------------------------------------- /optax/transforms/_combining_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_combining_test.py -------------------------------------------------------------------------------- /optax/transforms/_conditionality.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_conditionality.py -------------------------------------------------------------------------------- /optax/transforms/_conditionality_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_conditionality_test.py -------------------------------------------------------------------------------- /optax/transforms/_constraining.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_constraining.py -------------------------------------------------------------------------------- /optax/transforms/_constraining_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_constraining_test.py -------------------------------------------------------------------------------- /optax/transforms/_freezing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_freezing.py -------------------------------------------------------------------------------- /optax/transforms/_freezing_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_freezing_test.py -------------------------------------------------------------------------------- /optax/transforms/_layouts.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_layouts.py -------------------------------------------------------------------------------- /optax/transforms/_layouts_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_layouts_test.py -------------------------------------------------------------------------------- /optax/transforms/_masking.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_masking.py -------------------------------------------------------------------------------- /optax/transforms/_masking_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_masking_test.py -------------------------------------------------------------------------------- /optax/transforms/_monitoring.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_monitoring.py -------------------------------------------------------------------------------- /optax/transforms/_monitoring_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/transforms/_monitoring_test.py -------------------------------------------------------------------------------- /optax/tree/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/tree/__init__.py -------------------------------------------------------------------------------- /optax/tree_utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/tree_utils/__init__.py -------------------------------------------------------------------------------- /optax/tree_utils/_casting.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/tree_utils/_casting.py -------------------------------------------------------------------------------- /optax/tree_utils/_casting_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/tree_utils/_casting_test.py -------------------------------------------------------------------------------- /optax/tree_utils/_random.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/tree_utils/_random.py -------------------------------------------------------------------------------- /optax/tree_utils/_random_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/tree_utils/_random_test.py -------------------------------------------------------------------------------- /optax/tree_utils/_state_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/tree_utils/_state_utils.py -------------------------------------------------------------------------------- /optax/tree_utils/_state_utils_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/tree_utils/_state_utils_test.py -------------------------------------------------------------------------------- /optax/tree_utils/_tree_math.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/tree_utils/_tree_math.py -------------------------------------------------------------------------------- /optax/tree_utils/_tree_math_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/optax/tree_utils/_tree_math_test.py -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/pyproject.toml -------------------------------------------------------------------------------- /readthedocs.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/readthedocs.yml -------------------------------------------------------------------------------- /test.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/optax/HEAD/test.sh --------------------------------------------------------------------------------