<commit_before># Lint as: python3
# Copyright 2021 The TensorFlow Authors. 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.
# ==============================================================================
"""Module for all optimizers."""

import dataclasses
import functools
from typing import Any, Callable, NamedTuple, Optional, Sequence, Tuple, Union

from absl import logging
import jax
from jax import numpy as jnp
from lingvo.jax import gshard_utils
from lingvo.jax import py_utils
from lingvo.jax import pytypes
import optax

NestedMap = py_utils.NestedMap
InstantiableParams = py_utils.InstantiableParams
JTensor = pytypes.JTensor
NestedJTensor = pytypes.NestedJTensor
NestedParams = pytypes.NestedParams

# Initializes sharding spec for the optimizer state variables.
TransformInitPartitionSpecFn = Callable[[NestedParams],
                                        Union[NestedParams,
                                              Sequence[NestedParams]]]


# Extension of optax.GradientTransformation that supports spmd sharding and
# explicit annotation of sharding specs for the optimizer state variables.
@dataclasses.dataclass(frozen=True)
class ShardedGradientTransformation:
  """GradientTransformation that supports spmd."""
  # init_fn and update_fn are the same as in optax.GradientTransformation
  init: optax.TransformInitFn
  update: optax.TransformUpdateFn
  # Input is the sharding specs of the variables used in the forward
  # computation.  Output is the sharding specs of the optimizer state variables.
  #
  # Constraints: output from this function should be of identical structure as
  # that of the init() function.
  init_partition_spec: TransformInitPartitionSpecFn


def count_init_fn(_):
  """Common init_fn that initializes a count for global step."""
  return NestedMap(count=jnp.zeros([], jnp.int32))


def count_init_partition_spec_fn(var_params):
  """Init partition spec for only partitioning the count/step."""
  var_spec_flattened, _ = jax.tree_flatten(var_params)
  assert var_spec_flattened
  first_var = var_spec_flattened[0]
  assert isinstance(first_var, py_utils.Params)
  device_mesh = first_var.device_mesh
  return NestedMap(
      count=py_utils.weight_params(
          shape=[],
          init=None,
          dtype=jnp.int32,
          collections=None,
          device_mesh=device_mesh,
          tensor_split_dims_mapping=[]))


def sharded_sgd(learning_rate_fn: optax.Schedule, momentum: Optional[float],
                nesterov: bool) -> ShardedGradientTransformation:
  """A canonical Stochastic Gradient Descent optimiser that supports spmd ...

  sharding.

  This implements stochastic gradient descent. It also includes support for
  momentum, and nesterov acceleration, as these are standard practice when
  using stochastic gradient descent to train deep neural networks.

  References:
    Sutskever et al, 2013: http://proceedings.mlr.press/v28/sutskever13.pdf

  Args:
    learning_rate_fn: a callable that given the current training step, returns
      the learning rate to apply.
    momentum: (default `None`), the `decay` rate used by the momentum term, when
      it is set to `None`, then momentum is not used at all.
    nesterov (default `False`): whether nesterov momentum is used.

  Returns:
    A `ShardedGradientTransformation`.
  """
  # TODO(yonghui): support momentum.
  assert momentum is None
  del nesterov

  def update_fn(updates, state, params=None):
    del params
    step_size = -1.0 * learning_rate_fn(state.count)
    updates = jax.tree_map(lambda g: jnp.array(step_size, dtype=g.dtype) * g,
                           updates)
    updated_states = NestedMap(count=state.count + jnp.array(1, jnp.int32))
    return updates, updated_states

  return ShardedGradientTransformation(
      init=count_init_fn,
      update=update_fn,
      init_partition_spec=count_init_partition_spec_fn)


class _AdamOptState:

  def __init__(self, *, m, v):
    self.m = m
    self.v = v


