<commit_before># Copyright 2021 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.
# ==============================================================================
"""Optimizers and Schedulers, inspired by the PyTorch API."""
from collections import ChainMap  # pylint:disable=g-importing-member
from typing import Callable
import haiku as hk
import jax
import jax.numpy as jnp
import tree
from nfnets import utils


class Optimizer(object):
  """Optimizer base class."""

  def __init__(self, params, defaults):
    # Flag indicating if parameters have been broadcasted
    self._broadcasted = False
    # Optimizer hyperparameters; this is a dict to support using param_groups
    self._hyperparameters = {}
    # Mapping from model parameters to optimizer hyperparameters
    self._params2hyperparams = {}
    # Assign defaults
    self._hyperparameters = dict(**defaults)
    # Prepare parameter groups and mappings
    self.create_param_groups(params, defaults)
    # Join params at top-level if params is a list of groups
    if isinstance(params, list):
      flatmap = type(hk.data_structures.to_immutable_dict({}))
      if any([isinstance(group['params'], flatmap) for group in params]):
        params = hk.data_structures.merge(*[group['params']
                                            for group in params])
      else:
        params = dict(ChainMap(*[group['params'] for group in params]))
    # Prepare states
    create_buffers = lambda k, v: self.create_buffers('/'.join(k), v)
    self._states = tree.map_structure_with_path(create_buffers, params)

  def add_hyperparam_group(self, group, suffix, defaults):
    """Adds new hyperparameters to the hyperparams dict."""
    # Use default hyperparams unless overridden by group hyperparams
    group_dict = {key: key for key in defaults if key not in group}
    for key in group:
      if key != 'params':  # Reserved keyword 'params'
        group_dict[key] = '%s_%s' % (key, suffix)
        self._hyperparameters[group_dict[key]] = group[key]
    # Set up params2hyperparams
    def set_p2h(k, _):
      self._params2hyperparams['/'.join(k)] = group_dict
    tree.map_structure_with_path(set_p2h, group['params'])

  def create_param_groups(self, params, defaults):
    """Creates param-hyperparam mappings."""
    if isinstance(params, list):
      for group_index, group in enumerate(params):
        # Add group to hyperparams and get this group's full hyperparameters
        self.add_hyperparam_group(group, group_index, defaults)
    else:
      mapping = {key: key for key in self._hyperparameters}
      def set_p2h(k, _):
        self._params2hyperparams['/'.join(k)] = mapping
      tree.map_structure_with_path(set_p2h, params)

  def create_buffers(self, name, params):
    """Method to be overridden by child classes."""
    pass

  def get_opt_params(self, param_name, itr):
    """Returns hyperparams corresponding to param_name."""
    mapping = self._params2hyperparams[param_name]
    output = {}
    for key in mapping:
      hyper = self._hyperparameters[mapping[key]]
      # Handle the case where a hyper is a class, for hybrids
      if isinstance(hyper, Callable) and not isinstance(hyper, type):
        output[key] = hyper(itr)
      else:
        output[key] = hyper
    return output

  def get_hyper(self, param_name, hyper_name):
    """Get an individual hyperparam for a given param."""
    mapping = self._params2hyperparams[param_name]
    return self._hyperparameters[mapping[hyper_name]]

  def plugin(self, states):
    self._states = states

  def states(self):
    return self._states

  def broadcast(self):
    """Brodcasts all buffers and parameters."""
    self._broadcasted = True
    for name, state in self._states.items():
      self._states[name] = {key: utils.broadcast(state[key]) for key in state}

  def gather(self):
    """Gathers state (if broadcasted) for saving."""
    states = {}
    for name in self._states:
      state = self._states[name]
      states[name] = {key: state[key] if state[key] is None else state[key][0]
                      for key in state}
    return states

  def __setattr__(self, name, value):
    """Overrides the object's set-attribute function to register states, etc."""
    if '_hyperparameters' in self.__dict__ and name in self._hyperparameters:
      self._hyperparameters[name] = value
    elif '_states' in self.__dict__ and name in self._states:
      self._states[name] = value
    else:
      object.__setattr__(self, name, value)

  def __getattr__(self, name):
    """Override the object's get-attribute function to return states, etc."""
    if '_hyperparameters' in self.__dict__ and name in self._hyperparameters:
      return self._hyperparameters[name]
    elif '_states' in self.__dict__ and name in self._states:
      return self._states[name]
    else:
      object.__getattribute__(self, name)

  def step(self, params, grads, states, itr=None):
    """Takes a single optimizer step.

    Args:
      params: a dict containing the parameters to be updated.
      grads: a dict containing the gradients for each parameter in params.
      states: a dict containing any optimizer buffers (momentum, etc) for
        each parameter in params.
      itr: an optional integer indicating the current step, for scheduling.
    Returns:
       The updated params and optimizer buffers.
    """
    get_hyper = lambda k, v: self.get_opt_params('/'.join(k), itr)
    hypers = tree.map_structure_with_path(get_hyper, params)
    outs = tree.map_structure_up_to(params, self.update_param,
                                    params, grads, states, hypers)
    return utils.split_tree(outs, params, 2)


