src/kfac_jax/_src/optimizer.py

# Modifications copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0

# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""The kfac_jax optimizer (supporting K-FAC and other methods)."""

import functools
from typing import Any, Callable, Generic, Iterator, Sequence
from absl import logging

import jax
from jax import lax
import jax.numpy as jnp
from kfac_jax._src import curvature_estimator
from kfac_jax._src import utils
from typing_extensions import Self


# Types for annotation
Array = utils.Array
PRNGKey = utils.PRNGKey
Numeric = utils.Numeric
Params = utils.Params
Batch = utils.Batch
FuncState = Any
FuncAux = utils.FuncAux
Scalar = utils.Scalar
ScheduleType = utils.ScheduleType

FuncArgsVariants = (
    tuple[Params, Batch] |
    tuple[Params, FuncState, Batch] |
    tuple[Params, PRNGKey, Batch] |
    tuple[Params, FuncState, PRNGKey, Batch]
)
FuncOutputs = (
    Array |
    tuple[Array, FuncState] |
    tuple[Array, FuncAux] |
    tuple[Array, tuple[FuncState, FuncAux]]
)
ValueFunc = Callable[..., FuncOutputs]
ValueAndGradFunc = Callable[..., tuple[FuncOutputs, Params]]
SharedForwardFunc = Callable[..., tuple[Array, Array]]
BlockDiagonalCurvature = curvature_estimator.BlockDiagonalCurvature

ReturnEither = (
    tuple[Params, "Optimizer.State", FuncState, dict[str, Numeric]] |
    tuple[Params, "Optimizer.State", dict[str, Numeric]]
)

QuadModelParams = tuple[Array, Array, Array, Array]
# The quadratic model is given as
#  Q(w) = w^T V^T (C + damping * I + reg * L) V w / 2.0 + w^T V^T g
#  where (n - number of vectors, d - dimensions of each vector):
#   damping  - the damping value at the current iteration
#   reg      - the L2 regularization coefficient
#   w (n,)   - the vector of free weights (learning rate and momentum)
#   V (d, n) - the matrix of proposed vectors for each weight
#   C (d, d) - the curvature matrix (GGN/Fisher/Hessian)
#   L (d, d) - the L2 regularization matrix. L is diagonal, with 1 on diagonal
#              if the corresponding parameter is L2 regularised, and 0
#              otherwise.
#   g (d,)   - the gradient
#
# In QuadModelParams, we have the tuple (A, D, R, b) where:
#   A = V^T C V
#   D = V^T I V  (for damping)
#   R = V^T L V  (for L2 regularization)
#   b = V^T g
#
# See Optimizer._solve_quad_model for how these are used, and
# Optimizer._compute_exact_quad_model for how they are computed.


# Various lists of parameters that are biases and norms, to be
# used for registering parameters that are excluded from l2 regularization
# in Optimizer.
# "b" and "bias" for biases, "scale" for RMSNorm and LayerNorm, and
# "offset" for LayerNorm.

HAIKU_BIASES = "b,bias"
HAIKU_BIASES_AND_NORMS = "b,bias,scale,offset"


