<commit_before># Copyright 2018 Google LLC
#
# 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
#
#     https://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.

"""A composable gradient processing and optimization library for JAX.

The ``optix`` module implements a number of composable gradient transformations,
typically used in the context of optimizing neural nets.

Each transformation defines:

* an ``init_fn``, to initialize a (possibly empty) set of statistics, or ``state``.
* an ``update_fn`` to transform an input gradient and update the state.

An (optional) ``chain`` utility can be used to build custom optimizers by
chaining arbitrary sequences of transformations. For any sequence of
transformations ``chain`` returns a single ``init_fn`` and ``update_fn``.

An (optional) ``apply_updates`` function can be used to eventually apply the
transformed gradients to the set of parameters of interest.

Separating gradient transformations from the parameter update allows to flexibly
chain a sequence of transformations of the same gradients, as well as combine
multiple updates to the same parameters (e.g. in multi-task settings where the
different tasks may benefit from different sets of gradient transformations).

Many popular optimizers can be implemented using ``optix`` as one-liners, and,
for convenience, we provide aliases for some of the most popular ones.
"""


import collections

from jax import numpy as jnp
from jax import random as jrandom

from jax.tree_util import tree_leaves
from jax.tree_util import tree_multimap
from jax.tree_util import tree_structure
from jax.tree_util import tree_unflatten


### Composable gradient transformations. ###


InitUpdate = collections.namedtuple("InitUpdate", ("init", "update"))
ClipState = collections.namedtuple("ClipState", "")


def clip(max_delta):
  """Clip updates element-wise, to be between -max_delta and +max_delta.

  Args:
    max_delta: the maximum absolute value for each element in the update.

  Returns:
    An (init_fn, update_fn) tuple.
  """

  def init_fn(_):
    return ClipState()

  def update_fn(updates, state):
    updates = tree_multimap(
        lambda g: jnp.clip(g, -max_delta, max_delta), updates)
    return updates, state

  return InitUpdate(init_fn, update_fn)


ClipByGlobalNormState = collections.namedtuple("ClipByGlobalNormState", "")


def global_norm(items):
  return jnp.sqrt(jnp.sum([jnp.sum(x**2) for x in tree_leaves(items)]))
_global_norm = global_norm  # TODO(mtthss): remove when google code updated


def clip_by_global_norm(max_norm):
  """Clip updates using their global norm.

  References:
    [Pascanu et al, 2012](https://arxiv.org/abs/1211.5063)

  Args:
    max_norm: the maximum global norm for an update.

  Returns:
    An (init_fn, update_fn) tuple.
  """

  def init_fn(_):
    return ClipByGlobalNormState()

  def update_fn(updates, state):
    g_norm = global_norm(updates)
    trigger = g_norm < max_norm
    updates = tree_multimap(
        lambda t: jnp.where(trigger, t, t * (max_norm / g_norm)), updates)
    return updates, state

  return InitUpdate(init_fn, update_fn)


TraceState = collections.namedtuple("TraceState", "trace")


def trace(decay, nesterov):
  """Compute a trace of past updates.

  Args:
    decay: the decay rate for the tracing of past updates.
    nesterov: whether to use nesterov momentum.

  Returns:
    An (init_fn, update_fn) tuple.
  """
<commit_msg>Make gradient clipping by norm more numerically safe<commit_after># Copyright 2018 Google LLC
#
# 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
#
#     https://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.

"""A composable gradient processing and optimization library for JAX.

The ``optix`` module implements a number of composable gradient transformations,
typically used in the context of optimizing neural nets.

Each transformation defines:

* an ``init_fn``, to initialize a (possibly empty) set of statistics, or ``state``.
* an ``update_fn`` to transform an input gradient and update the state.

An (optional) ``chain`` utility can be used to build custom optimizers by
chaining arbitrary sequences of transformations. For any sequence of
transformations ``chain`` returns a single ``init_fn`` and ``update_fn``.

An (optional) ``apply_updates`` function can be used to eventually apply the
transformed gradients to the set of parameters of interest.

Separating gradient transformations from the parameter update allows to flexibly
chain a sequence of transformations of the same gradients, as well as combine
multiple updates to the same parameters (e.g. in multi-task settings where the
different tasks may benefit from different sets of gradient transformations).

Many popular optimizers can be implemented using ``optix`` as one-liners, and,
for convenience, we provide aliases for some of the most popular ones.
"""


import collections

from jax import numpy as jnp
from jax import random as jrandom

from jax.tree_util import tree_leaves
from jax.tree_util import tree_multimap
from jax.tree_util import tree_structure
from jax.tree_util import tree_unflatten


### Composable gradient transformations. ###


InitUpdate = collections.namedtuple("InitUpdate", ("init", "update"))
ClipState = collections.namedtuple("ClipState", "")


def clip(max_delta):
  """Clip updates element-wise, to be between -max_delta and +max_delta.

  Args:
    max_delta: the maximum absolute value for each element in the update.

  Returns:
    An (init_fn, update_fn) tuple.
  """

  def init_fn(_):
    return ClipState()

  def update_fn(updates, state):
    updates = tree_multimap(
        lambda g: jnp.clip(g, -max_delta, max_delta), updates)
    return updates, state

  return InitUpdate(init_fn, update_fn)


ClipByGlobalNormState = collections.namedtuple("ClipByGlobalNormState", "")


def global_norm(items):
  return jnp.sqrt(jnp.sum([jnp.sum(x**2) for x in tree_leaves(items)]))


def clip_by_global_norm(max_norm):
  """Clip updates using their global norm.

  References:
    [Pascanu et al, 2012](https://arxiv.org/abs/1211.5063)

  Args:
    max_norm: the maximum global norm for an update.

  Returns:
    An (init_fn, update_fn) tuple.
  """

  def init_fn(_):
    return ClipByGlobalNormState()

  def update_fn(updates, state):
    g_norm = global_norm(updates)
    trigger = g_norm < max_norm
    updates = tree_multimap(
        lambda t: jnp.where(trigger, t, (t / g_norm) * max_norm), updates)
    return updates, state

  return InitUpdate(init_fn, update_fn)


TraceState = collections.namedtuple("TraceState", "trace")


def trace(decay, nesterov):
  """Compute a trace of past updates.

  Args:
    decay: the decay rate for the tracing of past updates.
    nesterov: whether to use nesterov momentum.

  Returns:
    An (init_fn, update_fn) tuple.
  """

{'fraction_non_alphanumeric': 0.08327903036621921, 'fraction_numerical': 0.00547373908510361, 'mean_word_length': 4.423321554770318, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 4, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6507124', 'n_tokens_mistral': 2241, 'n_tokens_neox': 2139, 'n_words': 897}