class _ShardedAdamHelper:
  """A helper class facilitates the creation of sharded_adam_optimizer."""

  def opt_state_sharding_spec(self,
                              var_params: py_utils.Params) -> _AdamOptState:
    """Returns optimizer sharding spec for one particular variable."""
    m_var_params = var_params.Copy()
    m_var_params.init = None
    v_var_params = var_params.Copy()
    v_var_params.init = None
    # m and v simply share the same sharding.
    return _AdamOptState(m=m_var_params, v=v_var_params)

  def init_opt_state(self, var_params: py_utils.Params) -> _AdamOptState:
    """Returns optimizer state for one particular variable."""
    return _AdamOptState(
        m=jnp.zeros_like(var_params), v=jnp.zeros_like(var_params))

  def sanitize_values(self, array: JTensor, replacement: float = 0.0):
    """Sanitizes NaN and Infinity values."""
    return jnp.nan_to_num(
        array, nan=replacement, posinf=replacement, neginf=replacement)

  def bias_corrected_decay(self, step: JTensor, decay: float) -> JTensor:
    """Incorporates bias correction into decay.

    Please see section 7.1 in https://arxiv.org/pdf/1804.04235.pdf for the
    derivation of the formulas below. With bias-corrected decay, we can simply
    do

    m_{t} = decay1 * m_{t-1} + (1 - decay1) * g
    v_{t} = decay2 * v_{t-1} + (1 - decay2) * g ^ 2

    without further bias correction.

    Args:
      step: current step, 0-based.
      decay: the raw decay. As t -> infinity, bias corrected decay converges to
        this value.

    Returns:
      Bias corrected decay.
    """
    t = step.astype(jnp.float32) + 1.
    return decay * (1. - jnp.power(decay, t - 1.)) / (1. - jnp.power(decay, t))

  def update_moments(self, step: JTensor, update: JTensor,
                     moments: _AdamOptState, beta1: float,
                     beta2: float) -> _AdamOptState:
    """Updates momentum values."""
    beta1_decay = self.bias_corrected_decay(step, beta1)
    beta2_decay = self.bias_corrected_decay(step, beta2)
    m = (1.0 - beta1_decay) * update + beta1_decay * moments.m
    v = (1.0 - beta2_decay) * (update**2) + beta2_decay * moments.v
    return _AdamOptState(m=m, v=v)

  def clip_update(self, update: JTensor, clip_threshold: float) -> JTensor:
    mean_update = self.sanitize_values(reduce_rms(update), 1.0)
    clip_threshold = jnp.array(clip_threshold, dtype=update.dtype)
    denom = jnp.maximum(1.0, mean_update / clip_threshold)
    return update / denom


def sharded_chain(
    *args: Union[optax.GradientTransformation, ShardedGradientTransformation]
) -> ShardedGradientTransformation:
  """Applies a list of (possibly sharded) chainable update transformations.

  Given a sequence of chainable transforms, `sharded_chain` returns an `init_fn`
  that constructs a `state` by concatenating the states of the individual
  transforms, and returns an `update_fn` which chains the update transformations
  feeding the appropriate state to each. In addition, it differs from the optax
  `chain` function, by also supporting ShardedGradientTransformation by chaining
  also the `init_partition_spec_fn`. If there are no
  ShardedGradientTransformations in the chain, the sharding specs will be
  None, meaning all the variables are replicated.

  Args:
    *args: a sequence of chainable GradientTransformations or
      ShardedGradientTransformations or a combination of both.

  Returns:
    A single chained ShardedGradientTransformation.
  """

  def init_fn(params):
    return tuple(fn.init(params) for fn in args)

  def update_fn(updates, state, params=None):
    if len(args) != len(state):
      raise ValueError('The number of updates and states has to be the same in '
                       'sharded chain.')

    new_state = []
    for s, fn in zip(state, args):
      updates, new_s = fn.update(updates, s, params)
      new_state.append(new_s)
    return updates, tuple(new_state)

  def init_partition_spec_fn(mdl_vars):
    partition_specs = []
    for fn in args:
      if isinstance(fn, ShardedGradientTransformation):
        nmap = fn.init_partition_spec(mdl_vars)
        partition_specs.append(nmap)
      else:
        # Replicate the states.
        partition_specs.append(None)
    return tuple(partition_specs)

  return ShardedGradientTransformation(
      init=init_fn,
      update=update_fn,
      init_partition_spec=init_partition_spec_fn)


def apply_l2_weight_decay(
    learning_rate_fn: optax.Schedule,
    l2_regularizer_weight: Optional[float] = 0.
) -> ShardedGradientTransformation:
  """Applies L2 weight decay.

  Args:
    learning_rate_fn: An optax schedule that infers the lr given the step.
    l2_regularizer_weight: Weight for L2 regularization.

  Returns:
    A ShardedGradientTransformation applying L2 weight decay.
  """

  def update_fn(updates, state, params):
    count = state.count
    lr_multiplier = learning_rate_fn(count)
    if l2_regularizer_weight:
      if params is None:
        raise ValueError('Params must not be empty when applying weight decay.')
      updates = jax.tree_multimap(
          lambda g, p: g - lr_multiplier * l2_regularizer_weight * p, updates,
          params)
    updated_state = NestedMap(count=count + 1)
    return updates, updated_state

  return ShardedGradientTransformation(
      init=count_init_fn,
      update=update_fn,
      init_partition_spec=count_init_partition_spec_fn)


def apply_l1_weight_decay(
    learning_rate_fn: optax.Schedule,
    l1_regularizer_weight: Optional[float] = 0.
) -> ShardedGradientTransformation:
  """Applies L1 weight decay.

  Args:
    learning_rate_fn: An optax schedule that infers the lr given the step.
    l1_regularizer_weight: Weight for L1 regularization.

  Returns:
    A ShardedGradientTransformation applying L1 weight decay.
  """

  def update_fn(updates, state, params):
    count = state.count
    lr_multiplier = learning_rate_fn(count)
    if l1_regularizer_weight:
      if params is None:
        raise ValueError('Params must not be empty when applying weight decay.')
      updates = jax.tree_multimap(
          lambda g, p: g - lr_multiplier * l1_regularizer_weight * jnp.sign(p),
          updates, params)
    updated_state = NestedMap(count=count + 1)
    return updates, updated_state

  return ShardedGradientTransformation(
      init=count_init_fn,
      update=update_fn,
      init_partition_spec=count_init_partition_spec_fn)


