optax 0.2.2__tar.gz → 0.2.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: optax
-Version: 0.2.2
+Version: 0.2.3
 Summary: A gradient processing and optimisation library in JAX.
 Keywords: python,machine learning,reinforcement-learning
 Author-email: Google DeepMind <optax-dev@google.com>
@@ -18,9 +18,10 @@ 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: 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
@@ -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
 

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

@@ -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
 

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

@@ -14,12 +14,16 @@
 # ==============================================================================
 """Optax: composable gradient processing and optimization, in JAX."""
 
+# pylint: disable=wrong-import-position
+# pylint: disable=g-importing-member
+
 from optax import contrib
 from optax import losses
 from optax import monte_carlo
 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
@@ -33,6 +37,7 @@ 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
@@ -86,7 +91,10 @@ 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
@@ -99,10 +107,10 @@ 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
@@ -110,6 +118,7 @@ from optax._src.transform import scale_by_adamax
 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
@@ -130,6 +139,7 @@ from optax._src.transform import ScaleByAdaDeltaState
 from optax._src.transform import ScaleByAdamState
 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 +152,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 +160,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 +176,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
@@ -199,8 +215,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 +226,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.3"
 
 __all__ = (
     "adabelief",
@@ -245,6 +310,10 @@ __all__ = (
     "clip",
     "ClipByGlobalNormState",
     "ClipState",
+    "conditionally_mask",
+    "ConditionallyMaskState",
+    "conditionally_transform",
+    "ConditionallyTransformState",
     "constant_schedule",
     "ctc_loss",
     "ctc_loss_with_forward_probs",
@@ -278,6 +347,7 @@ __all__ = (
     "l2_loss",
     "lamb",
     "lars",
+    "lbfgs",
     "lion",
     "linear_onecycle_schedule",
     "linear_schedule",
@@ -301,6 +371,7 @@ __all__ = (
     "noisy_sgd",
     "novograd",
     "NonNegativeParamsState",
+    "ntxent",
     "OptState",
     "Params",
     "periodic_update",
@@ -324,6 +395,7 @@ __all__ = (
     "scale_by_amsgrad",
     "scale_by_backtracking_linesearch",
     "scale_by_belief",
+    "scale_by_lbfgs",
     "scale_by_lion",
     "scale_by_factored_rms",
     "scale_by_novograd",
@@ -339,6 +411,7 @@ __all__ = (
     "scale_by_stddev",
     "scale_by_trust_ratio",
     "scale_by_yogi",
+    "scale_by_zoom_linesearch",
     "scale_gradient",
     "scale",
     "ScaleByAdaDeltaState",
@@ -346,6 +419,7 @@ __all__ = (
     "ScaleByAmsgradState",
     "ScaleByBacktrackingLinesearchState",
     "ScaleByBeliefState",
+    "ScaleByLBFGSState",
     "ScaleByLionState",
     "ScaleByNovogradState",
     "ScaleByRmsState",
@@ -355,6 +429,7 @@ __all__ = (
     "ScaleByScheduleState",
     "ScaleBySM3State",
     "ScaleByTrustRatioState",
+    "ScaleByZoomLinesearchState",
     "ScaleState",
     "Schedule",
     "set_to_zero",
@@ -382,6 +457,7 @@ __all__ = (
     "yogi",
     "zero_nans",
     "ZeroNansState",
+    "ZoomLinesearchInfo",
 )
 
 #  _________________________________________

{optax-0.2.2 → optax-0.2.3}/optax/_src/alias.py

@@ -23,6 +23,7 @@ from optax._src import base
 from optax._src import clipping
 from optax._src import combine
 from optax._src import factorized
+from optax._src import linesearch as _linesearch
 from optax._src import transform
 from optax._src import wrappers
 
@@ -530,15 +531,50 @@ def adamw(
     *,
     nesterov: bool = False,
 ) -> base.GradientTransformation:
-  """Adam with weight decay regularization.
+  r"""Adam with weight decay regularization.
 
   AdamW uses weight decay to regularize learning towards small weights, as
   this leads to better generalization. In SGD you can also use L2 regularization
   to implement this as an additive loss term, however L2 regularization
-  does not behave as intended for adaptive gradient algorithms such as Adam.
+  does not behave as intended for adaptive gradient algorithms such as Adam,
+  see [Loshchilov et al, 2019].
+
+  Let :math:`\alpha_t` represent the learning rate and :math:`\beta_1, \beta_2`,
+  :math:`\varepsilon`, :math:`\bar{\varepsilon}` represent the arguments
+  ``b1``, ``b2``, ``eps`` and ``eps_root`` respectively. The learning rate is
+  indexed by :math:`t` since the learning rate may also be provided by a
+  schedule function. Let :math:`\lambda` be the weight decay and 
+  :math:`\theta_t` the parameter vector at time :math:`t`.
+
+  The ``init`` function of this optimizer initializes an internal state
+  :math:`S_0 := (m_0, v_0) = (0, 0)`, representing initial estimates for the
+  first and second moments. In practice these values are stored as pytrees
+  containing all zeros, with the same shape as the model updates.
+  At step :math:`t`, the ``update`` function of this optimizer takes as
+  arguments the incoming gradients :math:`g_t`, the optimizer state :math:`S_t` 
+  and the parameters :math:`\theta_t` and computes updates :math:`u_t` and 
+  new state :math:`S_{t+1}`. Thus, for :math:`t > 0`, we have,
+
+  .. math::
+
+    \begin{align*}
+      m_t &\leftarrow \beta_1 \cdot m_{t-1} + (1-\beta_1) \cdot g_t \\
+      v_t &\leftarrow \beta_2 \cdot v_{t-1} + (1-\beta_2) \cdot {g_t}^2 \\
+      \hat{m}_t &\leftarrow m_t / {(1-\beta_1^t)} \\
+      \hat{v}_t &\leftarrow v_t / {(1-\beta_2^t)} \\
+      u_t &\leftarrow -\alpha_t \cdot \left( \hat{m}_t / \left({\sqrt{\hat{v}_t 
+      + \bar{\varepsilon}} + \varepsilon} \right) + \lambda \theta_{t} \right)\\
+      S_t &\leftarrow (m_t, v_t).
+    \end{align*}
 
   This implementation can incorporate a momentum a la Nesterov introduced by
   [Dozat 2016]. The resulting optimizer is then often referred as NAdamW.
+  With the keyword argument `nesterov=True`, the optimizer uses Nesterov
+  momentum, replacing the above :math:`\hat{m}_t` with
+
+  .. math::
+      \hat{m}_t \leftarrow
+        \beta_1 m_t / {(1-\beta_1^{t+1})} + (1 - \beta_1) g_t / {(1-\beta_1^t)}.
   
   Examples:
     >>> import optax
@@ -1206,7 +1242,9 @@ def radam(
     b2: float = 0.999,
     eps: float = 1e-8,
     eps_root: float = 0.0,
-    threshold: float = 5.0
+    threshold: float = 5.0,
+    *,
+    nesterov: bool = False,
 ) -> base.GradientTransformation:
   """The Rectified Adam optimizer.
 
@@ -1250,13 +1288,20 @@ def radam(
       in RMSProp), to avoid dividing by zero when rescaling. This is needed for
       instance when computing (meta-)gradients through Adam.
     threshold: Threshold for variance tractability.
+    nesterov: Whether to use Nesterov momentum.
 
   Returns:
     The corresponding `GradientTransformation`.
   """
   return combine.chain(
       transform.scale_by_radam(
-          b1=b1, b2=b2, eps=eps, eps_root=eps_root, threshold=threshold),
+          b1=b1,
+          b2=b2,
+          eps=eps,
+          eps_root=eps_root,
+          threshold=threshold,
+          nesterov=nesterov,
+      ),
       transform.scale_by_learning_rate(learning_rate),
   )
 
@@ -1373,12 +1418,12 @@ def sgd(
   .. math::
 
     \begin{align*}
-      m_t & \leftarrow \begin{cases}
-        g_t + \mu m_t & \mathrm{if \ \texttt{nesterov=False}}\\
-        g_t + \mu (g_t + \mu m_{t-1}) & \mathrm{if \ \texttt{nesterov=True}}
-      \end{cases} \\
-      u_t & \leftarrow - \alpha_t m_t \\
-      S_t & = m_t,
+      m_t &\leftarrow g_t + \mu m_{t-1} \\
+      u_t &\leftarrow \begin{cases}
+        -\alpha_t m_t & \text{ if } \texttt{nesterov = False} \\
+        -\alpha_t (g_t + \mu m_t) & \text{ if } \texttt{nesterov = True}
+        \end{cases} \\
+      S_t &\leftarrow m_t,
     \end{align*}
 
   where :math:`\mu` is the ``momentum`` parameter and :math:`S_t` is the state
@@ -1546,7 +1591,37 @@ def adamax(
     b2: float = 0.999,
     eps: float = 1e-8,
 ) -> base.GradientTransformation:
-  """A variant of the Adam optimizer that uses the infinity norm.
+  r"""A variant of the Adam optimizer that uses the infinity norm.
+  
+  AdaMax is a variant of the :func:`optax.adam` optimizer. By generalizing 
+  Adam's :math:`L^2` norm to an :math:`L^p` norm and taking the limit as
+  :math:`p \rightarrow \infty`, we obtain a simple and stable update rule.
+
+  Let :math:`\alpha_t` represent the learning rate and :math:`\beta_1, \beta_2`,
+  :math:`\varepsilon` represent the arguments
+  ``b1``, ``b2`` and ``eps`` respectively. The learning rate is
+  indexed by :math:`t` since the learning rate may also be provided by a
+  schedule function.
+
+  The ``init`` function of this optimizer initializes an internal state
+  :math:`S_0 := (m_0, v_0) = (0, 0)`, representing initial estimates for the
+  first and second moments. In practice these values are stored as pytrees
+  containing all zeros, with the same shape as the model updates.
+  At step :math:`t`, the ``update`` function of this optimizer takes as
+  arguments the incoming gradients :math:`g_t` and optimizer state :math:`S_t`
+  and computes updates :math:`u_t` and new state :math:`S_{t+1}`. Thus, for
+  :math:`t > 0`, we have,
+
+  .. math::
+
+    \begin{align*}
+      m_t &\leftarrow \beta_1 \cdot m_{t-1} + (1-\beta_1) \cdot g_t \\
+      v_t &\leftarrow \max(\left| g_t \right| + \varepsilon, \beta_2 \cdot
+      v_{t-1}) \\
+      \hat{m}_t &\leftarrow m_t / (1-\beta_1^t) \\
+      u_t &\leftarrow -\alpha_t \cdot \hat{m}_t / v_t \\
+      S_t &\leftarrow (m_t, v_t).
+    \end{align*}
   
   Examples:
     >>> import optax
@@ -1700,11 +1775,14 @@ def rprop(
     Objective function: 1.38E+01
 
   References:
-    PyTorch implementation:
-      https://pytorch.org/docs/stable/generated/torch.optim.Rprop.html
-    Riedmiller and Braun, 1993: https://ieeexplore.ieee.org/document/298623
-    Igel and Hüsken, 2003:
-      https://www.sciencedirect.com/science/article/abs/pii/S0925231201007007
+    Riedmiller and Braun. `A direct adaptive method for faster backpropagation 
+    learning: the RPROP algorithm 
+    <https://ieeexplore.ieee.org/document/298623>`_, 1993
+
+    Igel and Hüsken.  `Empirical evaluation of the improved Rprop learning 
+    algorithms
+    <https://www.sciencedirect.com/science/article/abs/pii/S0925231201007007>`_,
+    2003
 
   Args:
     learning_rate: The initial step size.
@@ -1809,3 +1887,136 @@ def polyak_sgd(
           max_learning_rate=max_learning_rate, f_min=f_min, eps=eps
       ),
   )
+
+
+def lbfgs(
+    learning_rate: Optional[base.ScalarOrSchedule] = None,
+    memory_size: int = 10,
+    scale_init_precond: bool = True,
+    linesearch: Optional[
+        base.GradientTransformationExtraArgs
+    ] = _linesearch.scale_by_zoom_linesearch(max_linesearch_steps=15),
+) -> base.GradientTransformationExtraArgs:
+  r"""L-BFGS optimizer.
+
+  L-BFGS is a quasi-Newton method that multiplies the update (gradient)
+  with an approximation of the inverse Hessian. This algorithm does not need
+  access to the Hessian, as this approximation is constructed from the gradient
+  evaluations seen during optimization. L-BFGS is a limited-memory variant of
+  the Broyden-Fletcher-Goldfarb-Shanno (BFGS) algorithm. The BFGS algorithm
+  requires storing a matrix of size :math:`p \times p` with :math:`p` the
+  dimension of the parameters.
+  The limited variant circuments this issue by computing the approximation of
+  the inverse using only :math:`m` (``memory_size``) past differences of
+  parameters/gradients. Namely, the approximation of the Hessian inverse is
+  denoted :math:`P_k = P_{k, k}`, where
+
+  .. math::
+
+    \begin{align*}
+      P_{k, j+1} & = V_j^\top P_{k, j} V_j + \rho_j \delta w_j \delta w_j^\top
+      \quad \text{for} \ j \in \{k-m, \ldots, k-1\}\\
+      P_{k, k-m} & = \gamma_k I \\
+      V_k & = I - \rho_k \delta u_k \delta w_k^\top \\
+      \rho_k & = 1/(\delta u_k^\top \delta w_k) \\
+      \delta w_k & = w_{k+1} - w_k \\
+      \delta u_k & = u_{k+1} - u_k \\
+      \gamma_k & =
+        \begin{cases}
+          (\delta w_{k-1}^\top \delta u_{k-1}) /
+          (\delta u_{k-1}^\top \delta u_{k-1})
+          & \text{if} \ \texttt{scale\_init\_hess} \\
+          1 & \text{otherwise}
+        \end{cases},
+    \end{align*}
+
+  for
+  :math:`u_k` the gradients/updates at iteration :math:`k`,
+  :math:`w_k` the parameters at iteration :math:`k`.
+
+  The formula for updating :math:`P_k` is obtained by computing the optimal
+  preconditioning matrix subject to some secant condition, see references
+  for more details. Computing :math:`P_k u_k` can be done by a sequence of 
+  vector operations using past differences of parameters and gradients stored in
+  a memory bufffer.
+
+  The present function just outputs the LBFGS direction :math:`P_k u_k`.
+  It can be chained with a linesearch ensuring sufficient decrease and low
+  curvature, such as a zoom linesearch. The linesearch computes a stepsize
+  :math:`\eta_k`, such that the updated parameters
+  (using :func:`optax.apply_updates`) take the form
+  :math:`w_{k+1} = w_k - \eta_k P_k u_k`.
+
+  Example:
+    >>> import optax
+    >>> import jax
+    >>> import jax.numpy as jnp
+    >>> def f(x): return jnp.sum(x ** 2)
+    >>> solver = optax.lbfgs()
+    >>> params = jnp.array([1., 2., 3.])
+    >>> print('Objective function: ', f(params))
+    Objective function:  14.0
+    >>> opt_state = solver.init(params)
+    >>> value_and_grad = optax.value_and_grad_from_state(f)
+    >>> for _ in range(5):
+    ...   value, grad = value_and_grad(params, state=opt_state)
+    ...   updates, opt_state = solver.update(
+    ...      grad, opt_state, params, value=value, grad=grad, value_fn=f
+    ...   )
+    ...   params = optax.apply_updates(params, updates)
+    ...   print('Objective function: ', f(params))
+    Objective function:  0.0
+    Objective function:  0.0
+    Objective function:  0.0
+    Objective function:  0.0
+    Objective function:  0.0
+
+  .. warning::
+    This method is memory intensive optimizer, best used for small to medium
+    scale problems.
+
+  .. warning::
+    This optimizer works best with a linesearch (current default is a
+    zoom linesearch). See example above for best use in a non-stochastic
+    setting, where we can recycle gradients computed by the linesearch using
+    :func:`optax.value_and_grad_from_state`.
+
+  References:
+    Algorithms 7.4, 7.5 (page 199) of Nocedal et al, `Numerical Optimization
+    <https://www.math.uci.edu/~qnie/Publications/NumericalOptimization.pdf>`_
+    , 1999
+
+    Liu et al., `On the limited memory BFGS method for large scale optimization
+    <https://users.iems.northwestern.edu/~nocedal/PDFfiles/limited-memory.pdf>`_
+    , 1989.
+
+  Args:
+    learning_rate: optional global scaling factor, either fixed or evolving
+      along iterations with a scheduler, see 
+      :func:`optax.scale_by_learning_rate`. By default the learning rate is
+      handled by a linesearch.
+    memory_size: number of past updates to keep in memory to approximate the
+      Hessian inverse.
+    scale_init_precond: whether to use a scaled identity as the initial
+      preconditioner, see formula above.
+    linesearch: an instance of :class:`optax.GradientTransformationExtraArgs`
+      such as :func:`optax.scale_by_zoom_linesearch` that computes a
+      learning rate, a.k.a. stepsize, to satisfy some criterion such as a
+      sufficient decrease of the objective by additional calls to the objective.
+
+  Returns:
+    A :class:`optax.GradientTransformationExtraArgs` object.
+  """
+  if learning_rate is None:
+    base_scaling = transform.scale(-1.0)
+  else:
+    base_scaling = transform.scale_by_learning_rate(learning_rate)
+  if linesearch is None:
+    linesearch = base.identity()
+  return combine.chain(
+      transform.scale_by_lbfgs(
+          memory_size=memory_size, scale_init_precond=scale_init_precond
+      ),
+      base_scaling,
+      linesearch,
+  )