optax 0.2.2__tar.gz → 0.2.4__tar.gz

{optax-0.2.2 → optax-0.2.4}/PKG-INFO

@@ -1,7 +1,7 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: optax
-Version: 0.2.2
-Summary: A gradient processing and optimisation library in JAX.
+Version: 0.2.4
+Summary: A gradient processing and optimization library in JAX.
 Keywords: python,machine learning,reinforcement-learning
 Author-email: Google DeepMind <optax-dev@google.com>
 Requires-Python: >=3.9
@@ -17,10 +17,11 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Dist: absl-py>=0.7.1
-Requires-Dist: chex>=0.1.86
-Requires-Dist: jax>=0.1.55
-Requires-Dist: jaxlib>=0.1.37
+Requires-Dist: chex>=0.1.87
+Requires-Dist: jax>=0.4.27
+Requires-Dist: jaxlib>=0.4.27
 Requires-Dist: numpy>=1.18.0
+Requires-Dist: etils[epy]
 Requires-Dist: sphinx>=6.0.0 ; extra == "docs"
 Requires-Dist: sphinx-book-theme>=1.0.1 ; extra == "docs"
 Requires-Dist: sphinxcontrib-katex ; extra == "docs"
@@ -42,9 +43,12 @@ Requires-Dist: scipy>=1.7.1 ; extra == "dp-accounting"
 Requires-Dist: tensorflow-datasets>=4.2.0 ; extra == "examples"
 Requires-Dist: tensorflow>=2.4.0 ; extra == "examples"
 Requires-Dist: dp_accounting>=0.4 ; extra == "examples"
+Requires-Dist: ipywidgets ; extra == "examples"
 Requires-Dist: flax ; extra == "examples"
 Requires-Dist: dm-tree>=0.1.7 ; extra == "test"
 Requires-Dist: flax>=0.5.3 ; extra == "test"
+Requires-Dist: scipy>=1.7.1 ; extra == "test"
+Requires-Dist: scikit-learn ; extra == "test"
 Project-URL: documentation, https://optax.readthedocs.io/
 Project-URL: homepage, https://github.com/google-deepmind/optax
 Project-URL: repository, https://github.com/google-deepmind/optax
@@ -70,12 +74,12 @@ Our goals are to
 
 *   Provide simple, well-tested, efficient implementations of core components.
 *   Improve research productivity by enabling to easily combine low-level
-    ingredients into custom optimisers (or other gradient processing components).
+    ingredients into custom optimizers (or other gradient processing components).
 *   Accelerate adoption of new ideas by making it easy for anyone to contribute.
 
-We favour focusing on small composable building blocks that can be effectively
+We favor focusing on small composable building blocks that can be effectively
 combined into custom solutions. Others may build upon these basic components
-in more complicated abstractions. Whenever reasonable, implementations prioritise
+in more complicated abstractions. Whenever reasonable, implementations prioritize
 readability and structuring code to match standard equations, over code reuse.
 
 An initial prototype of this library was made available in JAX's experimental
@@ -132,7 +136,7 @@ updates, opt_state = optimizer.update(grads, opt_state)
 params = optax.apply_updates(params, updates)
 ```
 
-You can continue the quick start in [the Optax 101 notebook.](https://github.com/google-deepmind/optax/blob/main/docs/optax-101.ipynb)
+You can continue the quick start in [the Optax 🚀 Getting started notebook.](https://github.com/google-deepmind/optax/blob/main/docs/getting_started.ipynb)
 
 ## Development
 
@@ -150,7 +154,7 @@ git clone https://github.com/google-deepmind/optax.git
 To run the tests, please execute the following script.
 
 ```sh
-sh ./test.sh
+sh test.sh
 ```
 
 ### Documentation
@@ -161,7 +165,7 @@ pip install -e ".[docs]"
 ```
 Then, execute the following.
 ```sh
-cd docs/
+cd docs
 make html
 ```
 

{optax-0.2.2 → optax-0.2.4}/README.md

@@ -15,12 +15,12 @@ Our goals are to
 
 *   Provide simple, well-tested, efficient implementations of core components.
 *   Improve research productivity by enabling to easily combine low-level
-    ingredients into custom optimisers (or other gradient processing components).
+    ingredients into custom optimizers (or other gradient processing components).
 *   Accelerate adoption of new ideas by making it easy for anyone to contribute.
 
-We favour focusing on small composable building blocks that can be effectively
+We favor focusing on small composable building blocks that can be effectively
 combined into custom solutions. Others may build upon these basic components
-in more complicated abstractions. Whenever reasonable, implementations prioritise
+in more complicated abstractions. Whenever reasonable, implementations prioritize
 readability and structuring code to match standard equations, over code reuse.
 
 An initial prototype of this library was made available in JAX's experimental