def sharded_adam(learning_rate_fn: optax.Schedule, beta1: float, beta2: float,
                 epsilon: float, epsilon_root: float, update_capping: float,
                 weight_decay: float) -> ShardedGradientTransformation:
  """Standard Adam optimizer that also supports sharding.

  This Adam optimizer supports optional update capping when update_capping is >
  0. Update capping can help stabilizing model learning, avoiding excessive
  updates when gradient variance estimate is stale (e.g. when data distribution
  suddenly shifts).

  Args:
    learning_rate_fn: a callable that given the current training step, returns
      the learning rate to apply.
    beta1: decay rate to track the first moment.
    beta2: decay rate to track the second moment.
    epsilon: Small constant applied to the denominator outside of the square
      root to avoid dividing by zero when rescaling.
    epsilon_root: Small constant applied to the denominator inside of the square
      root to avoid dividing by zero when rescaling.
    update_capping: If > 0, cap mean update to at most this value.
    weight_decay: If > 0, weight decay to apply.

  Returns:
    A `ShardedGradientTransformation`.
  """
  helper = _ShardedAdamHelper()

  def init_fn(mdl_vars):
    slot_vars = jax.tree_map(helper.init_opt_state, mdl_vars)
    count = jnp.array(0, dtype=jnp.int32)
    return NestedMap(
        count=count,
        m=jax.tree_map(lambda x: x.m, slot_vars),
        v=jax.tree_map(lambda x: x.v, slot_vars))

  def init_partition_spec_fn(mdl_params):
    slot_vars = jax.tree_map(helper.opt_state_sharding_spec, mdl_params)
    count = py_utils.weight_params(
        shape=[], init=None, dtype=jnp.int32, collections=None)

    return NestedMap(
        count=count,
        m=jax.tree_map(lambda x: x.m, slot_vars),
        v=jax.tree_map(lambda x: x.v, slot_vars))

  def update_fn(updates, state, params=None):
    # Sanitize updates just in case.
    if weight_decay > 0:
      assert params is not None
    updates = jax.tree_multimap(helper.sanitize_values, updates)
    count = state.count

    def _update_momentum(g, m, v):
      return helper.update_moments(count, g, _AdamOptState(m=m, v=v), beta1,
                                   beta2)

    updated_moments = jax.tree_multimap(_update_momentum, updates, state.m,
                                        state.v)

    m = jax.tree_map(lambda x: x.m, updated_moments)
    v = jax.tree_map(lambda x: x.v, updated_moments)

    updates = jax.tree_multimap(
        lambda m, v: m / (jnp.sqrt(v + epsilon_root) + epsilon), m, v)

    if update_capping > 0:
      updates = jax.tree_map(lambda x: helper.clip_update(x, update_capping),
                             updates)

    if weight_decay > 0:
      updates = jax.tree_multimap(lambda x, v: x + weight_decay * v, updates,
                                  params)

    step_size = -1.0 * learning_rate_fn(count)
    # Finally, fold in step size.
    updates = jax.tree_map(lambda x: step_size * x, updates)

    updated_states = NestedMap(count=count + 1, m=m, v=v)
    return updates, updated_states

  return ShardedGradientTransformation(
      init=init_fn,
      update=update_fn,
      init_partition_spec=init_partition_spec_fn)


class BaseOptimizer:
  """Base class for all optimizers."""

  @classmethod
  def Params(cls) -> InstantiableParams:  # pylint: disable=invalid-name
    """Defines hyper-params for all optimizers."""
    p = InstantiableParams(cls)
    p.Define(
        'l2_regularizer_weight', None,
        'If not None, L2 regularization to apply to the model weights. '
        'Otherwise, disable L2 regularization.')
    p.Define(
        'l1_regularizer_weight', None,
        'If not None, L1 regularization to apply to the model weights. '
        'Otherwise, disable L1 regularization.')
    p.Define(
        'clip_gradient_norm_to_value', 0.0,
        'Clip gradient by global norm to this value. This is similar to '
        'the bahaviour of tf.clip_by_global_norm. If you are looking for '
        'tf.clip_by_norm refer to clip_gradient_single_norm_to_value. Note '
        'these are mutually exclusive.')
    p.Define(
        'clip_gradient_single_norm_to_value', 0.0,
        'Clip gradient by single tensor norm to this value. This is '
        'similar to the bahaviour of tf.clip_by_norm. Note this is mutually '
        'exclusive to using clip_gradient_norm_to_value.')
    p.Define('learning_rate', 0.0, 'learning rate to use.')
    p.Define('lr_schedule', None, 'Learning rate decay schedule.')
    return p

  def __init__(self, params: InstantiableParams) -> None:
    self._params = params.Copy()
    p = self._params
    self._lr_schedule = self._params.lr_schedule.Instantiate()
    # Should not mix L1 and L2 weight decay together.
    if p.l2_regularizer_weight and p.l1_regularizer_weight:
      raise ValueError('Should not mix L1 and L2 regularization together.')

  @property
  def params(self) -> InstantiableParams:
    return self._params

  def get_learning_rate(self, step_count: JTensor) -> JTensor:
    """Get the learning rate of this optimizer at a particular step."""
    return self._lr_schedule.value(step_count) * self.params.learning_rate

  def get_grad_transformation(
      self
  ) -> Union[optax.GradientTransformation, ShardedGradientTransformation]:
    """Get the grad transformation corresponds to this optimizer config.

    This is the final gradient transformation that incorporates all
    transformations.

    Returns:
      an optax.GradientTransformation or ShardedGradientTransformation.
    """
    # TODO(yonghui): respect gradient clipping, etc transformations.
    p = self.params
    return sharded_chain(
        self._get_raw_grad_transformation(self.get_learning_rate),
        apply_l1_weight_decay(
            self.get_learning_rate,
            l1_regularizer_weight=p.l1_regularizer_weight),
        apply_l2_weight_decay(
            self.get_learning_rate,
            l2_regularizer_weight=p.l2_regularizer_weight))

  def _get_raw_grad_transformation(
      self, lr: optax.Schedule
  ) -> Union[optax.GradientTransformation, ShardedGradientTransformation]:
    """Get the raw optimizer transformation without taking into other ...

    transformations such l1/l2 regularization, gradient norm clipping, etc.

    Args:
      lr: an optax schedule.

    Returns:
      an optax.GradientTransformation or ShardedGradientTransformation.
    """
    raise NotImplementedError()