class Schedule(object):
  """Hyperparameter scheduling objects."""


class CosineDecay(Schedule):
  """Cosine decay."""

  def __init__(self, min_val, max_val, num_steps):
    self.min_val = min_val
    self.max_val = max_val
    self.num_steps = num_steps

  def __call__(self, itr):
    cos = (1 + jnp.cos(jnp.pi * itr / self.num_steps))
    return 0.5 * (self.max_val - self.min_val) * cos + self.min_val


class WarmupCosineDecay(Schedule):
  """Cosine decay with linear warmup."""

  def __init__(self, start_val, min_val, max_val, num_steps, warmup_steps):
    self.start_val = start_val
    self.min_val = min_val
    self.max_val = max_val<commit_msg>Update NFNet optim to reflect Haiku default of dict instead of FlatMapping.<commit_after># Copyright 2021 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.
# ==============================================================================
"""Optimizers and Schedulers, inspired by the PyTorch API."""
from collections import ChainMap  # pylint:disable=g-importing-member
from typing import Callable, Mapping
import haiku as hk
import jax
import jax.numpy as jnp
import tree
from nfnets import utils


class Optimizer(object):
  """Optimizer base class."""

  def __init__(self, params, defaults):
    # Flag indicating if parameters have been broadcasted
    self._broadcasted = False
    # Optimizer hyperparameters; this is a dict to support using param_groups
    self._hyperparameters = {}
    # Mapping from model parameters to optimizer hyperparameters
    self._params2hyperparams = {}
    # Assign defaults
    self._hyperparameters = dict(**defaults)
    # Prepare parameter groups and mappings
    self.create_param_groups(params, defaults)
    # Join params at top-level if params is a list of groups
    if isinstance(params, list):
      if any(_is_non_empty_two_level_mapping(g['params']) for g in params):
        params = hk.data_structures.merge(*[g['params'] for g in params])
      else:
        params = dict(ChainMap(*[g['params'] for g in params]))
    # Prepare states
    create_buffers = lambda k, v: self.create_buffers('/'.join(k), v)
    self._states = tree.map_structure_with_path(create_buffers, params)

  def add_hyperparam_group(self, group, suffix, defaults):
    """Adds new hyperparameters to the hyperparams dict."""
    # Use default hyperparams unless overridden by group hyperparams
    group_dict = {key: key for key in defaults if key not in group}
    for key in group:
      if key != 'params':  # Reserved keyword 'params'
        group_dict[key] = '%s_%s' % (key, suffix)
        self._hyperparameters[group_dict[key]] = group[key]
    # Set up params2hyperparams
    def set_p2h(k, _):
      self._params2hyperparams['/'.join(k)] = group_dict
    tree.map_structure_with_path(set_p2h, group['params'])

  def create_param_groups(self, params, defaults):
    """Creates param-hyperparam mappings."""
    if isinstance(params, list):
      for group_index, group in enumerate(params):
        # Add group to hyperparams and get this group's full hyperparameters
        self.add_hyperparam_group(group, group_index, defaults)
    else:
      mapping = {key: key for key in self._hyperparameters}
      def set_p2h(k, _):
        self._params2hyperparams['/'.join(k)] = mapping
      tree.map_structure_with_path(set_p2h, params)

  def create_buffers(self, name, params):
    """Method to be overridden by child classes."""
    pass

  def get_opt_params(self, param_name, itr):
    """Returns hyperparams corresponding to param_name."""
    mapping = self._params2hyperparams[param_name]
    output = {}
    for key in mapping:
      hyper = self._hyperparameters[mapping[key]]
      # Handle the case where a hyper is a class, for hybrids
      if isinstance(hyper, Callable) and not isinstance(hyper, type):
        output[key] = hyper(itr)
      else:
        output[key] = hyper
    return output

  def get_hyper(self, param_name, hyper_name):
    """Get an individual hyperparam for a given param."""
    mapping = self._params2hyperparams[param_name]
    return self._hyperparameters[mapping[hyper_name]]

  def plugin(self, states):
    self._states = states

  def states(self):
    return self._states

  def broadcast(self):
    """Brodcasts all buffers and parameters."""
    self._broadcasted = True
    for name, state in self._states.items():
      self._states[name] = {key: utils.broadcast(state[key]) for key in state}

  def gather(self):
    """Gathers state (if broadcasted) for saving."""
    states = {}
    for name in self._states:
      state = self._states[name]
      states[name] = {key: state[key] if state[key] is None else state[key][0]
                      for key in state}
    return states

  def __setattr__(self, name, value):
    """Overrides the object's set-attribute function to register states, etc."""
    if '_hyperparameters' in self.__dict__ and name in self._hyperparameters:
      self._hyperparameters[name] = value
    elif '_states' in self.__dict__ and name in self._states:
      self._states[name] = value
    else:
      object.__setattr__(self, name, value)

  def __getattr__(self, name):
    """Override the object's get-attribute function to return states, etc."""
    if '_hyperparameters' in self.__dict__ and name in self._hyperparameters:
      return self._hyperparameters[name]
    elif '_states' in self.__dict__ and name in self._states:
      return self._states[name]
    else:
      object.__getattribute__(self, name)

  def step(self, params, grads, states, itr=None):
    """Takes a single optimizer step.

    Args:
      params: a dict containing the parameters to be updated.
      grads: a dict containing the gradients for each parameter in params.
      states: a dict containing any optimizer buffers (momentum, etc) for
        each parameter in params.
      itr: an optional integer indicating the current step, for scheduling.
    Returns:
       The updated params and optimizer buffers.
    """
    get_hyper = lambda k, v: self.get_opt_params('/'.join(k), itr)
    hypers = tree.map_structure_with_path(get_hyper, params)
    outs = tree.map_structure_up_to(params, self.update_param,
                                    params, grads, states, hypers)
    return utils.split_tree(outs, params, 2)