@@ -77,7 +77,7 @@ updates, opt_state = optimizer.update(grads, opt_state)
 params = optax.apply_updates(params, updates)
 ```
 
-You can continue the quick start in [the Optax 101 notebook.](https://github.com/google-deepmind/optax/blob/main/docs/optax-101.ipynb)
+You can continue the quick start in [the Optax 🚀 Getting started notebook.](https://github.com/google-deepmind/optax/blob/main/docs/getting_started.ipynb)
 
 ## Development
 
@@ -95,7 +95,7 @@ git clone https://github.com/google-deepmind/optax.git
 To run the tests, please execute the following script.
 
 ```sh
-sh ./test.sh
+sh test.sh
 ```
 
 ### Documentation
@@ -106,7 +106,7 @@ pip install -e ".[docs]"
 ```
 Then, execute the following.
 ```sh
-cd docs/
+cd docs
 make html
 ```
 

{optax-0.2.2 → optax-0.2.4}/optax/__init__.py

@@ -14,12 +14,18 @@
 # ==============================================================================
 """Optax: composable gradient processing and optimization, in JAX."""
 
+# pylint: disable=wrong-import-position
+# pylint: disable=g-importing-member
+
+from optax import assignment
 from optax import contrib
 from optax import losses
 from optax import monte_carlo
+from optax import perturbations
 from optax import projections
 from optax import schedules
 from optax import second_order
+from optax import transforms
 from optax import tree_utils
 from optax._src.alias import adabelief
 from optax._src.alias import adadelta
@@ -29,22 +35,26 @@ from optax._src.alias import adam
 from optax._src.alias import adamax
 from optax._src.alias import adamaxw
 from optax._src.alias import adamw
+from optax._src.alias import adan
 from optax._src.alias import amsgrad
 from optax._src.alias import fromage
 from optax._src.alias import lamb
 from optax._src.alias import lars
+from optax._src.alias import lbfgs
 from optax._src.alias import lion
 from optax._src.alias import MaskOrFn
 from optax._src.alias import nadam
 from optax._src.alias import nadamw
 from optax._src.alias import noisy_sgd
 from optax._src.alias import novograd
+from optax._src.alias import optimistic_adam
 from optax._src.alias import optimistic_gradient_descent
 from optax._src.alias import polyak_sgd
 from optax._src.alias import radam
 from optax._src.alias import rmsprop
 from optax._src.alias import rprop
 from optax._src.alias import sgd
+from optax._src.alias import sign_sgd
 from optax._src.alias import sm3
 from optax._src.alias import yogi
 from optax._src.base import EmptyState
@@ -86,10 +96,14 @@ from optax._src.linear_algebra import global_norm
 from optax._src.linear_algebra import matrix_inverse_pth_root
 from optax._src.linear_algebra import power_iteration
 from optax._src.linesearch import scale_by_backtracking_linesearch
+from optax._src.linesearch import scale_by_zoom_linesearch
 from optax._src.linesearch import ScaleByBacktrackingLinesearchState
+from optax._src.linesearch import ScaleByZoomLinesearchState
+from optax._src.linesearch import ZoomLinesearchInfo
 from optax._src.lookahead import lookahead
 from optax._src.lookahead import LookaheadParams
 from optax._src.lookahead import LookaheadState
+from optax._src.numerics import safe_increment
 from optax._src.numerics import safe_int32_increment
 from optax._src.numerics import safe_norm
 from optax._src.numerics import safe_root_mean_squares
@@ -99,17 +113,19 @@ from optax._src.transform import AddDecayedWeightsState
 from optax._src.transform import AddNoiseState
 from optax._src.transform import apply_every
 from optax._src.transform import ApplyEvery
-from optax._src.transform import bias_correction
 from optax._src.transform import centralize
 from optax._src.transform import ema
 from optax._src.transform import EmaState
+from optax._src.transform import normalize_by_update_norm
 from optax._src.transform import scale
 from optax._src.transform import scale_by_adadelta
 from optax._src.transform import scale_by_adam
 from optax._src.transform import scale_by_adamax
+from optax._src.transform import scale_by_adan
 from optax._src.transform import scale_by_amsgrad
 from optax._src.transform import scale_by_belief
 from optax._src.transform import scale_by_distance_over_gradients
+from optax._src.transform import scale_by_lbfgs
 from optax._src.transform import scale_by_learning_rate
 from optax._src.transform import scale_by_lion
 from optax._src.transform import scale_by_novograd
@@ -122,14 +138,17 @@ from optax._src.transform import scale_by_rms
 from optax._src.transform import scale_by_rprop
 from optax._src.transform import scale_by_rss
 from optax._src.transform import scale_by_schedule
+from optax._src.transform import scale_by_sign
 from optax._src.transform import scale_by_sm3
 from optax._src.transform import scale_by_stddev
 from optax._src.transform import scale_by_trust_ratio
 from optax._src.transform import scale_by_yogi
 from optax._src.transform import ScaleByAdaDeltaState
 from optax._src.transform import ScaleByAdamState