class Optimizer(utils.WithStagedMethods):
  """The kfac_jax optimizer (supporting K-FAC and other methods)."""

  @utils.register_state_class
  class State(Generic[Params], utils.State):
    r"""Persistent state of the optimizer.

    Attributes:
      velocities: The update to the parameters from the previous step -
        :math:`\theta_t - \theta_{t-1}`.
      estimator_state: The persistent state for the curvature estimator.
      damping: When using damping adaptation, this will contain the current
        value.
      data_seen: The number of training cases that the optimizer has processed.
      step_counter: An integer giving the current step number :math:`t`.
    """
    velocities: Params
    estimator_state: BlockDiagonalCurvature.State
    damping: Array
    data_seen: Numeric
    step_counter: Numeric

    @classmethod
    def from_dict(cls, dict_representation: dict[str, Any]) -> Self:
      dict_representation["estimator_state"] = (
          BlockDiagonalCurvature.State.from_dict(
              dict_representation["estimator_state"]
          )
      )
      return cls(**dict_representation)

  def __init__(
      self,
      value_and_grad_func: ValueAndGradFunc,
      l2_reg: Numeric,
      regularized_parameters_path_exclusions: str = "",
      value_func_has_aux: bool = False,
      value_func_has_state: bool = False,
      value_func_has_rng: bool = False,
      value_func_for_estimator: ValueFunc | None = None,
      use_adaptive_learning_rate: bool = False,
      learning_rate_schedule: ScheduleType | None = None,
      use_adaptive_momentum: bool = False,
      momentum_schedule: ScheduleType | None = None,
      use_adaptive_damping: bool = False,
      damping_schedule: ScheduleType | None = None,
      initial_damping: Numeric | None = None,
      use_initial_damping_calibration: bool = False,
      min_damping: Numeric = 1e-8,
      max_damping: Numeric = jnp.inf,
      include_damping_in_quad_change: bool = False,
      damping_adaptation_interval: int = 5,
      damping_adaptation_decay: Numeric = 0.9,
      damping_lower_threshold: Numeric = 0.25,
      damping_upper_threshold: Numeric = 0.75,
      always_use_exact_qmodel_for_damping_adjustment: bool = False,
      precon_damping_mult: Numeric = 1.0,
      precon_damping_schedule: ScheduleType | None = None,
      use_step_rejection: bool = False,
      reject_damping_increase_factor: float = 1.0,
      norm_constraint: Numeric | None = None,
      num_burnin_steps: int = 10,
      estimation_mode: str | None = None,
      custom_estimator_ctor: (
          Callable[..., BlockDiagonalCurvature] | None) = None,
      curvature_ema: Numeric = 0.95,
      curvature_update_period: int = 1,
      inverse_update_period: int = 5,
      use_exact_inverses: bool = False,
      batch_process_func: Callable[[Batch], Batch] | None = None,
      register_only_generic: bool = False,
      patterns_to_skip: Sequence[str] = (),
      use_automatic_registration: bool = True,
      auto_register_kwargs: dict[str, Any] | None = None,
      layer_tag_to_block_ctor: (
          dict[str, curvature_estimator.CurvatureBlockCtor] | None) = None,
      multi_device: bool = False,
      debug: bool = False,
      invalid_metric_value: Numeric = jnp.nan,
      batch_size_extractor: Callable[
          [Batch], Numeric
      ] = utils.default_batch_size_extractor,
      pmap_axis_name: str = "batch_axis",
      forbid_setting_attributes_after_finalize: bool = True,
      modifiable_attribute_exceptions: Sequence[str] = (),
      include_norms_in_stats: bool = False,
      include_per_param_norms_in_stats: bool = False,
      include_registered_loss_in_stats: bool = False,
      distributed_precon_apply: bool = True,
      distributed_inverses: bool = True,
      num_estimator_samples: int = 1,
      should_vmap_estimator_samples: bool = False,
      norm_to_scale_identity_weight_per_block: str | None = None,
      step_stats_hook: Callable[..., dict[str, Array]] | None = None,
      precon_power: Scalar = -1.0,
      exact_quad_model_matrix_type: str | None = None,
      value_func_for_shared_forward: SharedForwardFunc | None = None,
      share_curvature_and_grad_forward: bool = False,
  ):
    """Initializes the kfac_jax optimizer with the provided settings.

    NOTE: Please read the docstring for this constructor carefully. Especially
    the description of ``value_and_grad_func``.

    A note on the "damping" parameter:

    One of the main complications of using second-order optimizers like K-FAC is
    the "damping" parameter. This parameter is multiplied by the identity matrix
    and (approximately) added to the curvature matrix (i.e. the Fisher or GGN)
    before it is inverted and multiplied by the gradient when computing the
    update (before any learning rate scaling). The damping should follow the
    scale of the objective, so that if you multiply your loss by some factor you
    should do the same for the damping. Roughly speaking, larger damping values
    constrain the update vector to a smaller region around zero, which is needed
    in general since the second-order approximations that underlie second-order
    methods can break down for large updates. (In gradient descent the learning
    rate plays an analogous role.) The relationship between the damping
    parameter and the radius of this region is complicated and depends on the
    scale of the objective amongst other things.

    The optimizer provides a system for adjusting the damping automatically via
    the ``use_adaptive_damping`` argument, although this system is not reliable,
    especially for highly stochastic objectives. Using a fixed value or a
    manually tuned schedule can work as good or better for some problems, while
    it can be a very poor choice for others (like deep autoencoders).
    Empirically we have found that using a fixed value works well enough for
    common architectures like convnets and transformers.

    Args:
      value_and_grad_func: Python callable. This function should return the
        value of the loss to be optimized and its gradients, and optionally the
        model state and auxiliary information in the form of a a dict mapping
        strings to scalar arrays (usually statistics to log). Note that it
        should *not* be jitted/pmapped or otherwise compiled by JAX, as this can
        lead to errors. (Compilation is done internally by the optimizer.) The
        interface of this function should be: ``out_args, loss_grads =
        value_and_grad_func(*in_args)``. Here, ``in_args`` is ``(params,
        func_state, rng, batch)``, with ``rng`` omitted if
        ``value_func_has_rng`` is ``False``, and with ``func_state`` omitted if
        ``value_func_has_state`` is ``False``. Meanwhile, ``out_args`` is
        ``(loss, (func_state, aux))`` if ``value_func_has_state`` and
        ``value_func_has_aux`` are both ``True``, ``(loss, func_state)`` if
        ``value_func_has_state`` is ``True`` and ``value_func_has_aux`` is
        ``False``, ``(loss, aux)`` if ``value_func_has_state`` is ``False`` and
        ``value_func_has_aux`` is ``True``, and finally ``loss`` if
        ``value_func_has_state`` and ``value_func_has_aux`` are both ``False``.
        This should be consistent with how JAX's ``value_and_grad`` API function
        is typically used. Note that the value (and its gradient) should be
        normalized by the batch size, as is standard convention. Additional
        normalization, such as by the sequence length, is up to the user, but
        must by properly reported in the loss registration (by setting the
        ``weight`` arguments in the loss registration functions.)
      l2_reg: Scalar. Set this value to tell the optimizer what L2
        regularization coefficient you are using (if any). Note the coefficient
        appears in the regularizer as ``coeff / 2 * sum(param**2)``. This adds
        an additional diagonal term to the curvature and hence will affect the
        quadratic model when using adaptive damping. Note that the user is still
        responsible for adding regularization to the loss.
      regularized_parameters_path_exclusions: str. A comma-separated list
        specifying the names of parameters that should not be regularized.
        A number of convenience examples are given in this module, e.g.
        HAIKU_BIASES_AND_NORMS, which is ``"b,bias,scale,offset"``.
        (Default: ``""``)
      value_func_has_aux: Boolean. Specifies whether the provided callable
        ``value_and_grad_func`` returns auxiliary data. (Default: ``False``)
      value_func_has_state: Boolean. Specifies whether the provided callable
        ``value_and_grad_func`` has a persistent state that is passed in and
        out. (Default: ``False``)
      value_func_has_rng: Boolean. Specifies whether the provided callable
        ``value_and_grad_func`` additionally takes as input an rng key.
        (Default: ``False``)
      value_func_for_estimator: ValueFunc. If specified, this function will be
        used by the preconditioner estimator instead of ``value_and_grad_func``.
        This is useful for cases where the value function used for training is
        expensive to add to the preconditioner, e.g. because it has costly
        regularizers. (Default: ``None``)
      value_func_for_shared_forward: Tagged function returning
        ``(loss, gradient_surrogate)``. The surrogate's parameter gradient
        must equal the training gradient. Required when
        ``share_curvature_and_grad_forward=True``.
      use_adaptive_learning_rate: Boolean. Specifies whether to use the special
        rule from the original K-FAC paper for picking the learning rate at each
        step. Note that this won't work well for stochastic objectives. If this
        is ``False``, the user must use the ``learning_rate`` argument of the
        step function, or the constructor argument ``learning_rate_schedule``.
        (Default: ``False``)
      learning_rate_schedule: Callable. A schedule for the learning rate. This
        should take as input the current step number, and optionally the amount
        of data seen so far as a keyword argument ``data_seen``, and return a
        single array that represents the learning rate. (Default: ``None``)
      use_adaptive_momentum: Boolean. Specifies whether to use the special rule
        from the original K-FAC paper for picking the momentum "decay" parameter
        at each step. Note that this won't work well for stochastic objectives.
        If this is ``False``, the user must use the ``momentum`` argument of the
        step function, or the constructor argument ``momentum_schedule``.
        (Default: ``False``)
      momentum_schedule: Callable. A schedule for the momentum parameter. This
        should take as input the current step number, and optionally the amount
        of data seen so far as a keyword argument ``data_seen``, and return a
        single array that represents the momentum. (Default: ``None``)
      use_adaptive_damping: Boolean. Specifies whether the optimizer will use
        the Levenberg-Marquardt method to automatically adjust the damping every
        ``damping_adaptation_interval`` iterations. If this is set to ``False``
        the user must provide a value to the damping argument of the step
        function at each iteration, or use the ``damping_schedule`` constructor
        argument. Note that the effectiveness of this technique seems to vary
        between problems. (Default: ``False``)
      damping_schedule: Callable. A schedule for the damping. This should take
        as input the current step number, and optionally the amount of data seen
        so far as a keyword argument ``data_seen``, and return a single array
        that represents the learning rate. (Default: ``None``)
      initial_damping: Scalar or None. This specifies the initial value of the
        damping that the optimizer will use when using automatic damping
        adaptation. (Default: ``None``)
      use_initial_damping_calibration: Boolean. If ``True``, the initial damping
        value, used to initialize the adaptive damping method, will be first
        calibrated (after any burnin steps to estimate the preconditioner) so
        that its value wouldn't be changed after the first step of optimization.
        This calibration is done by essentially running the step function
        multiple times without actually updating the parameters or sampling a
        new mini-batch. ``num_burnin_steps`` must be greater than 0 to use this
        option. (Default: ``False``)
      min_damping: Scalar. Minimum value the damping parameter can take when
        using automatic damping adaptation. Note that the default value of 1e-8
        is quite arbitrary, and you may have to adjust this up or down for your
        particular problem. If you are using a non-zero value of l2_reg you
        *may* be able to set this to zero. (Default: ``1e-8``)
      max_damping: Scalar. Maximum value the damping parameter can take when
        using automatic damping adaptation. (Default: ``Infinity``)
      include_damping_in_quad_change: Boolean. Whether to include the
        contribution of the damping in the quadratic model for the purposes
        computing the reduction ration ("rho") in the Levenberg-Marquardt scheme
        used for adapting the damping. Note that the contribution from the
        ``l2_reg`` argument is always included. (Default: ``False``)
      damping_adaptation_interval: Int. The number of steps in between adapting
        the damping parameter. (Default: ``5``)
      damping_adaptation_decay: Scalar. The damping parameter will be adjusted
        up or down by ``damping_adaptation_decay **
        damping_adaptation_interval``, or remain unchanged, every
        ``damping_adaptation_interval`` number of iterations. (Default: ``0.9``)
      damping_lower_threshold: Scalar. The damping parameter is increased if the
        reduction ratio is below this threshold. (Default: ``0.25``)
      damping_upper_threshold: Scalar. The damping parameter is decreased if the
        reduction ratio is below this threshold. (Default: ``0.75``)
      always_use_exact_qmodel_for_damping_adjustment: Boolean. When using
        learning rate and/or momentum adaptation, the quadratic model change
        used for damping adaption is always computed using the exact curvature
        matrix. Otherwise, there is an option to use either the exact or
        approximate curvature matrix to compute the quadratic model change,
        which is what this argument controls. When True, the exact curvature
        matrix will be used, which is more expensive, but could possibly produce
        a better damping schedule. (Default: ``False``)
      precon_damping_mult: Scalar. When ``precon_damping_schedule`` is unset,
        the regular damping is used for the preconditioner damping, multiplied
        by this value. (Default: ``1.0``)
      precon_damping_schedule: Similar to ``damping_schedule``, but for the
        preconditioner only. If ``None``, the preconditioner will use the
        regular damping, multiplied by ``precon_damping_mult``.
        (Default: ``None``)
      use_step_rejection: Whether or not to reject the step whenever the loss
        on the current batch goes up after the update. This option offers
        robustness at the cost of doing more work per step (unless adaptive
        damping with Levenberg-Marquardt is used). (Default: ``False``)
      reject_damping_increase_factor: The damping parameter is increased by this
        factor if the step is rejected. (Default: ``1.0``)
      norm_constraint: Scalar. If specified, the update is scaled down so that
        its approximate squared Fisher norm ``v^T F v`` is at most the specified
        value. (Note that here ``F`` is the approximate curvature matrix, not
        the exact.) May only be used when ``use_adaptive_learning_rate`` is
        ``False``. (Default: ``None``)
      num_burnin_steps: Int. At the start of optimization, e.g. the first step,
        before performing the actual step the optimizer will perform this many
        times updates to the curvature approximation without updating the actual
        parameters. (Default: ``10``)
      estimation_mode: String. The type of estimator to use for the curvature
        matrix. See the documentation for :class:`~BlockDiagonalCurvature` for a
        detailed description of the possible options. If ``None`` will use
        default estimation_mode mode of the used CurvatureEstimator subclass,
        which is typically "ggn_curvature_prop". (Default: ``None``)
      custom_estimator_ctor: Optional constructor for subclass of
        :class:`~BlockDiagonalCurvature`. If specified, the optimizer will use
        this conastructor instead of the default
        :class:`~BlockDiagonalCurvature`. (Default: ``None``)
      curvature_ema: The decay factor used when calculating the covariance
        estimate moving averages. (Default: ``0.95``)
      curvature_update_period: Int. The number of steps in between updating the
        the curvature estimates. (Default: ``1``)
      inverse_update_period: Int. The number of steps in between updating the
        the computation of the inverse curvature approximation. (Default: ``5``)
      use_exact_inverses: Bool. If ``True``, preconditioner inverses are
        computed "exactly" without the pi-adjusted factored damping approach.
        Note that this involves the use of eigendecompositions, which can
        sometimes be much more expensive. (Default: ``False``)
      batch_process_func: Callable. A function which to be called on each batch
        before feeding to the KFAC on device. This could be useful for specific
        device input optimizations. (Default: ``None``)
      register_only_generic: Boolean. Whether when running the auto-tagger to
        register only generic parameters, or allow it to use the graph matcher
        to automatically pick up any kind of layer tags. (Default: ``False``)
      patterns_to_skip: tuple. A list of any patterns that should be skipped by
        the graph matcher when auto-tagging. (Default: ``()``)
      use_automatic_registration: Bool. If ``True``, the optimizer will try to
        automatically register the layers of your network. (Default: ``True``)
      auto_register_kwargs: Any additional kwargs to be passed down to
        :func:`~auto_register_tags`, which is called by the curvature estimator.
        (Default: ``None``)
      layer_tag_to_block_ctor: dictionary. A mapping from layer tags to block
        classes which to override the default choices of block approximation for
        that specific tag. See the documentation for
        :class:`~CurvatureEstimator` for a more detailed description. (Default:
        ``None``)
      multi_device: Boolean. Whether to use pmap and run the optimizer on
        multiple devices. (Default: ``False``)
      debug: Boolean. If neither the step or init functions should be jitted.
        Note that this also overrides ``multi_device`` and prevents using pmap,
        instead using a "simulated pmap" that loops over the device index and
        does everything on the default device. (Default: ``False``)
      invalid_metric_value: Numeric. Certain metrics returned from the step
        function are not always computed at each iteration, or may otherwise
        be invalid. In such cases we need to return a value anyway. jnp.nan is
        a natural choice, but can sometimes cause problems (e.g. false positives
        JAX's automatic NaN checker). This argument allows the user to specify a
        different value to return in such cases. (Default: ``jnp.nan``)
      batch_size_extractor: A function that takes as input the function
        arguments and returns the batch size for a single device. (Default:
        ``kfac.utils.default_batch_size_extractor``)
      pmap_axis_name: String. The name of the pmap axis to use when
        ``multi_device`` is set to True. (Default: ``batch_axis``)
      forbid_setting_attributes_after_finalize: Boolean. By default, after the
        object is finalized, you can not set any of its properties. This is done
        in order to protect the user from making changes to the object
        attributes that would not be picked up by various internal methods after
        they have been compiled. However, if you are extending this class, and
        clearly understand the risks of modifying attributes, setting this to
        ``False`` will remove the restriction. (Default: ``True``)
      modifiable_attribute_exceptions: Sequence of strings. Gives a list of
        names for attributes that can be modified after finalization even when
        ``forbid_setting_attributes_after_finalize`` is ``True``. (Default:
        ``()``)
      include_norms_in_stats: Boolean. It True, the vector norms of the
        gradient, preconditioned gradient, and parameter update are included in
        the statistics returned by the step function. (Default: ``False``)
      include_per_param_norms_in_stats: Boolean. It True, the per-parameter
        vector norms of the gradient, preconditioned gradient, and parameter
        update are included in the statistics returned by the step function.
        (Default: ``False``)
      include_registered_loss_in_stats: Boolean. If True, we include the loss,
        as computed from the registered losses, in the stats. Also included is
        the relative difference between this as the loss computed from
        ``value_and_grad_func``. This is useful for debugging registration
        errors. Note this for this option to work it's required that the targets
        are passed for each loss function registration. (Default: ``False``)
      distributed_precon_apply: Boolean. Whether to distribute the application
        of the preconditioner across the different devices in a layer-wise
        fashion. If False, each device will (redundantly) perform the required
        operations for all the layers. (Default: True)
      distributed_inverses: Boolean. Whether to distribute the inverse
        computations (required to compute the preconditioner) across the
        different devices in a layer-wise fashion. If False, each device will
        (redundantly) perform the required computations for all the layers.
        (Default: True)
      num_estimator_samples: Number of samples (per case) to use when computing
        stochastic curvature matrix estimates. This option is only used when
        ``estimation_mode == 'fisher_gradients'`` or ``estimation_mode ==
        '[fisher,ggn]_curvature_prop'``. (Default: 1)
      should_vmap_estimator_samples: Whether to use ``jax.vmap`` to compute
        samples when ``num_estimator_samples > 1``. (Default: False)
      share_curvature_and_grad_forward: Reuse the exact-Fisher tagged model
        primal evaluation for the ordinary loss and training gradient on
        curvature-update steps. (Default: ``False``)
      norm_to_scale_identity_weight_per_block: The name of a norm to use to
        compute extra per-block scaling for the damping. See psd_matrix_norm()
        in utils/math.py for the definition of these. Note that this will not
        affect the exact quadratic model that is used as part of the "adaptive"
        learning rate, momentum, and damping methods. (Default: None)
      step_stats_hook: Optional callable ``(estimator, grads,
        preconditioned_gradient) -> dict`` invoked inside ``_step`` on the
        PRE-norm-constraint preconditioned gradient; returned scalars are
        merged into the step stats dict. Runs inside the step's jit — must
        be trace-safe and cheap. (Default: None)
      precon_power: The matrix power to use when computing the preconditioner.
        K-FAC use -1 by default, but ``kfac_jax`` can simulate other optimizers
        like RMSProp by using -0.5 (along with appropriate changes to
        ``layer_tag_to_block_ctor`` and  ``estimation_mode``). (Default: -1)
      exact_quad_model_matrix_type: The type of matrix to use when computing the
        exact quadratic model (used in the adaptive learning rate and momentum).
        Can be ``'fisher'``, ``'ggn'``, or None. If None, will use the value
        implied by ``estimation_mode``. (Default: None)
    """

    super().__init__(
        multi_device=multi_device,
        pmap_axis_name=pmap_axis_name if multi_device else None,
        debug=debug,
        forbid_setting_attributes_after_finalize=
        forbid_setting_attributes_after_finalize,
        excluded_attribute_names=modifiable_attribute_exceptions,
    )

    if use_adaptive_damping and initial_damping is None:
      raise ValueError("When use_adaptive_damping is True you must provide a "
                       "value for initial_damping.")
    if use_adaptive_learning_rate and learning_rate_schedule is not None:
      raise ValueError("If you are using adaptive learning rate then "
                       "`learning_rate_schedule` should be None.")
    if use_adaptive_momentum and momentum_schedule is not None:
      raise ValueError("If you are using adaptive momentum then "
                       "`momentum_schedule` should be None.")
    if use_adaptive_damping and damping_schedule is not None:
      raise ValueError("If you are using adaptive damping then "
                       "`damping_schedule` should be None.")

    if num_burnin_steps <= 0 and use_initial_damping_calibration:
      raise ValueError("num_burnin_steps must be > 0 if "
                       "use_initial_damping_calibration is True.")

    self._value_and_grad_func = value_and_grad_func
    self._value_func_has_aux = value_func_has_aux
    self._value_func_has_state = value_func_has_state
    self._value_func_has_rng = value_func_has_rng
    if share_curvature_and_grad_forward:
      incompatible = []
      if estimation_mode != "fisher_exact":
        incompatible.append("estimation_mode must be 'fisher_exact'")
      if value_func_for_estimator is not None:
        incompatible.append("value_func_for_estimator must be None")
      if value_func_for_shared_forward is None:
        incompatible.append("value_func_for_shared_forward must be provided")
      if custom_estimator_ctor is not None:
        incompatible.append("custom_estimator_ctor must be None")
      if value_func_has_aux:
        incompatible.append("value_func_has_aux must be False")
      if value_func_has_state:
        incompatible.append("value_func_has_state must be False")
      if value_func_has_rng:
        incompatible.append("value_func_has_rng must be False")
      if include_registered_loss_in_stats:
        incompatible.append(
            "include_registered_loss_in_stats must be False"
        )
      if incompatible:
        raise ValueError(
            "`share_curvature_and_grad_forward=True` is incompatible with: "
            + "; ".join(incompatible)
            + "."
        )
    self._share_curvature_and_grad_forward = (
        share_curvature_and_grad_forward
    )
    self._value_func: ValueFunc = convert_value_and_grad_to_value_func(
        value_and_grad_func,
        has_aux=value_func_has_aux or value_func_has_state,
    )

    self._l2_reg = l2_reg
    self._regularized_parameters_path_exclusions = (
        regularized_parameters_path_exclusions.split(","))

    self._use_adaptive_learning_rate = use_adaptive_learning_rate
    self._learning_rate_schedule = learning_rate_schedule
    self._use_adaptive_momentum = use_adaptive_momentum
    self._momentum_schedule = momentum_schedule

    self._use_adaptive_damping = use_adaptive_damping
    self._damping_schedule = damping_schedule
    self._initial_damping = initial_damping
    self._use_initial_damping_calibration = use_initial_damping_calibration
    self._min_damping = min_damping
    self._max_damping = max_damping
    self._include_damping_in_quad_change = include_damping_in_quad_change
    self._damping_adaptation_decay = damping_adaptation_decay
    self._damping_adaptation_interval = damping_adaptation_interval
    self._damping_lower_threshold = damping_lower_threshold
    self._damping_upper_threshold = damping_upper_threshold
    self._always_use_exact_qmodel_for_damping_adjustment = (
        always_use_exact_qmodel_for_damping_adjustment)
    self._precon_damping_mult = precon_damping_mult
    self._precon_damping_schedule = precon_damping_schedule

    self._use_step_rejection = use_step_rejection
    self._reject_damping_increase_factor = reject_damping_increase_factor

    self._norm_constraint = norm_constraint
    self._num_burnin_steps = num_burnin_steps
    self._curvature_ema = curvature_ema
    if curvature_update_period > inverse_update_period:
      raise ValueError(
          "curvature_update_period ({}) cannot be larger than"
          " inverse_update_period ({}) as the identical matrix inversion would"
          " be redundantly performed. Set inverse_update_period larger instead."
          .format(curvature_update_period, inverse_update_period)
      )
    self._curvature_update_period = curvature_update_period
    self._inverse_update_period = inverse_update_period
    self._layer_tag_to_block_cls = layer_tag_to_block_ctor
    self._patterns_to_skip = patterns_to_skip
    self._batch_process_func = batch_process_func or (lambda x: x)
    self._include_norms_in_stats = include_norms_in_stats
    self._include_per_param_norms_in_stats = include_per_param_norms_in_stats
    self._include_registered_loss_in_stats = include_registered_loss_in_stats
    self._batch_size_extractor = batch_size_extractor

    self.__invalid_metric_value = invalid_metric_value

    self._use_cached_inverses = (self._inverse_update_period != 1)
    self._use_exact_inverses = use_exact_inverses

    self._norm_to_scale_identity_weight_per_block = (
        norm_to_scale_identity_weight_per_block
    )

    # Optional ``hook(estimator, grads, preconditioned_gradient) -> dict``
    # called inside ``_step`` on the PRE-norm-constraint preconditioned
    # gradient; the returned scalars are merged into the step stats.
    # Runs inside the step's jit — implementations must be trace-safe
    # and cheap (reductions only). Lets clients log per-block/per-family
    # update-allocation diagnostics without subclassing ``_step``.
    self._step_stats_hook = step_stats_hook

    self._precon_power = precon_power

    self._exact_quad_model_matrix_type = exact_quad_model_matrix_type

    self._params_index = 0
    batch_index = int(value_func_has_state + value_func_has_rng + 1)

    if (norm_to_scale_identity_weight_per_block is not None
        and norm_to_scale_identity_weight_per_block != "none"):

      assert (not use_adaptive_learning_rate and not use_adaptive_momentum
              and not use_adaptive_damping)  # not currently supported

    estimator_ctor = (custom_estimator_ctor or BlockDiagonalCurvature)

    auto_register_kwargs = auto_register_kwargs or {}
    auto_register_kwargs.update(dict(
        register_only_generic=register_only_generic,
        patterns_to_skip=patterns_to_skip,
    ))

    if value_func_for_estimator is None:
      # The reason we pass value_and_grad_func to the estimator here and not
      # value_func, is that the latter is usually produced using the primal
      # computation which is part of jax.grad. For whatever reason, when JAX
      # takes the gradient of this, it produces a slightly different graph than
      # if we apply jax.grad directly to the original function. This then makes
      # it impossible for XLA to merge the two computations, defeating the
      # purpose of the estimation modes "fisher_empirical_direct[_synced]".
      func_and_grad_for_estimator = convert_value_and_grad_to_clean_value_and_grad(  # pylint: disable=line-too-long
          value_and_grad_func,
          has_aux=value_func_has_aux or value_func_has_state,
      )

    else:
      func_and_grad_for_estimator = None

    estimator_extra_kwargs = {}
    if share_curvature_and_grad_forward:
      estimator_extra_kwargs["shared_forward_value_func"] = (
          value_func_for_shared_forward
      )

    # Curvature estimator
    self._estimator = estimator_ctor(
        func=value_func_for_estimator,
        func_and_grad=func_and_grad_for_estimator,
        default_estimation_mode=estimation_mode,
        params_index=self._params_index,
        batch_index=batch_index,
        layer_tag_to_block_ctor=layer_tag_to_block_ctor,
        distributed_multiplies=distributed_precon_apply,
        distributed_cache_updates=distributed_inverses,
        num_samples=num_estimator_samples,
        should_vmap_samples=should_vmap_estimator_samples,
        auto_register_tags=use_automatic_registration,
        auto_register_kwargs=auto_register_kwargs,
        **estimator_extra_kwargs,
    )
    self._implicit = curvature_estimator.ImplicitExactCurvature(
        self._value_func,
        params_index=self._params_index,
        batch_size_extractor=batch_size_extractor,
    )

    # Each subclass should call finalize on its own, so this gets called only
    # for instances of exactly this class type.
    if type(self) == Optimizer:  # pylint: disable=unidiomatic-typecheck
      self.finalize()

  @property
  def _invalid_metric_value(self) -> Array:
    return jnp.array(self.__invalid_metric_value, dtype=float)

  @property
  def _damping_decay_factor(self) -> Numeric:
    """How fast to decay the damping, when using damping adaptation."""
    return self._damping_adaptation_decay ** self._damping_adaptation_interval

  @property
  def _exact_powers_to_cache(self) -> Numeric | Sequence[Numeric] | None:
    if self._use_exact_inverses and self._use_cached_inverses:
      return self._precon_power
    else:
      return None

  @property
  def _approx_powers_to_cache(self) -> Numeric | Sequence[Numeric] | None:
    if not self._use_exact_inverses and self._use_cached_inverses:
      return self._precon_power
    else:
      return None

  @property
  def _mat_type_for_exact_quad_model(self) -> str:
    if self._exact_quad_model_matrix_type is None:
      return self._estimator.default_mat_type
    return self._exact_quad_model_matrix_type

  def _should_update_damping(self, step_counter: int) -> bool:
    """Whether at the current step the optimizer should update the damping."""
    return ((step_counter + 1) % self._damping_adaptation_interval == 0) and (
        self._use_adaptive_damping
    )

  def _should_update_estimate_curvature(self, step_counter: int) -> bool:
    """Whether at the current step the optimizer should update the curvature estimates."""
    return step_counter % self._curvature_update_period == 0

  def _should_update_inverse_cache(
      self,
      state: State,
      inverse_update_period: Numeric | None = None,
  ) -> Array | bool:
    """Whether at the current step the optimizer should update the inverse curvature approximation."""
    period = (self._inverse_update_period if inverse_update_period is None
              else inverse_update_period)
    return self._use_cached_inverses and (
        state.step_counter % period == 0)

  def _should_sync_estimator(
      self,
      state: State,
      inverse_update_period: Numeric | None = None,
  ) -> Array | bool:
    """Whether at the current step the optimizer should update the inverse curvature approximation."""

    if self._use_cached_inverses:
      return self._should_update_inverse_cache(state, inverse_update_period)

    return True

  def set_live_hparams(
      self,
      *,
      curvature_ema: Numeric | None = None,
      curvature_update_period: int | None = None,
      inverse_update_period: int | None = None,
  ) -> None:
    """Updates cadence/EMA hyperparameters on a live optimizer instance.

    All three take effect from the next call to ``step`` without triggering
    recompilation: ``curvature_update_period`` only enters Python-side
    executable selection, while ``curvature_ema`` and ``inverse_update_period``
    are threaded into the compiled step function as runtime scalars.

    ``inverse_update_period`` cannot be changed on an optimizer constructed
    with ``inverse_update_period=1``, since that construction permanently
    disables the inverse cache in the estimator state.
    """
    new_curv = (self._curvature_update_period if curvature_update_period is None
                else int(curvature_update_period))
    new_inv = (self._inverse_update_period if inverse_update_period is None
               else int(inverse_update_period))
    if new_curv < 1 or new_inv < 1:
      raise ValueError("Update periods must be positive integers.")
    if new_inv != self._inverse_update_period and not self._use_cached_inverses:
      raise ValueError(
          "Cannot change inverse_update_period on an optimizer constructed "
          "with inverse_update_period=1 (inverse cache disabled).")
    if new_curv > new_inv:
      raise ValueError(
          "curvature_update_period ({}) cannot be larger than"
          " inverse_update_period ({}).".format(new_curv, new_inv))
    if curvature_ema is not None and not 0.0 <= float(curvature_ema) <= 1.0:
      raise ValueError("curvature_ema must be in [0, 1].")
    self.unlock_attributes()
    try:
      self._curvature_update_period = new_curv
      self._inverse_update_period = new_inv
      if curvature_ema is not None:
        self._curvature_ema = float(curvature_ema)
    finally:
      self.lock_attributes()

  def _live_step_scalars(self) -> tuple[Array, Array]:
    """Current curvature_ema / inverse_update_period as traced step args.

    Plain rank-0 arrays, matching how callers pass learning_rate /
    momentum / damping into ``step``: the staging layer broadcasts them,
    so no per-device replication is required here (and
    ``device_put_replicated`` no longer exists on modern JAX anyway).
    """
    ema = jnp.asarray(self._curvature_ema, dtype=jnp.float32)
    period = jnp.asarray(self._inverse_update_period, dtype=jnp.int32)
    return ema, period

  @functools.partial(utils.staged, static_argnums=1)
  def _rng_split(self, rng: PRNGKey, num: int) -> tuple[Array, ...]:
    """Splits the ``rng`` key."""
    return tuple(jax.random.split(rng, num))

  @utils.auto_scope_method
  def _compute_loss_value(self, func_args: FuncArgsVariants) -> Array:
    """Computes the value of the loss function being optimized."""
    return self._value_func(*func_args)

  def _verify_args_and_get_step_counter(
      self,
      step_counter: Array,
      learning_rate: Array | None = None,
      momentum: Array | None = None,
      damping: Array | None = None,
      global_step_int: int | None = None,
  ) -> int:
    """Verifies that the arguments passed to the step function are correct."""

    # Verify correct arguments invocation
    if self._use_adaptive_learning_rate and learning_rate is not None:
      raise ValueError("When use_adaptive_learning_rate is set to True you "
                       "should not pass a value to the step function.")

    elif not self._use_adaptive_learning_rate and (
        self._learning_rate_schedule is None and learning_rate is None):
      raise ValueError("When `use_adaptive_learning_rate` is set to False and "
                       "`learning_rate_schedule` is None you must provide a "
                       "value to the step function.")

    elif self._learning_rate_schedule is not None and learning_rate is not None:
      raise ValueError("When you have passed a `learning_rate_schedule` you "
                       "should not pass a value to the step function.")

    if self._use_adaptive_momentum and momentum is not None:
      raise ValueError("When `use_adaptive_momentum` is set to True you "
                       "should not pass a value to the step function.")

    elif not self._use_adaptive_momentum and (
        self._momentum_schedule is None and momentum is None):
      raise ValueError("When `use_adaptive_momentum` is set to False and "
                       "`momentum_schedule` is None you must provide a value to"
                       " the step function.")

    elif self._momentum_schedule is not None and momentum is not None:
      raise ValueError("When you have passed a `momentum_schedule` you should "
                       "not pass a value to the step function.")

    if self._use_adaptive_damping and damping is not None:
      raise ValueError("When `use_adaptive_damping` is set to True you "
                       "should not pass a value to the step function.")

    elif not self._use_adaptive_damping and (
        self._damping_schedule is None and damping is None):
      raise ValueError("When `use_adaptive_damping` is set to False and "
                       "`damping_schedule` is None you must provide a value to "
                       "the step function.")

    elif self._damping_schedule is not None and damping is not None:
      raise ValueError("When you have passed a `damping_schedule` you should "
                       "not pass a value to the step function.")

    if global_step_int is None:
      return int(self.get_first(step_counter))

    return global_step_int

  @utils.staged
  def _setup_state_and_schedules(
      self,
      learning_rate: Array | None,
      momentum: Array | None,
      damping: Array | None,
      step_counter: Array,
      data_seen: Array,
  ) -> tuple[Numeric | None, Numeric | None, Numeric, Numeric]:
    """Helper function for setting up learning rate, momentum and damping."""

    # Compute schedules if applicable
    if self._learning_rate_schedule is not None:

      assert learning_rate is None
      learning_rate = utils.call_func_with_conditional_kwargs(
          self._learning_rate_schedule, step_counter, data_seen=data_seen)

    if self._momentum_schedule is not None:
      assert momentum is None
      momentum = utils.call_func_with_conditional_kwargs(
          self._momentum_schedule, step_counter, data_seen=data_seen)

    if self._damping_schedule is not None:
      assert damping is None
      damping = utils.call_func_with_conditional_kwargs(
          self._damping_schedule, step_counter, data_seen=data_seen)

    else:
      assert damping is not None

    if self._precon_damping_schedule is not None:
      precon_damping = utils.call_func_with_conditional_kwargs(
          self._precon_damping_schedule, step_counter, data_seen=data_seen)
    else:
      precon_damping = damping * self._precon_damping_mult

    return learning_rate, momentum, damping, precon_damping

  def _setup_func_args_and_rng(
      self,
      params: Params,
      rng: PRNGKey,
      batch: Batch,
      func_state: FuncState | None,
  ) -> tuple[FuncArgsVariants, Array]:
    """Helper function for setting up the model function arguments correctly."""

    # Preprocess the batch and construct correctly the function arguments
    batch = self._batch_process_func(batch)

    # Correctly split rng
    if self._value_func_has_rng:
      rng, func_rng = jax.random.split(rng)
    else:
      func_rng = None

    # Make the function args
    func_args = make_func_args(
        params=params,
        func_state=func_state,
        rng=func_rng,
        batch=batch,
        has_state=self._value_func_has_state,
        has_rng=self._value_func_has_rng,
    )

    return func_args, rng

  def _update_estimator_curvature(
      self,
      estimator_state: BlockDiagonalCurvature.State,
      func_args: FuncArgsVariants,
      rng: PRNGKey,
      ema_old: Numeric,
      ema_new: Numeric,
      precon_damping: Numeric,
      sync: Array | bool = True
  ) -> BlockDiagonalCurvature.State:
    """Updates the curvature estimator state."""

    state = self._estimator.update_curvature_matrix_estimate(
        state=estimator_state,
        ema_old=ema_old,
        ema_new=ema_new,
        identity_weight=self._l2_reg + precon_damping,
        # Note that the batch is always the last entry of FuncArgsVariantsdef
        batch_size=self._batch_size_extractor(func_args[-1]),
        rng=rng,
        func_args=func_args,
        pmap_axis_name=self.pmap_axis_name,
    )
    return jax.lax.cond(
        sync,
        functools.partial(self._estimator.sync,
                          pmap_axis_name=self.pmap_axis_name),
        lambda state_: state_,
        state,
    )

  def _update_estimator_curvature_and_value_and_grad(
      self,
      estimator_state: BlockDiagonalCurvature.State,
      func_args: FuncArgsVariants,
      rng: PRNGKey,
      ema_old: Numeric,
      ema_new: Numeric,
      precon_damping: Numeric,
      sync: Array | bool = True,
  ) -> tuple[BlockDiagonalCurvature.State, Array, Params]:
    """Updates exact-Fisher curvature and returns its shared loss/gradient."""

    state, loss, grads = (
        self._estimator.update_curvature_matrix_estimate_and_value_and_grad(
            state=estimator_state,
            ema_old=ema_old,
            ema_new=ema_new,
            identity_weight=self._l2_reg + precon_damping,
            batch_size=self._batch_size_extractor(func_args[-1]),
            rng=rng,
            func_args=func_args,
            pmap_axis_name=self.pmap_axis_name,
        )
    )
    state = jax.lax.cond(
        sync,
        functools.partial(
            self._estimator.sync,
            pmap_axis_name=self.pmap_axis_name,
        ),
        lambda state_: state_,
        state,
    )
    return state, loss, grads

  @utils.auto_scope_method
  def _compute_loss_and_grads(
      self,
      func_args: FuncArgsVariants,
      state: State | None = None,
  ) -> tuple[Array, Params, FuncState | None, FuncAux | None]:
    """Computes the model loss value and its gradients."""

    del state

    out, grads = self._value_and_grad_func(*func_args)

    loss, func_state, aux = extract_func_outputs(
        out, self._value_func_has_aux, self._value_func_has_state)

    if self._include_registered_loss_in_stats:
      aux = aux or {}
      aux["loss_registered"] = self._compute_loss_from_registrations(func_args)

    return loss, grads, func_state, aux

  @functools.partial(utils.staged, donate_argnums=0)
  def _maybe_update_inverse_cache(
      self,
      state: State,
      precon_damping: Array,
      inverse_update_period: Array,
  ) -> State:
    """Updates the estimator state cache if it is the right iteration."""

    # Copy this first since we mutate it later in this function.
    state = state.copy()

    state.estimator_state = lax.cond(
        self._should_update_inverse_cache(state, inverse_update_period),
        functools.partial(
            self._estimator.update_cache,
            identity_weight=self._l2_reg + precon_damping,
            exact_powers=self._exact_powers_to_cache,
            approx_powers=self._approx_powers_to_cache,
            eigenvalues=False,
            pmap_axis_name=self.pmap_axis_name,
        ),
        lambda state_: state_,
        state.estimator_state,
    )

    return state

  @functools.partial(utils.staged, static_argnums=3)
  def _compute_preconditioned_gradient(
      self,
      state: State,
      grads: Params,
      precon_damping: Array,
      can_distribute: bool = True,
  ) -> Params:
    """Computes the preconditioned gradient."""

    return self._estimator.multiply_matpower(
        state=state.estimator_state,
        parameter_structured_vector=grads,
        identity_weight=self._l2_reg + precon_damping,
        power=self._precon_power,
        exact_power=self._use_exact_inverses,
        use_cached=self._use_cached_inverses,
        pmap_axis_name=self.pmap_axis_name if can_distribute else None,
        norm_to_scale_identity_weight_per_block=self._norm_to_scale_identity_weight_per_block,
    )

  @utils.staged
  def _maybe_apply_norm_constraint(
      self, grads: Params, preconditioned_grads: Params, coefficient: Array
  ) -> tuple[Params, Params | None]:
    """Scales precon grad to have curvature-weighted norm <= norm_constraint."""

    if self._norm_constraint is None:
      return preconditioned_grads, None

    assert not self._use_adaptive_learning_rate

    sq_norm_grads = utils.inner_product(preconditioned_grads, grads)
    sq_norm_scaled_grads = sq_norm_grads * coefficient ** 2

    max_coefficient = jnp.sqrt(self._norm_constraint / sq_norm_scaled_grads)
    coefficient = jnp.minimum(max_coefficient, 1)

    precon_grad = utils.scalar_mul(preconditioned_grads, coefficient)

    return precon_grad, sq_norm_scaled_grads

  def _compute_quad_change_for_damping_adapt(
      self,
      state: State,
      delta: Params,
      grads: Params,
      damping: Array,
      func_args: FuncArgsVariants,
  ) -> Array:
    """The quadratic model change, when lr and momentum are non-adaptive."""

    assert not (self._use_adaptive_learning_rate or self._use_adaptive_momentum)

    if self._always_use_exact_qmodel_for_damping_adjustment:
      quad_model = self._compute_exact_quad_model_filtered(
          [delta], grads, func_args, state=state)
    else:
      quad_model = self._compute_approx_quad_model(state, [delta], grads)

    w = jnp.ones([])
    return self._solve_quad_model(quad_model, damping, [w])[1]

  def _coefficients_and_quad_change(
      self,
      state: State,
      vectors: Sequence[Params],
      grads: Params,
      learning_rate: Numeric | None,
      momentum: Numeric | None,
      damping: Numeric,
      func_args: FuncArgsVariants,
      should_update_damping: bool,
  ) -> tuple[tuple[Numeric, Numeric], Numeric]:
    """The correct update coefficients and corresponding quadratic change."""

    # Compute the coefficients of the update vectors
    # The learning rate is defined as the negative of the coefficient by which
    # we multiply the gradients, while the momentum is the coefficient by
    # which we multiply the velocities.
    neg_learning_rate = -learning_rate if learning_rate is not None else None
    fixed_coefficients = (neg_learning_rate, momentum)

    if self._use_adaptive_learning_rate or self._use_adaptive_momentum:

      assert fixed_coefficients[0] is None or fixed_coefficients[1] is None

      quad_model = self._compute_exact_quad_model_filtered(
          vectors, grads, func_args, state=state,
          fixed_coefficients=fixed_coefficients)

      return self._solve_quad_model(quad_model, damping, fixed_coefficients)

    else:

      assert all(c is not None for c in fixed_coefficients)
      fixed_coefficients: tuple[Numeric, Numeric]

      if should_update_damping:

        delta = self._weighted_sum_of_objects(vectors, fixed_coefficients)

        quad_change = self._compute_quad_change_for_damping_adapt(
            state, delta, grads, damping, func_args)

      else:
        quad_change = self._invalid_metric_value

      return fixed_coefficients, quad_change

  @utils.staged
  def _compute_loss_from_registrations(
      self,
      func_args: FuncArgsVariants
  ) -> Array:

    loss = self._estimator.compute_func_from_registered(
        func_args, self._batch_size_extractor(func_args[-1]))

    if self._l2_reg > 0.0:

      l2_reg_val = self._l2_reg / 2 * utils.squared_norm(
          func_args[self._params_index])

      loss += l2_reg_val

    return loss

  @utils.staged
  def _init(
      self,
      params: Params,
      rng: PRNGKey,
      batch: Batch,
      func_state: FuncState | None = None,
  ) -> State:
    """A staged function to initialize the optimizer state ."""

    # Note that we can reuse the rng in the func_args construction below, as
    # these are just dummy values used to perform the tracing.

    return Optimizer.State(
        velocities=jax.tree_util.tree_map(jnp.zeros_like, params),
        estimator_state=self._estimator.init(
            rng=rng,
            func_args=make_func_args(
                params=params,
                func_state=func_state,
                rng=rng,
                batch=self._batch_process_func(batch),
                has_state=self._value_func_has_state,
                has_rng=self._value_func_has_rng,
            ),
            exact_powers_to_cache=self._exact_powers_to_cache,
            approx_powers_to_cache=self._approx_powers_to_cache,
            cache_eigenvalues=False
        ),
        damping=jnp.array(
            (self._initial_damping if self._initial_damping is not None
             else -1e10), dtype=float),
        data_seen=jnp.array(0, dtype=int),
        step_counter=jnp.array(0, dtype=int)
    )

  def init(
      self,
      params: Params,
      rng: PRNGKey,
      batch: Batch,
      func_state: FuncState | None = None,
  ) -> State:
    """Initializes the optimizer and returns the appropriate optimizer state.

    NOTE: please do not jit/pmap or otherwise compile this function with JAX,
    as this can lead to errors. Compilation is handled internally by the
    optimizer.

    NOTE: when ``multi_device`` is ``True``, all of the JAX array arguments to
    this function (including arrays inside of trees), should have an extra
    leading axis the size of the number of local devices.

    Args:
      params: Example models parameters (used for tracing and shape info).
      rng: A Jax PRNG key. Unlike the ``rng`` in the step function, should be
        the same for each host and for each slice in the leading axis (i.e.
        corresponding to devices) when ``multi_device`` is ``True``.
      batch: An example batch of the same size as the one passed to ``step``
        (or returned from the ``data_iterator``). Used for tracing and shape
        info.
      func_state: Example function state (used for tracing and shape info).

    Returns:
      The initialized optimizer state.
    """

    if not self.finalized:
      self.finalize(params, rng, batch, func_state)

    # Check that mask_out_unregularized_params works as intended.
    _ = self._maybe_mask_out_unregularized_parameters(params, log_paths=True)

    return self._init(params, rng, batch, func_state)

  @functools.partial(utils.staged, donate_argnums=[1, 3, 5])  # pytype: disable=wrong-arg-types
  def _burnin(
      self,
      params: Params,
      state: State,
      rng: Array,
      batch: Batch,
      func_state: FuncState | None,
      damping: Array | None,
      accumulator: utils.MultiChunkAccumulator,
      sync: Array | bool,
  ) -> tuple[State, utils.MultiChunkAccumulator]:
    """A single burnin step, updating only the curvature estimate."""

    _, _, _, precon_damping = self._setup_state_and_schedules(
        None, None,
        state.damping if self._use_adaptive_damping else damping,
        state.step_counter, state.data_seen)

    # Copy this first since we mutate it later in this function.
    accumulator = accumulator.copy()

    func_args, rng = self._setup_func_args_and_rng(
        params, rng, batch, func_state)

    # Update curvature estimate
    state.estimator_state = self._update_estimator_curvature(
        state.estimator_state,
        func_args,
        rng,
        ema_old=1.0,
        ema_new=1.0,
        precon_damping=precon_damping,
        sync=sync,
    )

    # Optionally update func_state
    if func_state is not None:
      out, _ = self._value_and_grad_func(*func_args)
      _, func_state, _ = extract_func_outputs(
          out, self._value_func_has_aux, self._value_func_has_state)

    accumulator.add(func_state)

    return state, accumulator

  def _burnin_phase(
      self,
      num_steps: int,
      params: Params,
      state: State,
      rng: PRNGKey,
      data_iterator: Iterator[Batch],
      func_state: FuncState | None = None,
      damping: Array | None = None,
  ) -> tuple[State, FuncState | None]:
    """Runs all burnin steps required."""

    if num_steps > 0:

      rng = self._rng_split(rng, num_steps)

      accumulator = utils.MultiChunkAccumulator.zeros_like(
          func_state, self.multi_device)

      for i, rng_i in enumerate(rng):
        batch = next(data_iterator)

        state, accumulator = self._burnin(
            params, state, rng_i, batch, func_state, damping, accumulator,
            i == num_steps - 1)

      func_state = accumulator.value_and_clear()

    return state, func_state

  @functools.partial(
      utils.staged, donate_argnums=(0, 1, 4), static_argnums=(8, 9))
  @utils.auto_scope_method
  def _step(
      self,
      params: Params,
      state: State,
      rng: Array,
      batch: Batch,
      func_state: FuncState | None,
      learning_rate: Array | None,
      momentum: Array | None,
      damping: Array | None,
      should_update_estimate_curvature: bool,
      should_update_damping: bool,
      curvature_ema: Numeric,
      inverse_update_period: Numeric,
  )-> ReturnEither:
    """A single full step of the optimizer."""

    # Copy this first since we mutate it later in this function.
    state = state.copy()

    # Setup arguments
    (learning_rate, momentum, damping,
     precon_damping) = self._setup_state_and_schedules(
         learning_rate, momentum,
         state.damping if self._use_adaptive_damping else damping,
         state.step_counter, state.data_seen)

    func_args, rng = self._setup_func_args_and_rng(
        params, rng, batch, func_state)

    # Update curvature estimate
    if should_update_estimate_curvature:
      if self._share_curvature_and_grad_forward:
        (
            state.estimator_state,
            loss,
            grads,
        ) = self._update_estimator_curvature_and_value_and_grad(
            state.estimator_state,
            func_args,
            rng,
            ema_old=curvature_ema,
            ema_new=1.0,
            precon_damping=precon_damping,
            sync=self._should_sync_estimator(state, inverse_update_period),
        )
      else:
        state.estimator_state = self._update_estimator_curvature(
            state.estimator_state,
            func_args,
            rng,
            ema_old=curvature_ema,
            ema_new=1.0,
            precon_damping=precon_damping,
            sync=self._should_sync_estimator(state, inverse_update_period),
        )

    del rng  # should not be used after this point!

    # Compute loss and gradients
    if (
        should_update_estimate_curvature
        and self._share_curvature_and_grad_forward
    ):
      func_state = None
      aux = None
    else:
      loss, grads, func_state, aux = self._compute_loss_and_grads(
          func_args, state=state)

    # Sync
    loss, grads = utils.pmean_if_pmap((loss, grads), self.pmap_axis_name)

    # Update the inverse curvature
    state = self._maybe_update_inverse_cache(
        state, precon_damping, inverse_update_period)

    # Compute proposed directions
    preconditioned_gradient = self._compute_preconditioned_gradient(
        state, grads, precon_damping
    )

    # Client stats hook on the PRE-norm-constraint preconditioned
    # gradient (the clip below is a scalar rescale).
    if self._step_stats_hook is not None:
      hook_stats = self._step_stats_hook(
          self._estimator, grads, preconditioned_gradient)
    else:
      hook_stats = {}

    # constrain the norms
    preconditioned_gradient, scaled_grad_norm_sq = (
        self._maybe_apply_norm_constraint(
            grads, preconditioned_gradient, learning_rate,
        )
    )

    vectors = (preconditioned_gradient, state.velocities)

    # Compute the coefficients for the vectors
    coefficients, quad_model_change = self._coefficients_and_quad_change(
        state=state,
        vectors=vectors,
        grads=grads,
        learning_rate=learning_rate,
        momentum=momentum,
        damping=damping,
        func_args=func_args,
        should_update_damping=should_update_damping,
        )

    # Compute the parameter update (delta)
    delta = self._weighted_sum_of_objects(vectors, coefficients)

    # Update parameters
    new_params = jax.tree_util.tree_map(jnp.add, params, delta)

    if should_update_damping or self._use_step_rejection:

      new_loss = self._compute_loss_value((new_params,) + func_args[1:])
      # Sync
      new_loss = utils.pmean_if_pmap(new_loss, self.pmap_axis_name)

    else:
      new_loss = self._invalid_metric_value

    # Optionally compute the reduction ratio and update the damping
    if should_update_damping:

      state.damping, rho = self._compute_new_damping_and_rho(
          loss, new_loss, quad_model_change, state.damping)

    else:
      # If not adjusting the damping we don't compute these here and just set
      # them to self._invalid_metric_value.
      new_loss, rho = self._invalid_metric_value, self._invalid_metric_value

    if self._use_step_rejection:

      reject_step = jnp.logical_or(jnp.isnan(new_loss), new_loss > loss)

      params, state.velocities, state.damping = lax.cond(
          reject_step,
          lambda: (params, state.velocities,
                   self._reject_damping_increase_factor * state.damping),
          lambda: (new_params, delta, state.damping))

    else:
      # stop the linter from complaining about uninitialized variable
      reject_step = False
      params, state.velocities = new_params, delta

    # Compute per-device and total batch size
    batch_size = self._batch_size_extractor(func_args[-1])

    if self.multi_device:
      total_batch_size = batch_size * jax.device_count()
    else:
      total_batch_size = batch_size

    # Update data seen and step counter
    state.data_seen = state.data_seen + total_batch_size
    state.step_counter = state.step_counter + 1

    # Statistics with useful information
    # Unlike other norm stats, sq_norm_scaled_grads has to be computed if
    # norm_constraint is not None, so log it by default even if the other
    # norm stats are not logged. This reduces the overall computational cost if
    # no other grad stats are desired.
    stats = dict(
        step=state.step_counter,
        batch_size=jnp.asarray(total_batch_size, dtype=jnp.int32),
        data_seen=state.data_seen,
        loss=loss,
        new_loss=new_loss,
        learning_rate=-coefficients[0],
        momentum=coefficients[1],
        damping=damping,
        precon_damping=precon_damping,
        rho=rho,
        quad_model_change=quad_model_change,
        scaled_grad_norm_sq=scaled_grad_norm_sq,
    )

    if self._use_step_rejection:
      stats["step_rejected"] = reject_step

    stats.update(hook_stats)

    if aux is not None:
      aux = utils.pmean_if_pmap(aux, self.pmap_axis_name)
      stats["aux"] = aux

    if self._include_norms_in_stats:
      stats["param_norm"] = utils.norm(params)
      stats["grad_norm"] = utils.norm(grads)
      stats["precon_grad_norm"] = utils.norm(preconditioned_gradient)
      stats["update_norm"] = utils.norm(delta)

    if self._include_per_param_norms_in_stats:
      stats.update(utils.per_parameter_norm(params, "param_norm"))
      stats.update(utils.per_parameter_norm(grads, "grad_norm"))
      stats.update(
          utils.per_parameter_norm(preconditioned_gradient, "precon_grad_norm")
      )
      stats.update(utils.per_parameter_norm(delta, "update_norm"))

    if self._include_registered_loss_in_stats:
      assert aux is not None
      stats["loss_registered"] = aux.pop("loss_registered")
      stats["loss_registered"] = utils.pmean_if_pmap(stats["loss_registered"],
                                                     self.pmap_axis_name)
      stats["loss_registered_reldiff"] = (
          stats["loss_registered"] - loss) / loss

    if self._value_func_has_state:
      return params, state, func_state, stats

    assert func_state is None

    return params, state, stats

  def step(
      self,
      params: Params,
      state: State,
      rng: PRNGKey,
      data_iterator: Iterator[Batch] | None = None,
      batch: Batch | None = None,
      func_state: FuncState | None = None,
      learning_rate: Array | None = None,
      momentum: Array | None = None,
      damping: Array | None = None,
      global_step_int: int | None = None
  )-> ReturnEither:
    """Performs a single update step using the optimizer.

    NOTE: please do not jit/pmap or otherwise compile this function with JAX,
    as this can lead to errors. Compilation is handled internally by the
    optimizer.

    NOTE: when ``multi_device`` is ``True``, all of the JAX array arguments to
    this function (including arrays inside of trees), should have an extra
    leading axis the size of the number of local devices. Slices of ``batch``
    and ``rng`` should be different for each device, whereas the other arugments
    should be identical for each slice. Passing the arguments any other way will
    result in an exception, or possibly undefined behavior.

    Args:
      params: The current parameters of the model.
      state: The current state of the optimizer.
      rng: A Jax PRNG key. Should be different for each iteration, each host,
        and for each slice in the leading axis (i.e. corresponding to devices)
        when ``multi_device`` is ``True``.
      data_iterator: A data iterator to use (if not passing ``batch``).
      batch: A single batch used to compute the update. Should only pass one
        of ``data_iterator`` or ``batch``.
      func_state: Any function state that gets passed in and returned.
      learning_rate: Learning rate to use if the optimizer was created with
        ``use_adaptive_learning_rate=False`` and
        ``learning_rate_schedule=None``. Should be ``None`` otherwise.
      momentum: Momentum to use if the optimizer was created with
        ``use_adaptive_momentum=False`` and ``momentum_schedule=None``. Should
        be ``None`` otherwise.
      damping: Damping to use if the optimizer was created with
        ``use_adaptive_damping=False`` and ``damping_schedule=None``. Should be
        ``None`` otherwise. See discussion of constructor argument
        ``initial_damping`` for more information about damping.
      global_step_int: The global step as a python int. Note that this must
        match the step internal to the optimizer that is part of its state.

    Returns:
      (params, state, stats) if ``value_func_has_state=False`` and
      (params, state, func_state, stats) otherwise, where

          * params is the updated model parameters.

          * state is the updated optimizer state.

          * func_state is the updated function state.

          * stats is a dictionary of useful statistics including the loss.
    """

    if (data_iterator is None) == (batch is None):
      raise ValueError("Exactly one of the arguments ``data_iterator`` and "
                       "``batch`` must be provided.")

    step_counter_int = self._verify_args_and_get_step_counter(
        step_counter=state.step_counter,
        learning_rate=learning_rate,
        momentum=momentum,
        damping=damping,
        global_step_int=global_step_int,
    )

    if step_counter_int == 0:

      if self._num_burnin_steps > 0:

        if data_iterator is None:
          raise ValueError("If num_burnin_steps > 0, data_iterator must be "
                           "provided.")

        rng, burnin_rng = self._rng_split(rng, 2)

        state, func_state = self._burnin_phase(
            num_steps=self._num_burnin_steps,
            params=params,
            state=state,
            rng=burnin_rng,
            data_iterator=data_iterator,
            func_state=func_state,
            damping=damping,
        )

    if data_iterator is not None:
      batch = next(data_iterator)

    if (step_counter_int == 0 and self._use_adaptive_damping
        and self._use_initial_damping_calibration):

      assert self._num_burnin_steps > 0

      state = self._calibrate_initial_damping(
          params, state, rng, batch, func_state, learning_rate, momentum)

    should_update_estimate_curvature = self._should_update_estimate_curvature(
        step_counter_int
    )
    should_update_damping = self._should_update_damping(step_counter_int)

    curvature_ema, inverse_update_period = self._live_step_scalars()

    return self._step(
        params, state, rng, batch, func_state, learning_rate, momentum, damping,
        should_update_estimate_curvature, should_update_damping,
        curvature_ema, inverse_update_period)

  def _calibrate_initial_damping(
      self,
      params: Params,
      state: State,
      rng: PRNGKey,
      batch: Batch,
      func_state: FuncState | None = None,
      learning_rate: Array | None = None,
      momentum: Array | None = None,
  ) -> State:
    """Calibrates the initial damping parameter."""

    # Instead of writing a custom compiled function to compute rho and update
    # the damping, we're going to be lazy and just call the step function
    # repeatedly, throwing out the new optimizer state, params, and stats, while
    # keeping the rng and batch the same at each call. This is a bit hacky and
    # somewhat wasteful, both in terms of a few extra (minor) computations done
    # in step() that are pointless, as well as the extra memory required to
    # store temporary copies of the optimizer state and model params.

    # TODO(jamesmartens): Improve the implementation if this feature is commonly
    # used?

    while True:

      prev_damping = float(self.get_first(state.damping))

      # Note that we need to copy params and func_state since _step() will
      # donate them. A bette option might be to recompile _step() to not donate
      # these arguments.
      curvature_ema, inverse_update_period = self._live_step_scalars()

      ret = self._step(
          self.copy_obj(params), self.copy_obj(state), rng, batch,
          self.copy_obj(func_state), learning_rate, momentum, None, False, True,
          curvature_ema, inverse_update_period)

      new_state = ret[1]

      new_damping = float(self.get_first(new_state.damping))
      state.damping = new_state.damping

      del new_state

      if prev_damping == new_damping:
        return state

  @utils.auto_scope_method
  def _compute_exact_quad_model_filtered(
      self,
      vectors: Sequence[Params],
      grads: Params,
      func_args: FuncArgsVariants,
      state: State | None = None,
      fixed_coefficients: Sequence[Numeric | None] | None = None,
      **kwargs,
  ) -> QuadModelParams:
    """Computes the components of the exact quadratic model."""

    # We check the fixed_coefficients for zeros to save computing the expensive
    # matrix vector products for vectors that will eventually be multiplied by
    # zero. If fixed_coefficients is None, we assume that all coefficients are
    # free and compute the full model.

    if fixed_coefficients is None:  # can we get rid of this?
      return self._compute_exact_quad_model(
          vectors, grads, func_args, state=state, **kwargs)

    assert len(vectors) == len(fixed_coefficients)
    assert len(vectors) == 2  # only deal with the two vector case

    def if_momentum_coeff_zero():

      # Only pass in the vectors that won't be multiplied by zero
      quad_model = self._compute_exact_quad_model(
          vectors[:1], grads, func_args, state=state, **kwargs)

      # Repad the quad model with zeroes for the removed entries
      return tuple(
          jnp.pad(arr, [(0, 1)] * arr.ndim, constant_values=0.0)
          for arr in quad_model
      )

    # This saves compiling both branches in the static case
    if (isinstance(fixed_coefficients[1], float)
        and fixed_coefficients[1] == 0.0):

      return if_momentum_coeff_zero()

    # Due to how XLA cannot share computations across cond boundaries, such as
    # network forward and backwards passes, we cannot use a cond here and remain
    # efficient. If this behavior ever changes we can uncomment the block below.

    # return jax.lax.cond(
    #     fixed_coefficients[1] == 0.0,
    #     if_momentum_coeff_zero,
    #     lambda: self._compute_exact_quad_model(
    #         vectors, grads, func_args, state=state),
    # )

    return self._compute_exact_quad_model(
        vectors, grads, func_args, state=state, **kwargs)

  def _maybe_mask_out_unregularized_parameters(
      self, params: Params, log_paths: bool = False) -> Params:
    """Mask out parameters that are not l2 regularized."""

    if log_paths:
      logging.info("Unregularized parameters masking info (for curvature "
                   "calculations and L2 regularization)")

    def maybe_mask_out_single_param(
        path: tuple[Any, ...],
        param: Array
    ) -> Array:
      """Zero out a single parameter."""
      str_path = []
      for p in path:
        if isinstance(p, jax.tree_util.DictKey):
          str_path.append(p.key)
        elif isinstance(p, jax.tree_util.GetAttrKey):
          str_path.append(p.name)

      should_mask = any(
          p in str_path
          for p in self._regularized_parameters_path_exclusions
      )

      if log_paths:
        log_message = "Masking" if should_mask else "Not masking"
        logging.info("  %s out %s", log_message, path)

      return jnp.zeros_like(param) if should_mask else param

    return jax.tree.map_with_path(
        maybe_mask_out_single_param, params
    )

  @utils.auto_scope_method
  def _compute_exact_quad_model(
      self,
      vectors: Sequence[Params],
      grads: Params,
      func_args: FuncArgsVariants,
      state: State | None = None,
  ) -> QuadModelParams:
    """Computes the components of the exact quadratic model.

    See comments of QuadModelParams for a description of the returned tuple.

    Args:
      vectors: sequence of update vectors `V`.
      grads: The gradient `g` of the loss function.
      func_args: The arguments to the model's value function.
      state: The current optimizer state.

    Returns:
      A `QuadModelParams` tuple (A, D, R, b).
    """

    del state

    if self._mat_type_for_exact_quad_model == "fisher":
      c_factor_v = tuple(self._implicit.multiply_fisher_factor_transpose
                         (func_args, vi) for vi in vectors)
    elif self._mat_type_for_exact_quad_model == "ggn":
      c_factor_v = tuple(self._implicit.multiply_ggn_factor_transpose
                         (func_args, vi) for vi in vectors)
    else:
      raise ValueError(f"Unrecognized matrix type string for exact quad model:"
                       f"'{self._mat_type_for_exact_quad_model}'.")

    masked_vectors = tuple(self._maybe_mask_out_unregularized_parameters(vi)
                           for vi in vectors)

    # pylint: disable=invalid-name
    A = utils.matrix_of_inner_products(c_factor_v)
    D = utils.matrix_of_inner_products(vectors)
    R = utils.matrix_of_inner_products(masked_vectors)
    b = utils.vector_of_inner_products(grads, vectors)
    # pylint: enable=invalid-name

    quad_model_params = (A, D, R, b)

    return utils.pmean_if_pmap(quad_model_params, self.pmap_axis_name)

  @functools.partial(utils.staged, donate_argnums=2)
  @utils.auto_scope_method
  def _compute_approx_quad_model(
      self,
      state: State,
      vectors: Sequence[Params],
      grads: Params,
  ) -> QuadModelParams:
    """Computes the components of the approximate quadratic model."""

    # v_i^T C v_j
    def c_times_v(v):
      return self._estimator.multiply(
          state=state.estimator_state,
          parameter_structured_vector=v,
          identity_weight=0.0,
          exact_power=True,
          use_cached=False,
          pmap_axis_name=self.pmap_axis_name,
          norm_to_scale_identity_weight_per_block=self._norm_to_scale_identity_weight_per_block,
      )

    c_vectors = [c_times_v(v_i) for v_i in vectors]

    return (utils.symmetric_matrix_inner_products(c_vectors, vectors),
            utils.matrix_of_inner_products(vectors),
            utils.matrix_of_inner_products(vectors),
            utils.vector_of_inner_products(grads, vectors))

  def _evaluate_quadratic_model(
      self,
      a: Array,
      a_damped: Array,
      b: Array,
      w: Array,
  ) -> Array:
    """Computes the quadratic model value from the inputs provided."""

    a_final = a_damped if self._include_damping_in_quad_change else a

    return jnp.dot(w, jnp.dot(a_final, w)) / 2 + jnp.dot(w, b)

  @utils.staged
  def _solve_quad_model(
      self,
      quad_model_parameters: QuadModelParams,
      damping: Array,
      fixed_coefficients: Sequence[Numeric | None],
      reg_coeff: Numeric | None = None,
  ) -> tuple[tuple[Numeric, ...], Array]:
    """Solves for the optimal learning rate and momentum of the quadratic model.

    Args:
      quad_model_parameters: The computed matrices A, D, R, and vector b.
      damping: The damping to use for evaluating the quadratic model.
      fixed_coefficients: A list over the vectors of the fixed numerical values
        to use for their coefficients. For each of these that is None, the
        quadratic model is minimized to compute the 'optimal' coefficient value.
      reg_coeff: The L2 regularization parameter to use. If None, the default
        value from the optimizer is used.

    Returns:
     A tuple of coefficients which are the solution (and include any values that
     are not None from fixed_weights), and the value of the quadratic model
     function for this solution (as a scalar).

    Raises:
      The function currently supports only up to two vectors, hence if you
      provide more, it will raise a ``NotImplementedError``.
    """

    if reg_coeff is None:
      # use default l2 regularisation value.
      reg_coeff = self._l2_reg

    # pylint: disable=invalid-name
    A_no_diag, D, R, b = quad_model_parameters
    A = A_no_diag + reg_coeff * R
    A_damped = A + damping * D

    if all(c is None for c in fixed_coefficients):
      # Adapt all coefficients

      if len(fixed_coefficients) == 1:
        # This special case arises at the first iteration, because all
        # velocities are zeros.
        special_case = jnp.logical_and(A_damped[0, 0] == 0, b[0] == 0)
        w = -lax.cond(special_case, lambda: b, lambda: b / A_damped[0])

      elif len(fixed_coefficients) == 2:
        w = -utils.psd_solve_maybe_zero_last_idx(A_damped, b)

      else:
        raise NotImplementedError()

    elif all(c is not None for c in fixed_coefficients):
      # No coefficients adapted

      w = jnp.asarray(fixed_coefficients)

    elif len(fixed_coefficients) == 2:
      # Exactly one adapted coefficient

      w = [None, None]
      index = fixed_coefficients.index(None)
      w[1 - index] = fixed_coefficients[1 - index]

      b_extra = A_damped[1 - index, index] * w[1 - index]
      # pylint: enable=invalid-name

      w[index] = -(b[index] + b_extra) / A_damped[index, index]

    else:
      raise NotImplementedError()

    w = tuple(w)
    w: tuple[Numeric, ...]

    quad_model_change = self._evaluate_quadratic_model(
        A, A_damped, b, jnp.array(w))

    return w, quad_model_change

  @utils.staged
  def _compute_new_damping_and_rho(
      self,
      old_loss: Array,
      new_loss: Array,
      quad_change: Array,
      current_damping: Array,
  ) -> tuple[Array, Array]:
    """Computes the reduction ratio and the updated value of the damping."""

    # Reduction ratio
    rho = (new_loss - old_loss) / quad_change
    rho_not_nan = jnp.nan_to_num(rho, nan=-100.0)

    # Update damping
    should_increase = rho_not_nan < self._damping_lower_threshold
    increased_damping = current_damping / self._damping_decay_factor
    should_decrease = rho_not_nan > self._damping_upper_threshold
    decreased_damping = current_damping * self._damping_decay_factor

    damping = jnp.select([should_decrease, should_increase],
                         [decreased_damping, increased_damping],
                         default=current_damping)

    return jnp.clip(damping, self._min_damping, self._max_damping), rho

  @utils.staged
  def _weighted_sum_of_objects(
      self,
      objects: Sequence[utils.PyTree],
      coefficients: Sequence[Numeric],
  ) -> utils.PyTree:
    """Returns the weighted sum of the objects in the sequence."""
    return utils.weighted_sum_of_objects(objects, coefficients)