def _is_non_empty_two_level_mapping(obj):
  instof = lambda t: lambda v: isinstance(v, t)
  # Basically: isinstance(obj, Mapping[str, Mapping[str, Any]]) ...
  return (isinstance(obj, Mapping) and all(map(instof(str), obj.keys())) and
          all(map(instof(Mapping), obj.values())) and
          all(map(lambda v: all(map(instof(str), v.keys())), obj.values())) and
          # ... and has at least one leaf.
          bool(obj) and any(map(bool, obj.values())))


class Schedule(object):
  """Hyperparameter scheduling objects."""


class CosineDecay(Schedule):
  """Cosine decay."""

  def __init__(self, min_val, max_val, num_steps):
    self.min_val = min_val
    self.max_val = max_val
    self.num_steps = num_steps

  def __call__(self, itr):
    cos = (1 + jnp.cos(jnp.pi * itr / self.num_steps))
    return 0.5 * (self.max_val - self.min_val) * cos + self.min_val


class WarmupCosineDecay(Schedule):
  """Cosine decay with linear warmup."""

  def __init__(self, start_val, min_val, max_val, num_steps, warmup_steps):
    self.start_val = start_val
    self.min_val = min_val
    self.max_val = max_val
{'fraction_non_alphanumeric': 0.10937058989556872, 'fraction_numerical': 0.003245836861416878, 'mean_word_length': 3.5151322077094616, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3658860', 'n_tokens_mistral': 4156, 'n_tokens_neox': 4040, 'n_words': 1356}