class Sgd(BaseOptimizer):
  """Canonical SGD optimizer."""

  @classmethod
  def Params(cls) -> InstantiableParams:  # pylint: disable=invalid-name
    p = super().Params()
    p.Define(
        'momentum', None,
        'Decay rate used by the momentum term. If set to None, momentum is not '
        'used.')
    p.Define('nesterov', False, 'Whether Nesterov momentum is used or not.')
    return p

  def _get_raw_grad_transformation(
      self, lr: optax.Schedule) -> optax.GradientTransformation:
    p = self._params
    return optax.sgd(learning_rate=lr, momentum=p.momentum, nesterov=p.nesterov)


class ShardedSgd(BaseOptimizer):
  """Sharded SGD optimizer."""

  @classmethod
  def Params(cls) -> InstantiableParams:  # pylint: disable=invalid-name
    p = super().Params()
    p.Define(
        'momentum', None,
        'Decay rate used by the momentum term. If set to None, momentum is not '
        'used.')
    p.Define('nesterov', False, 'Whether Nesterov momentum is used or not.')
    return p

  def _get_raw_grad_transformation(
      self, lr: optax.Schedule) -> ShardedGradientTransformation:
    p = self._params
    return sharded_sgd(
        learning_rate_fn=lr, momentum=p.momentum, nesterov=p.nesterov)