def convert_value_and_grad_to_value_func(
    value_and_grad_func: ValueAndGradFunc,
    has_aux: bool = False,
) -> ValueFunc:
  """Converts a value_and_grad function to value_func only.

  Args:
    value_and_grad_func: The function which computes the loss value and the
      gradients w.r.t. parameters.
    has_aux: Similar to the meaning in :func:`jax.grad`, whether the
      ``value_and_grad_func`` returns with the loss value any auxiliary data.

  Returns:
    A function that returns only the loss value.
  """

  def value_func(*args, **kwargs) -> Array:
    out, _ = value_and_grad_func(*args, **kwargs)
    return out[0] if has_aux else out

  return value_func


def convert_value_and_grad_to_clean_value_and_grad(
    value_and_grad_func: ValueAndGradFunc,
    has_aux: bool = False,
) -> utils.ValueAndGradFunc:
  """Converts a value_and_grad function to return only (loss, grads).

  Args:
    value_and_grad_func: The function which computes the loss value and the
      gradients w.r.t. parameters.
    has_aux: Similar to the meaning in :func:`jax.grad`, whether the
      ``value_and_grad_func`` returns with the loss value any auxiliary data.

  Returns:
    A function that returns `(loss, grads)`.
  """

  def clean_value_and_grad_func(*args, **kwargs) -> tuple[Array, Params]:
    out, grads = value_and_grad_func(*args, **kwargs)
    loss = out[0] if has_aux else out
    return loss, grads

  return clean_value_and_grad_func