+from optax._src.transform import ScaleByAdanState
 from optax._src.transform import ScaleByAmsgradState
 from optax._src.transform import ScaleByBeliefState
+from optax._src.transform import ScaleByLBFGSState
 from optax._src.transform import ScaleByLionState
 from optax._src.transform import ScaleByNovogradState
 from optax._src.transform import ScaleByRmsState
@@ -142,9 +161,6 @@ from optax._src.transform import ScaleByTrustRatioState
 from optax._src.transform import ScaleState
 from optax._src.transform import trace
 from optax._src.transform import TraceState
-from optax._src.transform import update_infinity_moment
-from optax._src.transform import update_moment
-from optax._src.transform import update_moment_per_elem_norm
 from optax._src.update import apply_updates
 from optax._src.update import incremental_update
 from optax._src.update import periodic_update
@@ -153,6 +169,10 @@ from optax._src.utils import scale_gradient
 from optax._src.utils import value_and_grad_from_state
 from optax._src.wrappers import apply_if_finite
 from optax._src.wrappers import ApplyIfFiniteState
+from optax._src.wrappers import conditionally_mask
+from optax._src.wrappers import conditionally_transform
+from optax._src.wrappers import ConditionallyMaskState
+from optax._src.wrappers import ConditionallyTransformState
 from optax._src.wrappers import flatten
 from optax._src.wrappers import masked
 from optax._src.wrappers import MaskedNode
@@ -165,8 +185,13 @@ from optax._src.wrappers import ShouldSkipUpdateFunction
 from optax._src.wrappers import skip_large_updates
 from optax._src.wrappers import skip_not_finite
 
+
 # TODO(mtthss): remove tree_utils aliases after updates.
 tree_map_params = tree_utils.tree_map_params
+bias_correction = tree_utils.tree_bias_correction
+update_infinity_moment = tree_utils.tree_update_infinity_moment
+update_moment = tree_utils.tree_update_moment
+update_moment_per_elem_norm = tree_utils.tree_update_moment_per_elem_norm
 
 # TODO(mtthss): remove schedules alises from flat namespaces after user updates.
 constant_schedule = schedules.constant_schedule
@@ -182,6 +207,7 @@ piecewise_constant_schedule = schedules.piecewise_constant_schedule
 piecewise_interpolate_schedule = schedules.piecewise_interpolate_schedule
 polynomial_schedule = schedules.polynomial_schedule
 sgdr_schedule = schedules.sgdr_schedule
+warmup_constant_schedule = schedules.warmup_constant_schedule
 warmup_cosine_decay_schedule = schedules.warmup_cosine_decay_schedule
 warmup_exponential_decay_schedule = schedules.warmup_exponential_decay_schedule
 inject_stateful_hyperparams = schedules.inject_stateful_hyperparams
@@ -199,8 +225,10 @@ huber_loss = losses.huber_loss
 kl_divergence = losses.kl_divergence
 l2_loss = losses.l2_loss
 log_cosh = losses.log_cosh
+ntxent = losses.ntxent
 sigmoid_binary_cross_entropy = losses.sigmoid_binary_cross_entropy
 smooth_labels = losses.smooth_labels
