We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
__all__ is currently not useful in __init__.py
`__all__` in the `__init__.py` overrides which functions are imported when a user does `from optax import *`. It seems like every function in that file should be exposed through a wildcard import, so there is no need for the `__all__`. Besides being redundant, having it creates opportunities for bugs: right now, many of the functions (e.g. `maybe_update`, `keep_params_nonnegative`) are imported but not exposed in `__all__`. I believe it should be removed, and would be happy to create a PR if that makes sense.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `optax/__init__.py`
Content:
```
1 # Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ==============================================================================
15 """Optax: composable gradient processing and optimization, in JAX."""
16 
17 from optax._src.alias import adabelief
18 from optax._src.alias import adafactor
19 from optax._src.alias import adagrad
20 from optax._src.alias import adam
21 from optax._src.alias import adamw
22 from optax._src.alias import dpsgd
23 from optax._src.alias import fromage
24 from optax._src.alias import lamb
25 from optax._src.alias import lars
26 from optax._src.alias import noisy_sgd
27 from optax._src.alias import radam
28 from optax._src.alias import rmsprop
29 from optax._src.alias import sgd
30 from optax._src.alias import sm3
31 from optax._src.alias import yogi
32 from optax._src.base import EmptyState
33 from optax._src.base import GradientTransformation
34 from optax._src.base import identity
35 from optax._src.base import OptState
36 from optax._src.base import Params
37 from optax._src.base import Schedule
38 from optax._src.base import TransformInitFn
39 from optax._src.base import TransformUpdateFn
40 from optax._src.base import Updates
41 from optax._src.clipping import adaptive_grad_clip
42 from optax._src.clipping import AdaptiveGradClipState
43 from optax._src.clipping import clip
44 from optax._src.clipping import clip_by_block_rms
45 from optax._src.clipping import clip_by_global_norm
46 from optax._src.clipping import ClipByGlobalNormState
47 from optax._src.clipping import ClipState
48 from optax._src.combine import chain
49 from optax._src.combine import multi_transform
50 from optax._src.combine import MultiTransformState
51 from optax._src.constrain import keep_params_nonnegative
52 from optax._src.constrain import NonNegativeParamsState
53 from optax._src.constrain import zero_nans
54 from optax._src.constrain import ZeroNansState
55 from optax._src.control_variates import control_delta_method
56 from optax._src.control_variates import control_variates_jacobians
57 from optax._src.control_variates import moving_avg_baseline
58 from optax._src.factorized import FactoredState
59 from optax._src.factorized import scale_by_factored_rms
60 from optax._src.linear_algebra import global_norm
61 from optax._src.linear_algebra import matrix_inverse_pth_root
62 from optax._src.linear_algebra import power_iteration
63 from optax._src.lookahead import lookahead
64 from optax._src.lookahead import LookaheadParams
65 from optax._src.lookahead import LookaheadState
66 from optax._src.loss import cosine_distance
67 from optax._src.loss import cosine_similarity
68 from optax._src.loss import huber_loss
69 from optax._src.loss import l2_loss
70 from optax._src.loss import log_cosh
71 from optax._src.loss import sigmoid_binary_cross_entropy
72 from optax._src.loss import smooth_labels
73 from optax._src.loss import softmax_cross_entropy
74 from optax._src.privacy import differentially_private_aggregate
75 from optax._src.privacy import DifferentiallyPrivateAggregateState
76 from optax._src.schedule import constant_schedule
77 from optax._src.schedule import cosine_decay_schedule
78 from optax._src.schedule import cosine_onecycle_schedule
79 from optax._src.schedule import exponential_decay
80 from optax._src.schedule import inject_hyperparams
81 from optax._src.schedule import InjectHyperparamsState
82 from optax._src.schedule import join_schedules
83 from optax._src.schedule import linear_onecycle_schedule
84 from optax._src.schedule import linear_schedule
85 from optax._src.schedule import piecewise_constant_schedule
86 from optax._src.schedule import piecewise_interpolate_schedule
87 from optax._src.schedule import polynomial_schedule
88 from optax._src.schedule import sgdr_schedule
89 from optax._src.schedule import warmup_cosine_decay_schedule
90 from optax._src.schedule import warmup_exponential_decay_schedule
91 from optax._src.second_order import fisher_diag
92 from optax._src.second_order import hessian_diag
93 from optax._src.second_order import hvp
94 from optax._src.stochastic_gradient_estimators import measure_valued_jacobians
95 from optax._src.stochastic_gradient_estimators import pathwise_jacobians
96 from optax._src.stochastic_gradient_estimators import score_function_jacobians
97 from optax._src.transform import add_decayed_weights
98 from optax._src.transform import add_noise
99 from optax._src.transform import AddDecayedWeightsState
100 from optax._src.transform import additive_weight_decay
101 from optax._src.transform import AdditiveWeightDecayState
102 from optax._src.transform import AddNoiseState
103 from optax._src.transform import apply_every
104 from optax._src.transform import ApplyEvery
105 from optax._src.transform import centralize
106 from optax._src.transform import ema
107 from optax._src.transform import EmaState
108 from optax._src.transform import scale
109 from optax._src.transform import scale_by_adam
110 from optax._src.transform import scale_by_belief
111 from optax._src.transform import scale_by_param_block_norm
112 from optax._src.transform import scale_by_param_block_rms
113 from optax._src.transform import scale_by_radam
114 from optax._src.transform import scale_by_rms
115 from optax._src.transform import scale_by_rss
116 from optax._src.transform import scale_by_schedule
117 from optax._src.transform import scale_by_sm3
118 from optax._src.transform import scale_by_stddev
119 from optax._src.transform import scale_by_trust_ratio
120 from optax._src.transform import scale_by_yogi
121 from optax._src.transform import ScaleByAdamState
122 from optax._src.transform import ScaleByFromageState
123 from optax._src.transform import ScaleByRmsState
124 from optax._src.transform import ScaleByRssState
125 from optax._src.transform import ScaleByRStdDevState
126 from optax._src.transform import ScaleByScheduleState
127 from optax._src.transform import ScaleBySM3State
128 from optax._src.transform import ScaleByTrustRatioState
129 from optax._src.transform import ScaleState
130 from optax._src.transform import trace
131 from optax._src.transform import TraceState
132 from optax._src.update import apply_updates
133 from optax._src.update import incremental_update
134 from optax._src.update import periodic_update
135 from optax._src.utils import multi_normal
136 from optax._src.wrappers import apply_if_finite
137 from optax._src.wrappers import ApplyIfFiniteState
138 from optax._src.wrappers import flatten
139 from optax._src.wrappers import masked
140 from optax._src.wrappers import MaskedState
141 from optax._src.wrappers import maybe_update
142 from optax._src.wrappers import MaybeUpdateState
143 from optax._src.wrappers import MultiSteps
144 from optax._src.wrappers import MultiStepsState
145 
146 __version__ = "0.0.9"
147 
148 __all__ = (
149     "adabelief",
150     "adafactor",
151     "adagrad",
152     "adam",
153     "adamw",
154     "AdaptiveGradClipState",
155     "adaptive_grad_clip",
156     "add_decayed_weights",
157     "add_noise",
158     "AddDecayedWeightsState",
159     "additive_weight_decay",
160     "AdditiveWeightDecayState",
161     "AddNoiseState",
162     "apply_if_finite",
163     "apply_every",
164     "apply_updates",
165     "ApplyEvery",
166     "ApplyIfFiniteState",
167     "centralize",
168     "chain",
169     "clip",
170     "clip_by_block_rms",
171     "clip_by_global_norm",
172     "ClipByGlobalNormState",
173     "ClipState",
174     "constant_schedule",
175     "control_delta_method",
176     "control_variates_jacobians",
177     "cosine_decay_schedule",
178     "cosine_distance",
179     "cosine_onecycle_schedule",
180     "cosine_similarity",
181     "dpsgd",
182     "differentially_private_aggregate",
183     "DifferentiallyPrivateAggregateState",
184     "ema",
185     "EmaState",
186     "EmptyState",
187     "exponential_decay",
188     "FactoredState",
189     "fisher_diag",
190     "flatten",
191     "fromage",
192     "global_norm",
193     "GradientTransformation",
194     "hessian_diag",
195     "huber_loss",
196     "hvp",
197     "identity",
198     "incremental_update",
199     "inject_hyperparams",
200     "InjectHyperparamsState",
201     "join_schedules",
202     "lamb",
203     "lars",
204     "log_cosh",
205     "lookahead",
206     "LookaheadParams",
207     "LookaheadState",
208     "l2_loss",
209     "linear_onecycle_schedule",
210     "linear_schedule",
211     "matrix_inverse_pth_root",
212     "masked",
213     "MaskedState",
214     "multi_transform",
215     "MultiTransformState",
216     "measure_valued_jacobians",
217     "moving_avg_baseline",
218     "multi_normal",
219     "noisy_sgd",
220     "OptState",
221     "Params",
222     "pathwise_jacobians",
223     "periodic_update",
224     "piecewise_constant_schedule",
225     "piecewise_interpolate_schedule",
226     "polynomial_schedule",
227     "power_iteration",
228     "radam",
229     "rmsprop",
230     "scale",
231     "scale_by_adam",
232     "scale_by_belief",
233     "scale_by_param_block_norm",
234     "scale_by_param_block_rms",
235     "scale_by_radam",
236     "scale_by_rms",
237     "scale_by_rss",
238     "scale_by_schedule",
239     "scale_by_sm3",
240     "scale_by_stddev",
241     "scale_by_trust_ratio",
242     "scale_by_yogi",
243     "ScaleByAdamState",
244     "ScaleByFromageState",
245     "ScaleByRmsState",
246     "ScaleByRssState",
247     "ScaleByRStdDevState",
248     "ScaleByScheduleState",
249     "ScaleBySM3State",
250     "ScaleByTrustRatioState",
251     "ScaleState",
252     "Schedule",
253     "score_function_jacobians",
254     "sgd",
255     "sgdr_schedule",
256     "sm3",
257     "sigmoid_binary_cross_entropy",
258     "smooth_labels",
259     "softmax_cross_entropy",
260     "trace",
261     "TraceState",
262     "TransformInitFn",
263     "TransformUpdateFn",
264     "Updates",
265     "warmup_cosine_decay_schedule",
266     "warmup_exponential_decay_schedule",
267     "yogi",
268 )
269 
270 #  _________________________________________
271 # / Please don't use symbols in `_src` they \
272 # \ are not part of the Optax public API.   /
273 #  -----------------------------------------
274 #         \   ^__^
275 #          \  (oo)\_______
276 #             (__)\       )\/\
277 #                 ||----w |
278 #                 ||     ||
279 #
280
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/optax/__init__.py b/optax/__init__.py
--- a/optax/__init__.py
+++ b/optax/__init__.py
@@ -151,24 +151,24 @@
     "adagrad",
     "adam",
     "adamw",
-    "AdaptiveGradClipState",
     "adaptive_grad_clip",
+    "AdaptiveGradClipState",
     "add_decayed_weights",
     "add_noise",
     "AddDecayedWeightsState",
     "additive_weight_decay",
     "AdditiveWeightDecayState",
     "AddNoiseState",
-    "apply_if_finite",
     "apply_every",
+    "apply_if_finite",
     "apply_updates",
     "ApplyEvery",
     "ApplyIfFiniteState",
     "centralize",
     "chain",
-    "clip",
     "clip_by_block_rms",
     "clip_by_global_norm",
+    "clip",
     "ClipByGlobalNormState",
     "ClipState",
     "constant_schedule",
@@ -178,9 +178,9 @@
     "cosine_distance",
     "cosine_onecycle_schedule",
     "cosine_similarity",
-    "dpsgd",
     "differentially_private_aggregate",
     "DifferentiallyPrivateAggregateState",
+    "dpsgd",
     "ema",
     "EmaState",
     "EmptyState",
@@ -199,24 +199,30 @@
     "inject_hyperparams",
     "InjectHyperparamsState",
     "join_schedules",
+    "keep_params_nonnegative",
+    "l2_loss",
     "lamb",
     "lars",
+    "linear_onecycle_schedule",
+    "linear_schedule",
     "log_cosh",
     "lookahead",
     "LookaheadParams",
     "LookaheadState",
-    "l2_loss",
-    "linear_onecycle_schedule",
-    "linear_schedule",
-    "matrix_inverse_pth_root",
     "masked",
     "MaskedState",
-    "multi_transform",
-    "MultiTransformState",
+    "matrix_inverse_pth_root",
+    "maybe_update",
+    "MaybeUpdateState",
     "measure_valued_jacobians",
     "moving_avg_baseline",
     "multi_normal",
+    "multi_transform",
+    "MultiSteps",
+    "MultiStepsState",
+    "MultiTransformState",
     "noisy_sgd",
+    "NonNegativeParamsState",
     "OptState",
     "Params",
     "pathwise_jacobians",
@@ -227,9 +233,9 @@
     "power_iteration",
     "radam",
     "rmsprop",
-    "scale",
     "scale_by_adam",
     "scale_by_belief",
+    "scale_by_factored_rms",
     "scale_by_param_block_norm",
     "scale_by_param_block_rms",
     "scale_by_radam",
@@ -240,6 +246,7 @@
     "scale_by_stddev",
     "scale_by_trust_ratio",
     "scale_by_yogi",
+    "scale",
     "ScaleByAdamState",
     "ScaleByFromageState",
     "ScaleByRmsState",
@@ -253,8 +260,8 @@
     "score_function_jacobians",
     "sgd",
     "sgdr_schedule",
-    "sm3",
     "sigmoid_binary_cross_entropy",
+    "sm3",
     "smooth_labels",
     "softmax_cross_entropy",
     "trace",
@@ -265,6 +272,8 @@
     "warmup_cosine_decay_schedule",
     "warmup_exponential_decay_schedule",
     "yogi",
+    "zero_nans",
+    "ZeroNansState",
 )
 
 #  _________________________________________

{"golden_diff": "diff --git a/optax/__init__.py b/optax/__init__.py\n--- a/optax/__init__.py\n+++ b/optax/__init__.py\n@@ -151,24 +151,24 @@\n     \"adagrad\",\n     \"adam\",\n     \"adamw\",\n-    \"AdaptiveGradClipState\",\n     \"adaptive_grad_clip\",\n+    \"AdaptiveGradClipState\",\n     \"add_decayed_weights\",\n     \"add_noise\",\n     \"AddDecayedWeightsState\",\n     \"additive_weight_decay\",\n     \"AdditiveWeightDecayState\",\n     \"AddNoiseState\",\n-    \"apply_if_finite\",\n     \"apply_every\",\n+    \"apply_if_finite\",\n     \"apply_updates\",\n     \"ApplyEvery\",\n     \"ApplyIfFiniteState\",\n     \"centralize\",\n     \"chain\",\n-    \"clip\",\n     \"clip_by_block_rms\",\n     \"clip_by_global_norm\",\n+    \"clip\",\n     \"ClipByGlobalNormState\",\n     \"ClipState\",\n     \"constant_schedule\",\n@@ -178,9 +178,9 @@\n     \"cosine_distance\",\n     \"cosine_onecycle_schedule\",\n     \"cosine_similarity\",\n-    \"dpsgd\",\n     \"differentially_private_aggregate\",\n     \"DifferentiallyPrivateAggregateState\",\n+    \"dpsgd\",\n     \"ema\",\n     \"EmaState\",\n     \"EmptyState\",\n@@ -199,24 +199,30 @@\n     \"inject_hyperparams\",\n     \"InjectHyperparamsState\",\n     \"join_schedules\",\n+    \"keep_params_nonnegative\",\n+    \"l2_loss\",\n     \"lamb\",\n     \"lars\",\n+    \"linear_onecycle_schedule\",\n+    \"linear_schedule\",\n     \"log_cosh\",\n     \"lookahead\",\n     \"LookaheadParams\",\n     \"LookaheadState\",\n-    \"l2_loss\",\n-    \"linear_onecycle_schedule\",\n-    \"linear_schedule\",\n-    \"matrix_inverse_pth_root\",\n     \"masked\",\n     \"MaskedState\",\n-    \"multi_transform\",\n-    \"MultiTransformState\",\n+    \"matrix_inverse_pth_root\",\n+    \"maybe_update\",\n+    \"MaybeUpdateState\",\n     \"measure_valued_jacobians\",\n     \"moving_avg_baseline\",\n     \"multi_normal\",\n+    \"multi_transform\",\n+    \"MultiSteps\",\n+    \"MultiStepsState\",\n+    \"MultiTransformState\",\n     \"noisy_sgd\",\n+    \"NonNegativeParamsState\",\n     \"OptState\",\n     \"Params\",\n     \"pathwise_jacobians\",\n@@ -227,9 +233,9 @@\n     \"power_iteration\",\n     \"radam\",\n     \"rmsprop\",\n-    \"scale\",\n     \"scale_by_adam\",\n     \"scale_by_belief\",\n+    \"scale_by_factored_rms\",\n     \"scale_by_param_block_norm\",\n     \"scale_by_param_block_rms\",\n     \"scale_by_radam\",\n@@ -240,6 +246,7 @@\n     \"scale_by_stddev\",\n     \"scale_by_trust_ratio\",\n     \"scale_by_yogi\",\n+    \"scale\",\n     \"ScaleByAdamState\",\n     \"ScaleByFromageState\",\n     \"ScaleByRmsState\",\n@@ -253,8 +260,8 @@\n     \"score_function_jacobians\",\n     \"sgd\",\n     \"sgdr_schedule\",\n-    \"sm3\",\n     \"sigmoid_binary_cross_entropy\",\n+    \"sm3\",\n     \"smooth_labels\",\n     \"softmax_cross_entropy\",\n     \"trace\",\n@@ -265,6 +272,8 @@\n     \"warmup_cosine_decay_schedule\",\n     \"warmup_exponential_decay_schedule\",\n     \"yogi\",\n+    \"zero_nans\",\n+    \"ZeroNansState\",\n )\n \n #  _________________________________________\n", "issue": "__all__ is currently not useful in __init__.py\n`__all__` in the `__init__.py` overrides which functions are imported when a user does `from optax import *`. It seems like every function in that file should be exposed through a wildcard import, so there is no need for the `__all__`. Besides being redundant, having it creates opportunities for bugs: right now, many of the functions (e.g. `maybe_update`, `keep_params_nonnegative`) are imported but not exposed in `__all__`. I believe it should be removed, and would be happy to create a PR if that makes sense.\n", "before_files": [{"content": "# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Optax: composable gradient processing and optimization, in JAX.\"\"\"\n\nfrom optax._src.alias import adabelief\nfrom optax._src.alias import adafactor\nfrom optax._src.alias import adagrad\nfrom optax._src.alias import adam\nfrom optax._src.alias import adamw\nfrom optax._src.alias import dpsgd\nfrom optax._src.alias import fromage\nfrom optax._src.alias import lamb\nfrom optax._src.alias import lars\nfrom optax._src.alias import noisy_sgd\nfrom optax._src.alias import radam\nfrom optax._src.alias import rmsprop\nfrom optax._src.alias import sgd\nfrom optax._src.alias import sm3\nfrom optax._src.alias import yogi\nfrom optax._src.base import EmptyState\nfrom optax._src.base import GradientTransformation\nfrom optax._src.base import identity\nfrom optax._src.base import OptState\nfrom optax._src.base import Params\nfrom optax._src.base import Schedule\nfrom optax._src.base import TransformInitFn\nfrom optax._src.base import TransformUpdateFn\nfrom optax._src.base import Updates\nfrom optax._src.clipping import adaptive_grad_clip\nfrom optax._src.clipping import AdaptiveGradClipState\nfrom optax._src.clipping import clip\nfrom optax._src.clipping import clip_by_block_rms\nfrom optax._src.clipping import clip_by_global_norm\nfrom optax._src.clipping import ClipByGlobalNormState\nfrom optax._src.clipping import ClipState\nfrom optax._src.combine import chain\nfrom optax._src.combine import multi_transform\nfrom optax._src.combine import MultiTransformState\nfrom optax._src.constrain import keep_params_nonnegative\nfrom optax._src.constrain import NonNegativeParamsState\nfrom optax._src.constrain import zero_nans\nfrom optax._src.constrain import ZeroNansState\nfrom optax._src.control_variates import control_delta_method\nfrom optax._src.control_variates import control_variates_jacobians\nfrom optax._src.control_variates import moving_avg_baseline\nfrom optax._src.factorized import FactoredState\nfrom optax._src.factorized import scale_by_factored_rms\nfrom optax._src.linear_algebra import global_norm\nfrom optax._src.linear_algebra import matrix_inverse_pth_root\nfrom optax._src.linear_algebra import power_iteration\nfrom optax._src.lookahead import lookahead\nfrom optax._src.lookahead import LookaheadParams\nfrom optax._src.lookahead import LookaheadState\nfrom optax._src.loss import cosine_distance\nfrom optax._src.loss import cosine_similarity\nfrom optax._src.loss import huber_loss\nfrom optax._src.loss import l2_loss\nfrom optax._src.loss import log_cosh\nfrom optax._src.loss import sigmoid_binary_cross_entropy\nfrom optax._src.loss import smooth_labels\nfrom optax._src.loss import softmax_cross_entropy\nfrom optax._src.privacy import differentially_private_aggregate\nfrom optax._src.privacy import DifferentiallyPrivateAggregateState\nfrom optax._src.schedule import constant_schedule\nfrom optax._src.schedule import cosine_decay_schedule\nfrom optax._src.schedule import cosine_onecycle_schedule\nfrom optax._src.schedule import exponential_decay\nfrom optax._src.schedule import inject_hyperparams\nfrom optax._src.schedule import InjectHyperparamsState\nfrom optax._src.schedule import join_schedules\nfrom optax._src.schedule import linear_onecycle_schedule\nfrom optax._src.schedule import linear_schedule\nfrom optax._src.schedule import piecewise_constant_schedule\nfrom optax._src.schedule import piecewise_interpolate_schedule\nfrom optax._src.schedule import polynomial_schedule\nfrom optax._src.schedule import sgdr_schedule\nfrom optax._src.schedule import warmup_cosine_decay_schedule\nfrom optax._src.schedule import warmup_exponential_decay_schedule\nfrom optax._src.second_order import fisher_diag\nfrom optax._src.second_order import hessian_diag\nfrom optax._src.second_order import hvp\nfrom optax._src.stochastic_gradient_estimators import measure_valued_jacobians\nfrom optax._src.stochastic_gradient_estimators import pathwise_jacobians\nfrom optax._src.stochastic_gradient_estimators import score_function_jacobians\nfrom optax._src.transform import add_decayed_weights\nfrom optax._src.transform import add_noise\nfrom optax._src.transform import AddDecayedWeightsState\nfrom optax._src.transform import additive_weight_decay\nfrom optax._src.transform import AdditiveWeightDecayState\nfrom optax._src.transform import AddNoiseState\nfrom optax._src.transform import apply_every\nfrom optax._src.transform import ApplyEvery\nfrom optax._src.transform import centralize\nfrom optax._src.transform import ema\nfrom optax._src.transform import EmaState\nfrom optax._src.transform import scale\nfrom optax._src.transform import scale_by_adam\nfrom optax._src.transform import scale_by_belief\nfrom optax._src.transform import scale_by_param_block_norm\nfrom optax._src.transform import scale_by_param_block_rms\nfrom optax._src.transform import scale_by_radam\nfrom optax._src.transform import scale_by_rms\nfrom optax._src.transform import scale_by_rss\nfrom optax._src.transform import scale_by_schedule\nfrom optax._src.transform import scale_by_sm3\nfrom optax._src.transform import scale_by_stddev\nfrom optax._src.transform import scale_by_trust_ratio\nfrom optax._src.transform import scale_by_yogi\nfrom optax._src.transform import ScaleByAdamState\nfrom optax._src.transform import ScaleByFromageState\nfrom optax._src.transform import ScaleByRmsState\nfrom optax._src.transform import ScaleByRssState\nfrom optax._src.transform import ScaleByRStdDevState\nfrom optax._src.transform import ScaleByScheduleState\nfrom optax._src.transform import ScaleBySM3State\nfrom optax._src.transform import ScaleByTrustRatioState\nfrom optax._src.transform import ScaleState\nfrom optax._src.transform import trace\nfrom optax._src.transform import TraceState\nfrom optax._src.update import apply_updates\nfrom optax._src.update import incremental_update\nfrom optax._src.update import periodic_update\nfrom optax._src.utils import multi_normal\nfrom optax._src.wrappers import apply_if_finite\nfrom optax._src.wrappers import ApplyIfFiniteState\nfrom optax._src.wrappers import flatten\nfrom optax._src.wrappers import masked\nfrom optax._src.wrappers import MaskedState\nfrom optax._src.wrappers import maybe_update\nfrom optax._src.wrappers import MaybeUpdateState\nfrom optax._src.wrappers import MultiSteps\nfrom optax._src.wrappers import MultiStepsState\n\n__version__ = \"0.0.9\"\n\n__all__ = (\n    \"adabelief\",\n    \"adafactor\",\n    \"adagrad\",\n    \"adam\",\n    \"adamw\",\n    \"AdaptiveGradClipState\",\n    \"adaptive_grad_clip\",\n    \"add_decayed_weights\",\n    \"add_noise\",\n    \"AddDecayedWeightsState\",\n    \"additive_weight_decay\",\n    \"AdditiveWeightDecayState\",\n    \"AddNoiseState\",\n    \"apply_if_finite\",\n    \"apply_every\",\n    \"apply_updates\",\n    \"ApplyEvery\",\n    \"ApplyIfFiniteState\",\n    \"centralize\",\n    \"chain\",\n    \"clip\",\n    \"clip_by_block_rms\",\n    \"clip_by_global_norm\",\n    \"ClipByGlobalNormState\",\n    \"ClipState\",\n    \"constant_schedule\",\n    \"control_delta_method\",\n    \"control_variates_jacobians\",\n    \"cosine_decay_schedule\",\n    \"cosine_distance\",\n    \"cosine_onecycle_schedule\",\n    \"cosine_similarity\",\n    \"dpsgd\",\n    \"differentially_private_aggregate\",\n    \"DifferentiallyPrivateAggregateState\",\n    \"ema\",\n    \"EmaState\",\n    \"EmptyState\",\n    \"exponential_decay\",\n    \"FactoredState\",\n    \"fisher_diag\",\n    \"flatten\",\n    \"fromage\",\n    \"global_norm\",\n    \"GradientTransformation\",\n    \"hessian_diag\",\n    \"huber_loss\",\n    \"hvp\",\n    \"identity\",\n    \"incremental_update\",\n    \"inject_hyperparams\",\n    \"InjectHyperparamsState\",\n    \"join_schedules\",\n    \"lamb\",\n    \"lars\",\n    \"log_cosh\",\n    \"lookahead\",\n    \"LookaheadParams\",\n    \"LookaheadState\",\n    \"l2_loss\",\n    \"linear_onecycle_schedule\",\n    \"linear_schedule\",\n    \"matrix_inverse_pth_root\",\n    \"masked\",\n    \"MaskedState\",\n    \"multi_transform\",\n    \"MultiTransformState\",\n    \"measure_valued_jacobians\",\n    \"moving_avg_baseline\",\n    \"multi_normal\",\n    \"noisy_sgd\",\n    \"OptState\",\n    \"Params\",\n    \"pathwise_jacobians\",\n    \"periodic_update\",\n    \"piecewise_constant_schedule\",\n    \"piecewise_interpolate_schedule\",\n    \"polynomial_schedule\",\n    \"power_iteration\",\n    \"radam\",\n    \"rmsprop\",\n    \"scale\",\n    \"scale_by_adam\",\n    \"scale_by_belief\",\n    \"scale_by_param_block_norm\",\n    \"scale_by_param_block_rms\",\n    \"scale_by_radam\",\n    \"scale_by_rms\",\n    \"scale_by_rss\",\n    \"scale_by_schedule\",\n    \"scale_by_sm3\",\n    \"scale_by_stddev\",\n    \"scale_by_trust_ratio\",\n    \"scale_by_yogi\",\n    \"ScaleByAdamState\",\n    \"ScaleByFromageState\",\n    \"ScaleByRmsState\",\n    \"ScaleByRssState\",\n    \"ScaleByRStdDevState\",\n    \"ScaleByScheduleState\",\n    \"ScaleBySM3State\",\n    \"ScaleByTrustRatioState\",\n    \"ScaleState\",\n    \"Schedule\",\n    \"score_function_jacobians\",\n    \"sgd\",\n    \"sgdr_schedule\",\n    \"sm3\",\n    \"sigmoid_binary_cross_entropy\",\n    \"smooth_labels\",\n    \"softmax_cross_entropy\",\n    \"trace\",\n    \"TraceState\",\n    \"TransformInitFn\",\n    \"TransformUpdateFn\",\n    \"Updates\",\n    \"warmup_cosine_decay_schedule\",\n    \"warmup_exponential_decay_schedule\",\n    \"yogi\",\n)\n\n#  _________________________________________\n# / Please don't use symbols in `_src` they \\\n# \\ are not part of the Optax public API.   /\n#  -----------------------------------------\n#         \\   ^__^\n#          \\  (oo)\\_______\n#             (__)\\       )\\/\\\n#                 ||----w |\n#                 ||     ||\n#\n", "path": "optax/__init__.py"}], "after_files": [{"content": "# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Optax: composable gradient processing and optimization, in JAX.\"\"\"\n\nfrom optax._src.alias import adabelief\nfrom optax._src.alias import adafactor\nfrom optax._src.alias import adagrad\nfrom optax._src.alias import adam\nfrom optax._src.alias import adamw\nfrom optax._src.alias import dpsgd\nfrom optax._src.alias import fromage\nfrom optax._src.alias import lamb\nfrom optax._src.alias import lars\nfrom optax._src.alias import noisy_sgd\nfrom optax._src.alias import radam\nfrom optax._src.alias import rmsprop\nfrom optax._src.alias import sgd\nfrom optax._src.alias import sm3\nfrom optax._src.alias import yogi\nfrom optax._src.base import EmptyState\nfrom optax._src.base import GradientTransformation\nfrom optax._src.base import identity\nfrom optax._src.base import OptState\nfrom optax._src.base import Params\nfrom optax._src.base import Schedule\nfrom optax._src.base import TransformInitFn\nfrom optax._src.base import TransformUpdateFn\nfrom optax._src.base import Updates\nfrom optax._src.clipping import adaptive_grad_clip\nfrom optax._src.clipping import AdaptiveGradClipState\nfrom optax._src.clipping import clip\nfrom optax._src.clipping import clip_by_block_rms\nfrom optax._src.clipping import clip_by_global_norm\nfrom optax._src.clipping import ClipByGlobalNormState\nfrom optax._src.clipping import ClipState\nfrom optax._src.combine import chain\nfrom optax._src.combine import multi_transform\nfrom optax._src.combine import MultiTransformState\nfrom optax._src.constrain import keep_params_nonnegative\nfrom optax._src.constrain import NonNegativeParamsState\nfrom optax._src.constrain import zero_nans\nfrom optax._src.constrain import ZeroNansState\nfrom optax._src.control_variates import control_delta_method\nfrom optax._src.control_variates import control_variates_jacobians\nfrom optax._src.control_variates import moving_avg_baseline\nfrom optax._src.factorized import FactoredState\nfrom optax._src.factorized import scale_by_factored_rms\nfrom optax._src.linear_algebra import global_norm\nfrom optax._src.linear_algebra import matrix_inverse_pth_root\nfrom optax._src.linear_algebra import power_iteration\nfrom optax._src.lookahead import lookahead\nfrom optax._src.lookahead import LookaheadParams\nfrom optax._src.lookahead import LookaheadState\nfrom optax._src.loss import cosine_distance\nfrom optax._src.loss import cosine_similarity\nfrom optax._src.loss import huber_loss\nfrom optax._src.loss import l2_loss\nfrom optax._src.loss import log_cosh\nfrom optax._src.loss import sigmoid_binary_cross_entropy\nfrom optax._src.loss import smooth_labels\nfrom optax._src.loss import softmax_cross_entropy\nfrom optax._src.privacy import differentially_private_aggregate\nfrom optax._src.privacy import DifferentiallyPrivateAggregateState\nfrom optax._src.schedule import constant_schedule\nfrom optax._src.schedule import cosine_decay_schedule\nfrom optax._src.schedule import cosine_onecycle_schedule\nfrom optax._src.schedule import exponential_decay\nfrom optax._src.schedule import inject_hyperparams\nfrom optax._src.schedule import InjectHyperparamsState\nfrom optax._src.schedule import join_schedules\nfrom optax._src.schedule import linear_onecycle_schedule\nfrom optax._src.schedule import linear_schedule\nfrom optax._src.schedule import piecewise_constant_schedule\nfrom optax._src.schedule import piecewise_interpolate_schedule\nfrom optax._src.schedule import polynomial_schedule\nfrom optax._src.schedule import sgdr_schedule\nfrom optax._src.schedule import warmup_cosine_decay_schedule\nfrom optax._src.schedule import warmup_exponential_decay_schedule\nfrom optax._src.second_order import fisher_diag\nfrom optax._src.second_order import hessian_diag\nfrom optax._src.second_order import hvp\nfrom optax._src.stochastic_gradient_estimators import measure_valued_jacobians\nfrom optax._src.stochastic_gradient_estimators import pathwise_jacobians\nfrom optax._src.stochastic_gradient_estimators import score_function_jacobians\nfrom optax._src.transform import add_decayed_weights\nfrom optax._src.transform import add_noise\nfrom optax._src.transform import AddDecayedWeightsState\nfrom optax._src.transform import additive_weight_decay\nfrom optax._src.transform import AdditiveWeightDecayState\nfrom optax._src.transform import AddNoiseState\nfrom optax._src.transform import apply_every\nfrom optax._src.transform import ApplyEvery\nfrom optax._src.transform import centralize\nfrom optax._src.transform import ema\nfrom optax._src.transform import EmaState\nfrom optax._src.transform import scale\nfrom optax._src.transform import scale_by_adam\nfrom optax._src.transform import scale_by_belief\nfrom optax._src.transform import scale_by_param_block_norm\nfrom optax._src.transform import scale_by_param_block_rms\nfrom optax._src.transform import scale_by_radam\nfrom optax._src.transform import scale_by_rms\nfrom optax._src.transform import scale_by_rss\nfrom optax._src.transform import scale_by_schedule\nfrom optax._src.transform import scale_by_sm3\nfrom optax._src.transform import scale_by_stddev\nfrom optax._src.transform import scale_by_trust_ratio\nfrom optax._src.transform import scale_by_yogi\nfrom optax._src.transform import ScaleByAdamState\nfrom optax._src.transform import ScaleByFromageState\nfrom optax._src.transform import ScaleByRmsState\nfrom optax._src.transform import ScaleByRssState\nfrom optax._src.transform import ScaleByRStdDevState\nfrom optax._src.transform import ScaleByScheduleState\nfrom optax._src.transform import ScaleBySM3State\nfrom optax._src.transform import ScaleByTrustRatioState\nfrom optax._src.transform import ScaleState\nfrom optax._src.transform import trace\nfrom optax._src.transform import TraceState\nfrom optax._src.update import apply_updates\nfrom optax._src.update import incremental_update\nfrom optax._src.update import periodic_update\nfrom optax._src.utils import multi_normal\nfrom optax._src.wrappers import apply_if_finite\nfrom optax._src.wrappers import ApplyIfFiniteState\nfrom optax._src.wrappers import flatten\nfrom optax._src.wrappers import masked\nfrom optax._src.wrappers import MaskedState\nfrom optax._src.wrappers import maybe_update\nfrom optax._src.wrappers import MaybeUpdateState\nfrom optax._src.wrappers import MultiSteps\nfrom optax._src.wrappers import MultiStepsState\n\n__version__ = \"0.0.9\"\n\n__all__ = (\n    \"adabelief\",\n    \"adafactor\",\n    \"adagrad\",\n    \"adam\",\n    \"adamw\",\n    \"adaptive_grad_clip\",\n    \"AdaptiveGradClipState\",\n    \"add_decayed_weights\",\n    \"add_noise\",\n    \"AddDecayedWeightsState\",\n    \"additive_weight_decay\",\n    \"AdditiveWeightDecayState\",\n    \"AddNoiseState\",\n    \"apply_every\",\n    \"apply_if_finite\",\n    \"apply_updates\",\n    \"ApplyEvery\",\n    \"ApplyIfFiniteState\",\n    \"centralize\",\n    \"chain\",\n    \"clip_by_block_rms\",\n    \"clip_by_global_norm\",\n    \"clip\",\n    \"ClipByGlobalNormState\",\n    \"ClipState\",\n    \"constant_schedule\",\n    \"control_delta_method\",\n    \"control_variates_jacobians\",\n    \"cosine_decay_schedule\",\n    \"cosine_distance\",\n    \"cosine_onecycle_schedule\",\n    \"cosine_similarity\",\n    \"differentially_private_aggregate\",\n    \"DifferentiallyPrivateAggregateState\",\n    \"dpsgd\",\n    \"ema\",\n    \"EmaState\",\n    \"EmptyState\",\n    \"exponential_decay\",\n    \"FactoredState\",\n    \"fisher_diag\",\n    \"flatten\",\n    \"fromage\",\n    \"global_norm\",\n    \"GradientTransformation\",\n    \"hessian_diag\",\n    \"huber_loss\",\n    \"hvp\",\n    \"identity\",\n    \"incremental_update\",\n    \"inject_hyperparams\",\n    \"InjectHyperparamsState\",\n    \"join_schedules\",\n    \"keep_params_nonnegative\",\n    \"l2_loss\",\n    \"lamb\",\n    \"lars\",\n    \"linear_onecycle_schedule\",\n    \"linear_schedule\",\n    \"log_cosh\",\n    \"lookahead\",\n    \"LookaheadParams\",\n    \"LookaheadState\",\n    \"masked\",\n    \"MaskedState\",\n    \"matrix_inverse_pth_root\",\n    \"maybe_update\",\n    \"MaybeUpdateState\",\n    \"measure_valued_jacobians\",\n    \"moving_avg_baseline\",\n    \"multi_normal\",\n    \"multi_transform\",\n    \"MultiSteps\",\n    \"MultiStepsState\",\n    \"MultiTransformState\",\n    \"noisy_sgd\",\n    \"NonNegativeParamsState\",\n    \"OptState\",\n    \"Params\",\n    \"pathwise_jacobians\",\n    \"periodic_update\",\n    \"piecewise_constant_schedule\",\n    \"piecewise_interpolate_schedule\",\n    \"polynomial_schedule\",\n    \"power_iteration\",\n    \"radam\",\n    \"rmsprop\",\n    \"scale_by_adam\",\n    \"scale_by_belief\",\n    \"scale_by_factored_rms\",\n    \"scale_by_param_block_norm\",\n    \"scale_by_param_block_rms\",\n    \"scale_by_radam\",\n    \"scale_by_rms\",\n    \"scale_by_rss\",\n    \"scale_by_schedule\",\n    \"scale_by_sm3\",\n    \"scale_by_stddev\",\n    \"scale_by_trust_ratio\",\n    \"scale_by_yogi\",\n    \"scale\",\n    \"ScaleByAdamState\",\n    \"ScaleByFromageState\",\n    \"ScaleByRmsState\",\n    \"ScaleByRssState\",\n    \"ScaleByRStdDevState\",\n    \"ScaleByScheduleState\",\n    \"ScaleBySM3State\",\n    \"ScaleByTrustRatioState\",\n    \"ScaleState\",\n    \"Schedule\",\n    \"score_function_jacobians\",\n    \"sgd\",\n    \"sgdr_schedule\",\n    \"sigmoid_binary_cross_entropy\",\n    \"sm3\",\n    \"smooth_labels\",\n    \"softmax_cross_entropy\",\n    \"trace\",\n    \"TraceState\",\n    \"TransformInitFn\",\n    \"TransformUpdateFn\",\n    \"Updates\",\n    \"warmup_cosine_decay_schedule\",\n    \"warmup_exponential_decay_schedule\",\n    \"yogi\",\n    \"zero_nans\",\n    \"ZeroNansState\",\n)\n\n#  _________________________________________\n# / Please don't use symbols in `_src` they \\\n# \\ are not part of the Optax public API.   /\n#  -----------------------------------------\n#         \\   ^__^\n#          \\  (oo)\\_______\n#             (__)\\       )\\/\\\n#                 ||----w |\n#                 ||     ||\n#\n", "path": "optax/__init__.py"}]}