optax 0.2.3__tar.gz → 0.2.4__tar.gz

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

@@ -1,7 +1,7 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: optax
-Version: 0.2.3
-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,7 +17,7 @@ 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: chex>=0.1.87
 Requires-Dist: jax>=0.4.27
 Requires-Dist: jaxlib>=0.4.27
 Requires-Dist: numpy>=1.18.0
@@ -74,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
@@ -154,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
@@ -165,7 +165,7 @@ pip install -e ".[docs]"
 ```
 Then, execute the following.
 ```sh
-cd docs/
+cd docs
 make html
 ```
 

{optax-0.2.3 → 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
@@ -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.3 → optax-0.2.4}/optax/__init__.py

@@ -17,9 +17,11 @@
 # 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
@@ -33,6 +35,7 @@ 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
@@ -44,12 +47,14 @@ 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
@@ -98,6 +103,7 @@ 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
@@ -115,6 +121,7 @@ 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
@@ -131,12 +138,14 @@ 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
@@ -198,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
@@ -280,7 +290,7 @@ del _typing
 # pylint: enable=g-importing-member
 
 
-__version__ = "0.2.3"
+__version__ = "0.2.4"
 
 __all__ = (
     "adabelief",
@@ -291,6 +301,7 @@ __all__ = (
     "adamax",
     "adamaxw",
     "adamw",
+    "adan",
     "adaptive_grad_clip",
     "AdaptiveGradClipState",
     "add_decayed_weights",
@@ -303,6 +314,7 @@ __all__ = (
     "apply_updates",
     "ApplyEvery",
     "ApplyIfFiniteState",
+    "assignment",
     "centralize",
     "chain",
     "clip_by_block_rms",
@@ -385,6 +397,7 @@ __all__ = (
     "radam",
     "rmsprop",
     "rprop",
+    "safe_increment",
     "safe_int32_increment",
     "safe_norm",
     "safe_root_mean_squares",
@@ -392,6 +405,7 @@ __all__ = (
     "scale_by_adadelta",
     "scale_by_adam",
     "scale_by_adamax",
+    "scale_by_adan",
     "scale_by_amsgrad",
     "scale_by_backtracking_linesearch",
     "scale_by_belief",
@@ -407,6 +421,7 @@ __all__ = (
     "scale_by_rprop",
     "scale_by_rss",
     "scale_by_schedule",
+    "scale_by_sign",
     "scale_by_sm3",
     "scale_by_stddev",
     "scale_by_trust_ratio",
@@ -416,6 +431,7 @@ __all__ = (
     "scale",
     "ScaleByAdaDeltaState",
     "ScaleByAdamState",
+    "ScaleByAdanState",
     "ScaleByAmsgradState",
     "ScaleByBacktrackingLinesearchState",
     "ScaleByBeliefState",
@@ -437,6 +453,7 @@ __all__ = (
     "sgdr_schedule",
     "ShouldSkipUpdateFunction",
     "sigmoid_binary_cross_entropy",
+    "sign_sgd",
     "skip_large_updates",
     "skip_not_finite",
     "sm3",