Repository: google-deepmind/optax

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.


Discussion:
`__all__` can be a useful way to communicate explicitly the intended public API, and is considered good practice according to [PEP8](https://www.python.org/dev/peps/pep-0008/#public-and-internal-interfaces).

This can help to avoid, for example, leaking global constants declared in `__init__.py` into the public interface.

You raise an interesting point that not all functions have been added to the `__all__`, though. I will look a little further into where this may be intention or not. 
I agree with the motivation in PEP8, and think `__all__`s are useful in other libraries, but I think the situation is different here. Optax already keeps all the source in the private `_src` folder and imports necessary symbols explicitly in `__init__`, which already communicates the intended API and prevents leaking constants, thus makes the `__all__` redundant. 

Resolving patch:
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",
 )
 
 #  _________________________________________


Test patch:
diff --git a/optax/_src/control_variates_test.py b/optax/_src/control_variates_test.py
--- a/optax/_src/control_variates_test.py
+++ b/optax/_src/control_variates_test.py
@@ -45,8 +45,7 @@ def _assert_equal(actual, expected, rtol=1e-2, atol=1e-2):
   # We get around the bug https://github.com/numpy/numpy/issues/13801
   zero_indices = np.argwhere(expected == 0)
   if not np.all(np.abs(actual[zero_indices]) <= atol):
-    raise AssertionError('Larger than {} diff in {}'.format(
-        atol, actual[zero_indices]))
+    raise AssertionError(f'Larger than {atol} diff in {actual[zero_indices]}')
 
   non_zero_indices = np.argwhere(expected != 0)
   np.testing.assert_allclose(
diff --git a/optax/_src/stochastic_gradient_estimators_test.py b/optax/_src/stochastic_gradient_estimators_test.py
--- a/optax/_src/stochastic_gradient_estimators_test.py
+++ b/optax/_src/stochastic_gradient_estimators_test.py
@@ -55,8 +55,7 @@ def _assert_equal(actual, expected, rtol=1e-2, atol=1e-2):
   # We get around the bug https://github.com/numpy/numpy/issues/13801
   zero_indices = np.argwhere(expected == 0)
   if not np.all(np.abs(actual[zero_indices]) <= atol):
-    raise AssertionError('Larger than {} diff in {}'.format(
-        atol, actual[zero_indices]))
+    raise AssertionError(f'Larger than {atol} diff in {actual[zero_indices]}')
 
   non_zero_indices = np.argwhere(expected != 0)
   np.testing.assert_allclose(