class Adam(BaseOptimizer):
  """Adam optimizer."""

  @classmethod
  def Params(cls) -> InstantiableParams:  # pylint: disable=invalid-name
    p = super().Params()
    p.Define(
        'beta1', 0.9,
        'Expenonential decay rate to track the first moment of past gradients.')
    p.Define(
        'beta2', 0.999,
        'Exponential decay rate to track the second moment of past gradients.')
    p.Define(
        'epsilon', 1e-6,
        'Small constant applied to the denominator outside of the square root '
        'to avoid dividing by zero when rescaling.')
    p.Define(
        'epsilon_root', 0.0,
        'Small constant applied to the denominator inside of the square root '
        'to avoid dividing by zero when rescaling.')
    # update clipping (as specified by clip_threshold) and weight decay features
    # are available only when sharded_adam is set to True.
    p.Define('clip_threshold', 1.0,
             'An optional float to clip raw adam updates to.')
    p.Define('weight_decay', 0.0, 'Decoupled weight decay to apply.')
    p.Define('sharded_adam', True, 'whether or not to use sharded_adam')
    return p

  @classmethod
  def ParamsA(cls) -> InstantiableParams:  # pylint: disable=invalid-name
    """Convenient method for a commonly used Adam config."""
    return cls.Params().Set(beta1=0.9, beta2=0.997, epsilon=1e-9)

  @classmethod
  def ParamsB(cls) -> InstantiableParams:  # pylint: disable=invalid-name
    """Convenient method for another commonly used Adam config."""
    return cls.Params().Set(beta1=0.9, beta2=0.98, epsilon=1e-9)

  def _get_raw_grad_transformation(
      self, lr: optax.Schedule) -> optax.GradientTransformation:
    p = self._params
    if p.sharded_adam:
      logging.info('Using sharded_adam.')
      return sharded_adam(
          learning_rate_fn=lr,
          beta1=p.beta1,
          beta2=p.beta2,<commit_msg>Set optax.OptState to chex.ArrayTree.<commit_after># Lint as: python3
# Copyright 2021 The TensorFlow Authors. 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.
# ==============================================================================
"""Module for all optimizers."""

import dataclasses
import functools
from typing import Any, Callable, NamedTuple, Optional, Sequence, Tuple, Union

from absl import logging
import jax
from jax import numpy as jnp
from lingvo.jax import gshard_utils
from lingvo.jax import py_utils
from lingvo.jax import pytypes
import optax

NestedMap = py_utils.NestedMap
InstantiableParams = py_utils.InstantiableParams
JTensor = pytypes.JTensor
NestedJTensor = pytypes.NestedJTensor
NestedParams = pytypes.NestedParams

# Initializes sharding spec for the optimizer state variables.
TransformInitPartitionSpecFn = Callable[[NestedParams],
                                        Union[NestedParams,
                                              Sequence[NestedParams]]]


# Extension of optax.GradientTransformation that supports spmd sharding and
# explicit annotation of sharding specs for the optimizer state variables.
@dataclasses.dataclass(frozen=True)
class ShardedGradientTransformation:
  """GradientTransformation that supports spmd."""
  # init_fn and update_fn are the same as in optax.GradientTransformation
  init: optax.TransformInitFn
  update: optax.TransformUpdateFn
  # Input is the sharding specs of the variables used in the forward
  # computation.  Output is the sharding specs of the optimizer state variables.
  #
  # Constraints: output from this function should be of identical structure as
  # that of the init() function.
  init_partition_spec: TransformInitPartitionSpecFn


def count_init_fn(_):
  """Common init_fn that initializes a count for global step."""
  return NestedMap(count=jnp.zeros([], jnp.int32))


def count_init_partition_spec_fn(var_params):
  """Init partition spec for only partitioning the count/step."""
  var_spec_flattened, _ = jax.tree_flatten(var_params)
  assert var_spec_flattened
  first_var = var_spec_flattened[0]
  assert isinstance(first_var, py_utils.Params)
  device_mesh = first_var.device_mesh
  return NestedMap(
      count=py_utils.weight_params(
          shape=[],
          init=None,
          dtype=jnp.int32,
          collections=None,
          device_mesh=device_mesh,
          tensor_split_dims_mapping=[]))


def sharded_sgd(learning_rate_fn: optax.Schedule, momentum: Optional[float],
                nesterov: bool) -> ShardedGradientTransformation:
  """A canonical Stochastic Gradient Descent optimiser that supports spmd ...

  sharding.

  This implements stochastic gradient descent. It also includes support for
  momentum, and nesterov acceleration, as these are standard practice when
  using stochastic gradient descent to train deep neural networks.

  References:
    Sutskever et al, 2013: http://proceedings.mlr.press/v28/sutskever13.pdf

  Args:
    learning_rate_fn: a callable that given the current training step, returns
      the learning rate to apply.
    momentum: (default `None`), the `decay` rate used by the momentum term, when
      it is set to `None`, then momentum is not used at all.
    nesterov (default `False`): whether nesterov momentum is used.

  Returns:
    A `ShardedGradientTransformation`.
  """
  # TODO(yonghui): support momentum.
  assert momentum is None
  del nesterov

  def update_fn(updates, state, params=None):
    del params
    step_size = -1.0 * learning_rate_fn(state.count)
    updates = jax.tree_map(lambda g: jnp.array(step_size, dtype=g.dtype) * g,
                           updates)
    updated_states = NestedMap(count=state.count + jnp.array(1, jnp.int32))
    return updates, updated_states

  return ShardedGradientTransformation(
      init=count_init_fn,
      update=update_fn,
      init_partition_spec=count_init_partition_spec_fn)


class _AdamOptState:

  def __init__(self, *, m, v):
    self.m = m
    self.v = v


class _ShardedAdamHelper:
  """A helper class facilitates the creation of sharded_adam_optimizer."""

  def opt_state_sharding_spec(self,
                              var_params: py_utils.Params) -> _AdamOptState:
    """Returns optimizer sharding spec for one particular variable."""
    m_var_params = var_params.Copy()
    m_var_params.init = None
    v_var_params = var_params.Copy()
    v_var_params.init = None
    # m and v simply share the same sharding.
    return _AdamOptState(m=m_var_params, v=v_var_params)

  def init_opt_state(self, var_params: py_utils.Params) -> _AdamOptState:
    """Returns optimizer state for one particular variable."""
    return _AdamOptState(
        m=jnp.zeros_like(var_params), v=jnp.zeros_like(var_params))

  def sanitize_values(self, array: JTensor, replacement: float = 0.0):
    """Sanitizes NaN and Infinity values."""
    return jnp.nan_to_num(
        array, nan=replacement, posinf=replacement, neginf=replacement)

  def bias_corrected_decay(self, step: JTensor, decay: float) -> JTensor:
    """Incorporates bias correction into decay.

    Please see section 7.1 in https://arxiv.org/pdf/1804.04235.pdf for the
    derivation of the formulas below. With bias-corrected decay, we can simply
    do

    m_{t} = decay1 * m_{t-1} + (1 - decay1) * g
    v_{t} = decay2 * v_{t-1} + (1 - decay2) * g ^ 2

    without further bias correction.

    Args:
      step: current step, 0-based.
      decay: the raw decay. As t -> infinity, bias corrected decay converges to
        this value.

    Returns:
      Bias corrected decay.
    """
    t = step.astype(jnp.float32) + 1.
    return decay * (1. - jnp.power(decay, t - 1.)) / (1. - jnp.power(decay, t))

  def update_moments(self, step: JTensor, update: JTensor,
                     moments: _AdamOptState, beta1: float,
                     beta2: float) -> _AdamOptState:
    """Updates momentum values."""
    beta1_decay = self.bias_corrected_decay(step, beta1)
    beta2_decay = self.bias_corrected_decay(step, beta2)
    m = (1.0 - beta1_decay) * update + beta1_decay * moments.m
    v = (1.0 - beta2_decay) * (update**2) + beta2_decay * moments.v
    return _AdamOptState(m=m, v=v)

  def clip_update(self, update: JTensor, clip_threshold: float) -> JTensor:
    mean_update = self.sanitize_values(reduce_rms(update), 1.0)
    clip_threshold = jnp.array(clip_threshold, dtype=update.dtype)
    denom = jnp.maximum(1.0, mean_update / clip_threshold)
    return update / denom


def sharded_chain(
    *args: Union[optax.GradientTransformation, ShardedGradientTransformation]
) -> ShardedGradientTransformation:
  """Applies a list of (possibly sharded) chainable update transformations.

  Given a sequence of chainable transforms, `sharded_chain` returns an `init_fn`
  that constructs a `state` by concatenating the states of the individual
  transforms, and returns an `update_fn` which chains the update transformations
  feeding the appropriate state to each. In addition, it differs from the optax
  `chain` function, by also supporting ShardedGradientTransformation by chaining
  also the `init_partition_spec_fn`. If there are no
  ShardedGradientTransformations in the chain, the sharding specs will be
  None, meaning all the variables are replicated.

  Args:
    *args: a sequence of chainable GradientTransformations or
      ShardedGradientTransformations or a combination of both.

  Returns:
    A single chained ShardedGradientTransformation.
  """

  def init_fn(params):
    return tuple(fn.init(params) for fn in args)

  def update_fn(updates, state, params=None):
    if len(args) != len(state):
      raise ValueError('The number of updates and states has to be the same in '
                       'sharded chain.')

    new_state = []
    for s, fn in zip(state, args):
      updates, new_s = fn.update(updates, s, params)
      new_state.append(new_s)
    return updates, tuple(new_state)

  def init_partition_spec_fn(mdl_vars):
    partition_specs = []
    for fn in args:
      if isinstance(fn, ShardedGradientTransformation):
        nmap = fn.init_partition_spec(mdl_vars)
        partition_specs.append(nmap)
      else:
        # Replicate the states.
        partition_specs.append(None)
    return tuple(partition_specs)

  return ShardedGradientTransformation(
      init=init_fn,
      update=update_fn,
      init_partition_spec=init_partition_spec_fn)


def apply_l2_weight_decay(
    learning_rate_fn: optax.Schedule,
    l2_regularizer_weight: Optional[float] = 0.
) -> ShardedGradientTransformation:
  """Applies L2 weight decay.

  Args:
    learning_rate_fn: An optax schedule that infers the lr given the step.
    l2_regularizer_weight: Weight for L2 regularization.

  Returns:
    A ShardedGradientTransformation applying L2 weight decay.
  """

  def update_fn(updates, state, params):
    count = state.count
    lr_multiplier = learning_rate_fn(count)
    if l2_regularizer_weight:
      if params is None:
        raise ValueError('Params must not be empty when applying weight decay.')
      updates = jax.tree_multimap(
          lambda g, p: g - lr_multiplier * l2_regularizer_weight * p, updates,
          params)
    updated_state = NestedMap(count=count + 1)
    return updates, updated_state

  return ShardedGradientTransformation(
      init=count_init_fn,
      update=update_fn,
      init_partition_spec=count_init_partition_spec_fn)


def apply_l1_weight_decay(
    learning_rate_fn: optax.Schedule,
    l1_regularizer_weight: Optional[float] = 0.
) -> ShardedGradientTransformation:
  """Applies L1 weight decay.

  Args:
    learning_rate_fn: An optax schedule that infers the lr given the step.
    l1_regularizer_weight: Weight for L1 regularization.

  Returns:
    A ShardedGradientTransformation applying L1 weight decay.
  """

  def update_fn(updates, state, params):
    count = state.count
    lr_multiplier = learning_rate_fn(count)
    if l1_regularizer_weight:
      if params is None:
        raise ValueError('Params must not be empty when applying weight decay.')
      updates = jax.tree_multimap(
          lambda g, p: g - lr_multiplier * l1_regularizer_weight * jnp.sign(p),
          updates, params)
    updated_state = NestedMap(count=count + 1)
    return updates, updated_state

  return ShardedGradientTransformation(
      init=count_init_fn,
      update=update_fn,
      init_partition_spec=count_init_partition_spec_fn)


def sharded_adam(learning_rate_fn: optax.Schedule, beta1: float, beta2: float,
                 epsilon: float, epsilon_root: float, update_capping: float,
                 weight_decay: float) -> ShardedGradientTransformation:
  """Standard Adam optimizer that also supports sharding.

  This Adam optimizer supports optional update capping when update_capping is >
  0. Update capping can help stabilizing model learning, avoiding excessive
  updates when gradient variance estimate is stale (e.g. when data distribution
  suddenly shifts).

  Args:
    learning_rate_fn: a callable that given the current training step, returns
      the learning rate to apply.
    beta1: decay rate to track the first moment.
    beta2: decay rate to track the second moment.
    epsilon: Small constant applied to the denominator outside of the square
      root to avoid dividing by zero when rescaling.
    epsilon_root: Small constant applied to the denominator inside of the square
      root to avoid dividing by zero when rescaling.
    update_capping: If > 0, cap mean update to at most this value.
    weight_decay: If > 0, weight decay to apply.

  Returns:
    A `ShardedGradientTransformation`.
  """
  helper = _ShardedAdamHelper()

  def init_fn(mdl_vars):
    slot_vars = jax.tree_map(helper.init_opt_state, mdl_vars)
    count = jnp.array(0, dtype=jnp.int32)
    return NestedMap(
        count=count,
        m=jax.tree_map(lambda x: x.m, slot_vars),
        v=jax.tree_map(lambda x: x.v, slot_vars))

  def init_partition_spec_fn(mdl_params):
    slot_vars = jax.tree_map(helper.opt_state_sharding_spec, mdl_params)
    count = py_utils.weight_params(
        shape=[], init=None, dtype=jnp.int32, collections=None)

    return NestedMap(
        count=count,
        m=jax.tree_map(lambda x: x.m, slot_vars),
        v=jax.tree_map(lambda x: x.v, slot_vars))

  def update_fn(updates, state, params=None):
    # Sanitize updates just in case.
    if weight_decay > 0:
      assert params is not None
    updates = jax.tree_multimap(helper.sanitize_values, updates)
    count = state.count

    def _update_momentum(g, m, v):
      return helper.update_moments(count, g, _AdamOptState(m=m, v=v), beta1,
                                   beta2)

    updated_moments = jax.tree_multimap(_update_momentum, updates, state.m,
                                        state.v)

    m = jax.tree_map(lambda x: x.m, updated_moments)
    v = jax.tree_map(lambda x: x.v, updated_moments)

    updates = jax.tree_multimap(
        lambda m, v: m / (jnp.sqrt(v + epsilon_root) + epsilon), m, v)

    if update_capping > 0:
      updates = jax.tree_map(lambda x: helper.clip_update(x, update_capping),
                             updates)

    if weight_decay > 0:
      updates = jax.tree_multimap(lambda x, v: x + weight_decay * v, updates,
                                  params)

    step_size = -1.0 * learning_rate_fn(count)
    # Finally, fold in step size.
    updates = jax.tree_map(lambda x: step_size * x, updates)

    updated_states = NestedMap(count=count + 1, m=m, v=v)
    return updates, updated_states

  return ShardedGradientTransformation(
      init=init_fn,
      update=update_fn,
      init_partition_spec=init_partition_spec_fn)


class BaseOptimizer:
  """Base class for all optimizers."""

  @classmethod
  def Params(cls) -> InstantiableParams:  # pylint: disable=invalid-name
    """Defines hyper-params for all optimizers."""
    p = InstantiableParams(cls)
    p.Define(
        'l2_regularizer_weight', None,
        'If not None, L2 regularization to apply to the model weights. '
        'Otherwise, disable L2 regularization.')
    p.Define(
        'l1_regularizer_weight', None,
        'If not None, L1 regularization to apply to the model weights. '
        'Otherwise, disable L1 regularization.')
    p.Define(
        'clip_gradient_norm_to_value', 0.0,
        'Clip gradient by global norm to this value. This is similar to '
        'the bahaviour of tf.clip_by_global_norm. If you are looking for '
        'tf.clip_by_norm refer to clip_gradient_single_norm_to_value. Note '
        'these are mutually exclusive.')
    p.Define(
        'clip_gradient_single_norm_to_value', 0.0,
        'Clip gradient by single tensor norm to this value. This is '
        'similar to the bahaviour of tf.clip_by_norm. Note this is mutually '
        'exclusive to using clip_gradient_norm_to_value.')
    p.Define('learning_rate', 0.0, 'learning rate to use.')
    p.Define('lr_schedule', None, 'Learning rate decay schedule.')
    return p

  def __init__(self, params: InstantiableParams) -> None:
    self._params = params.Copy()
    p = self._params
    self._lr_schedule = self._params.lr_schedule.Instantiate()
    # Should not mix L1 and L2 weight decay together.
    if p.l2_regularizer_weight and p.l1_regularizer_weight:
      raise ValueError('Should not mix L1 and L2 regularization together.')

  @property
  def params(self) -> InstantiableParams:
    return self._params

  def get_learning_rate(self, step_count: JTensor) -> JTensor:
    """Get the learning rate of this optimizer at a particular step."""
    return self._lr_schedule.value(step_count) * self.params.learning_rate

  def get_grad_transformation(
      self
  ) -> Union[optax.GradientTransformation, ShardedGradientTransformation]:
    """Get the grad transformation corresponds to this optimizer config.

    This is the final gradient transformation that incorporates all
    transformations.

    Returns:
      an optax.GradientTransformation or ShardedGradientTransformation.
    """
    # TODO(yonghui): respect gradient clipping, etc transformations.
    p = self.params
    return sharded_chain(
        self._get_raw_grad_transformation(self.get_learning_rate),
        apply_l1_weight_decay(
            self.get_learning_rate,
            l1_regularizer_weight=p.l1_regularizer_weight),
        apply_l2_weight_decay(
            self.get_learning_rate,
            l2_regularizer_weight=p.l2_regularizer_weight))

  def _get_raw_grad_transformation(
      self, lr: optax.Schedule
  ) -> Union[optax.GradientTransformation, ShardedGradientTransformation]:
    """Get the raw optimizer transformation without taking into other ...

    transformations such l1/l2 regularization, gradient norm clipping, etc.

    Args:
      lr: an optax schedule.

    Returns:
      an optax.GradientTransformation or ShardedGradientTransformation.
    """
    raise NotImplementedError()


class Sgd(BaseOptimizer):
  """Canonical SGD optimizer."""

  @classmethod
  def Params(cls) -> InstantiableParams:  # pylint: disable=invalid-name
    p = super().Params()
    p.Define(
        'momentum', None,
        'Decay rate used by the momentum term. If set to None, momentum is not '
        'used.')
    p.Define('nesterov', False, 'Whether Nesterov momentum is used or not.')
    return p

  def _get_raw_grad_transformation(
      self, lr: optax.Schedule) -> optax.GradientTransformation:
    p = self._params
    return optax.sgd(learning_rate=lr, momentum=p.momentum, nesterov=p.nesterov)


class ShardedSgd(BaseOptimizer):
  """Sharded SGD optimizer."""

  @classmethod
  def Params(cls) -> InstantiableParams:  # pylint: disable=invalid-name
    p = super().Params()
    p.Define(
        'momentum', None,
        'Decay rate used by the momentum term. If set to None, momentum is not '
        'used.')
    p.Define('nesterov', False, 'Whether Nesterov momentum is used or not.')
    return p

  def _get_raw_grad_transformation(
      self, lr: optax.Schedule) -> ShardedGradientTransformation:
    p = self._params
    return sharded_sgd(
        learning_rate_fn=lr, momentum=p.momentum, nesterov=p.nesterov)


class Adam(BaseOptimizer):
  """Adam optimizer."""

  @classmethod
  def Params(cls) -> InstantiableParams:  # pylint: disable=invalid-name
    p = super().Params()
    p.Define(
        'beta1', 0.9,
        'Expenonential decay rate to track the first moment of past gradients.')
    p.Define(
        'beta2', 0.999,
        'Exponential decay rate to track the second moment of past gradients.')
    p.Define(
        'epsilon', 1e-6,
        'Small constant applied to the denominator outside of the square root '
        'to avoid dividing by zero when rescaling.')
    p.Define(
        'epsilon_root', 0.0,
        'Small constant applied to the denominator inside of the square root '
        'to avoid dividing by zero when rescaling.')
    # update clipping (as specified by clip_threshold) and weight decay features
    # are available only when sharded_adam is set to True.
    p.Define('clip_threshold', 1.0,
             'An optional float to clip raw adam updates to.')
    p.Define('weight_decay', 0.0, 'Decoupled weight decay to apply.')
    p.Define('sharded_adam', True, 'whether or not to use sharded_adam')
    return p

  @classmethod
  def ParamsA(cls) -> InstantiableParams:  # pylint: disable=invalid-name
    """Convenient method for a commonly used Adam config."""
    return cls.Params().Set(beta1=0.9, beta2=0.997, epsilon=1e-9)

  @classmethod
  def ParamsB(cls) -> InstantiableParams:  # pylint: disable=invalid-name
    """Convenient method for another commonly used Adam config."""
    return cls.Params().Set(beta1=0.9, beta2=0.98, epsilon=1e-9)

  def _get_raw_grad_transformation(
      self, lr: optax.Schedule
  ) -> Union[optax.GradientTransformation, ShardedGradientTransformation]:
    p = self._params
    if p.sharded_adam:
      logging.info('Using sharded_adam.')
      return sharded_adam(
          learning_rate_fn=lr,
          beta1=p.beta1,
          beta2=p.beta2,
{'fraction_non_alphanumeric': 0.08792437240273154, 'fraction_numerical': 0.008748693771416073, 'mean_word_length': 3.3265692356219114, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 65, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7024471', 'n_tokens_mistral': 12939, 'n_tokens_neox': 12339, 'n_words': 3703}