+safe_softmax_cross_entropy = losses.safe_softmax_cross_entropy
 softmax_cross_entropy = losses.softmax_cross_entropy
 softmax_cross_entropy_with_integer_labels = (
     losses.softmax_cross_entropy_with_integer_labels
@@ -208,14 +236,61 @@ softmax_cross_entropy_with_integer_labels = (
 squared_error = losses.squared_error
 sigmoid_focal_loss = losses.sigmoid_focal_loss
 
+# pylint: disable=g-import-not-at-top
 # TODO(mtthss): remove contrib aliases from flat namespace once users updated.
-differentially_private_aggregate = contrib.differentially_private_aggregate
-DifferentiallyPrivateAggregateState = (
-    contrib.DifferentiallyPrivateAggregateState
-)
-dpsgd = contrib.dpsgd
+# Deprecated modules
+from optax.contrib import differentially_private_aggregate as _deprecated_differentially_private_aggregate
+from optax.contrib import DifferentiallyPrivateAggregateState as _deprecated_DifferentiallyPrivateAggregateState
+from optax.contrib import dpsgd as _deprecated_dpsgd
+
+_deprecations = {
+    # Added Apr 2024
+    "differentially_private_aggregate": (
+        (
+            "optax.differentially_private_aggregate is deprecated: use"
+            " optax.contrib.differentially_private_aggregate (optax v0.1.8 or"
+            " newer)."
+        ),
+        _deprecated_differentially_private_aggregate,
+    ),
+    "DifferentiallyPrivateAggregateState": (
+        (
+            "optax.DifferentiallyPrivateAggregateState is deprecated: use"
+            " optax.contrib.DifferentiallyPrivateAggregateState (optax v0.1.8"
+            " or newer)."
+        ),
+        _deprecated_DifferentiallyPrivateAggregateState,
+    ),
+    "dpsgd": (
+        (
+            "optax.dpsgd is deprecated: use optax.contrib.dpsgd (optax v0.1.8"
+            " or newer)."
+        ),
+        _deprecated_dpsgd,
+    ),
+}
+# pylint: disable=g-bad-import-order
+import typing as _typing
+
+if _typing.TYPE_CHECKING:
+  # pylint: disable=reimported
+  from optax.contrib import differentially_private_aggregate
+  from optax.contrib import DifferentiallyPrivateAggregateState
+  from optax.contrib import dpsgd
+  # pylint: enable=reimported
+
+else:
+  from optax._src.deprecations import deprecation_getattr as _deprecation_getattr
+
+  __getattr__ = _deprecation_getattr(__name__, _deprecations)
+  del _deprecation_getattr
+del _typing
+# pylint: enable=g-bad-import-order
+# pylint: enable=g-import-not-at-top
+# pylint: enable=g-importing-member
+
 
-__version__ = "0.2.2"
+__version__ = "0.2.4"
 
 __all__ = (
     "adabelief",
@@ -226,6 +301,7 @@ __all__ = (
     "adamax",
     "adamaxw",
     "adamw",
+    "adan",
     "adaptive_grad_clip",
     "AdaptiveGradClipState",
     "add_decayed_weights",
@@ -238,6 +314,7 @@ __all__ = (
     "apply_updates",
     "ApplyEvery",
     "ApplyIfFiniteState",
+    "assignment",
     "centralize",
     "chain",
     "clip_by_block_rms",
@@ -245,6 +322,10 @@ __all__ = (
     "clip",
     "ClipByGlobalNormState",
     "ClipState",
+    "conditionally_mask",
+    "ConditionallyMaskState",
+    "conditionally_transform",
+    "ConditionallyTransformState",
     "constant_schedule",
     "ctc_loss",
     "ctc_loss_with_forward_probs",
@@ -278,6 +359,7 @@ __all__ = (
     "l2_loss",
     "lamb",
     "lars",
+    "lbfgs",
     "lion",
     "linear_onecycle_schedule",
     "linear_schedule",
@@ -301,6 +383,7 @@ __all__ = (
     "noisy_sgd",
     "novograd",
     "NonNegativeParamsState",
+    "ntxent",
     "OptState",
     "Params",
     "periodic_update",
@@ -314,6 +397,7 @@ __all__ = (
     "radam",
     "rmsprop",
     "rprop",
+    "safe_increment",
     "safe_int32_increment",
     "safe_norm",
     "safe_root_mean_squares",
@@ -321,9 +405,11 @@ __all__ = (
     "scale_by_adadelta",
     "scale_by_adam",
     "scale_by_adamax",
+    "scale_by_adan",
     "scale_by_amsgrad",
     "scale_by_backtracking_linesearch",
     "scale_by_belief",
+    "scale_by_lbfgs",
     "scale_by_lion",
     "scale_by_factored_rms",
     "scale_by_novograd",
@@ -335,17 +421,21 @@ __all__ = (
     "scale_by_rprop",
     "scale_by_rss",
     "scale_by_schedule",
+    "scale_by_sign",
     "scale_by_sm3",
     "scale_by_stddev",
     "scale_by_trust_ratio",
     "scale_by_yogi",
+    "scale_by_zoom_linesearch",
     "scale_gradient",
     "scale",
     "ScaleByAdaDeltaState",
     "ScaleByAdamState",
+    "ScaleByAdanState",
     "ScaleByAmsgradState",
     "ScaleByBacktrackingLinesearchState",
     "ScaleByBeliefState",
+    "ScaleByLBFGSState",
     "ScaleByLionState",
     "ScaleByNovogradState",
     "ScaleByRmsState",
@@ -355,6 +445,7 @@ __all__ = (
     "ScaleByScheduleState",
     "ScaleBySM3State",
     "ScaleByTrustRatioState",
+    "ScaleByZoomLinesearchState",
     "ScaleState",
     "Schedule",
     "set_to_zero",
@@ -362,6 +453,7 @@ __all__ = (
     "sgdr_schedule",
     "ShouldSkipUpdateFunction",
     "sigmoid_binary_cross_entropy",
+    "sign_sgd",
     "skip_large_updates",
     "skip_not_finite",
     "sm3",
@@ -382,6 +474,7 @@ __all__ = (
     "yogi",
     "zero_nans",
     "ZeroNansState",
+    "ZoomLinesearchInfo",
 )
 
 #  _________________________________________