def make_func_args(
    params: Params,
    func_state: FuncState | None,
    rng: PRNGKey | None,
    batch: Batch,
    has_state: bool,
    has_rng: bool,
) -> FuncArgsVariants:
  """Constructs the arguments to the model function in the pre-assumed order.

  The model function is assumed to take arguments in the following order:
    params, func_state, rng, batch
  If it has no function state or does not use an rng, those two arguments are
  discarded.

  Args:
    params: The model parameters.
    func_state: The function state, if ``has_state`` is ``True``, ``None``
      otherwise.
    rng: The PRNG, if ``has_rng`` is ``True``, ``None`` otherwise.
    batch: The batch of data.
    has_state: Whether the function has a function state.
    has_rng: Whether the function uses an rng.

  Returns:
    The arguments that need to be passed to the model function.
  """
  if has_state and func_state is None:
    raise ValueError("`func_state=None`, but argument `has_state=True`.")

  if has_rng and rng is None:
    raise ValueError("`rng=None`, but argument `has_rng=True`.")

  if not has_state and not has_rng:
    return params, batch

  elif not has_rng:
    return params, func_state, batch

  elif not has_state:
    return params, rng, batch

  else:
    return params, func_state, rng, batch


def extract_func_outputs(
    raw_outputs: FuncOutputs,
    has_aux: bool,
    has_state: bool,
) -> tuple[Array, FuncState | None, FuncAux | None]:
  """Converts the raw output of the model function into loss,func_state and aux.

  Args:
    raw_outputs: The direct output of the model function.
    has_aux: Whether the model function returns also some auxiliary data.
    has_state: Whether the model function has a function state.

  Returns:
    A triple ``(loss, func_state, aux)``. If the model function does not return
    any auxiliary data than ``aux`` will be ``None`` and if it does not have a
    state ``func_state`` will be ``None``.
  """

  if not has_aux and not has_state:
    assert isinstance(raw_outputs, Array)
    return raw_outputs, None, None

  loss, other = raw_outputs

  if has_aux and has_state:
    func_state, aux = other
  elif has_aux:
    func_state, aux = None, other
  else:
    func_state, aux = other, None

  return loss, func_state, aux