text: #*
  # @file Different utility functions
  # Copyright (c) Zhewei Yao, Amir Gholami, Sheng Shen
  # All rights reserved.
  # This file is part of AdaHessian library.
  #
  # AdaHessian is free software: you can redistribute it and/or modify
  # it under the terms of the GNU General Public License as published by
  # the Free Software Foundation, either version 3 of the License, or
  # (at your option) any later version.
  #
  # AdaHessian is distributed in the hope that it will be useful,
  # but WITHOUT ANY WARRANTY; without even the implied warranty of
  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  # GNU General Public License for more details.
  #
  # You should have received a copy of the GNU General Public License
  # along with adahessian.  If not, see <http://www.gnu.org/licenses/>.
  #*
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import ops
  from tensorflow.python.keras import backend_config
  from tensorflow.python.keras.optimizer_v2 import optimizer_v2
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import state_ops
  from tensorflow.python.training import training_ops
  from tensorflow.python.util.tf_export import keras_export
  
  
  import abc
  import contextlib
  import functools
  
  import six
  import tensorflow as tf
  import numpy as np
  from tensorflow.python.distribute import distribution_strategy_context as distribute_ctx
  from tensorflow.python.distribute import parameter_server_strategy
  from tensorflow.python.distribute import reduce_util as ds_reduce_util
  from tensorflow.python.distribute import values as ds_values
  from tensorflow.python.eager import backprop
  from tensorflow.python.eager import context
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_util
  from tensorflow.python.keras import backend
  from tensorflow.python.keras import initializers
  from tensorflow.python.keras.engine import base_layer_utils
  from tensorflow.python.keras.optimizer_v2 import learning_rate_schedule
  from tensorflow.python.keras.utils import generic_utils
  from tensorflow.python.keras.utils import tf_utils
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import clip_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import gradients
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import resource_variable_ops
  from tensorflow.python.ops import variables as tf_variables
  from tensorflow.python.platform import tf_logging as logging
  from tensorflow.python.saved_model import revived_types
  from tensorflow.python.training.tracking import base as trackable
  from tensorflow.python.training.tracking import tracking
  from tensorflow.python.util import nest
  from tensorflow.python.util import tf_inspect
  from tensorflow.python.util.tf_export import keras_export
  
  class AdaHessian(optimizer_v2.OptimizerV2):
  
      _HAS_AGGREGATE_GRAD = True
  
      def __init__(self,
                 learning_rate=0.1,
                 beta_1=0.9,
                 beta_2=0.999,
                 epsilon=1e-4,
                 weight_decay = 0.,
                 hessian_power=1.0,
                 name='AdaHessian',
                 average_size_1d=None,
                 average_size_2d=None,
                 average_size_3d=-1,
                 average_size_4d=-1,
                 **kwargs):
          """Construct a new AdaHessian optimizer.
          Args:
              learning_rate: A `Tensor`, floating point value, or a schedule that is a
              `tf.keras.optimizers.schedules.LearningRateSchedule`, or a callable that
              takes no arguments and returns the actual value to use, The learning
              rate. Defaults to 0.1.
              beta_1: A float value or a constant float tensor, or a callable that takes
              no arguments and returns the actual value to use. The exponential decay
              rate for the 1st moment estimates. Defaults to 0.9.
              beta_2: A float value or a constant float tensor, or a callable that takes
              no arguments and returns the actual value to use, The exponential decay
              rate for the 2nd moment estimates. Defaults to 0.999.
              epsilon: A small constant for numerical stability. This epsilon is
              "epsilon hat" in the Kingma and Ba paper (in the formula just before
              Section 2.1), not the epsilon in Algorithm 1 of the paper. Defaults to
              1e-7.
              weight_decay: We are using AdamW's weight decay scheme. Defaults to 0.
              name: Optional name for the operations created when applying gradients.
              Defaults to "Adam".
              hessian_power: Hessian power to control the optimizer more similar to first/second 
              order method (default: 1). You can also try 0.5. For some tasks we found this 
              to result in better performance.
              **kwargs: keyword arguments. Allowed to be {`clipnorm`, `clipvalue`, `lr`,
              `decay`}. `clipnorm` is clip gradients by norm; `clipvalue` is clip
              gradients by value, `decay` is included for backward compatibility to
              allow time inverse decay of learning rate. `lr` is included for backward
              compatibility, recommended to use `learning_rate` instead.
              # average_size_{1,2,3,4}d: 
                          None: use no spatial averaging
                          -1: use suggested spatial averaging (recommended for conv kernels)
                          >= 1: use customized size
          """
  
          super(AdaHessian, self).__init__(name, **kwargs)
          self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
          self._set_hyper('decay', self._initial_decay)
          self._set_hyper('beta_1', beta_1)
          self._set_hyper('beta_2', beta_2)
          self.epsilon = epsilon or backend_config.epsilon()
          self.weight_decay = weight_decay
          self.hessian_power = hessian_power
          self.average_size_1d = average_size_1d
          self.average_size_2d = average_size_2d
          self.average_size_3d = average_size_3d
          self.average_size_4d = average_size_4d
  
      def _create_slots(self, var_list):
          # Create slots for the first and second moments.
          # Separate for-loops to respect the ordering of slot variables from v1.
          for var in var_list:
              self.add_slot(var, 'm')
          for var in var_list:
              self.add_slot(var, 'v')
  
      def _prepare_local(self, var_device, var_dtype, apply_state):
          super(AdaHessian, self)._prepare_local(var_device, var_dtype, apply_state)
  
          local_step = math_ops.cast(self.iterations + 1, var_dtype)
          beta_1_t = array_ops.identity(self._get_hyper('beta_1', var_dtype))
          beta_2_t = array_ops.identity(self._get_hyper('beta_2', var_dtype))
          beta_1_power = math_ops.pow(beta_1_t, local_step)
          beta_2_power = math_ops.pow(beta_2_t, local_step)
          lr = (
              apply_state[(var_device, var_dtype)]['lr_t'] *
              (math_ops.sqrt(1 - beta_2_power) / (1 - beta_1_power)))
          apply_state[(var_device, var_dtype)].update(
              dict(
                  lr=lr,
                  epsilon=ops.convert_to_tensor_v2(self.epsilon, var_dtype),
                  beta_1_t=beta_1_t,
                  beta_1_power=beta_1_power,
                  one_minus_beta_1_t=1 - beta_1_t,
                  beta_2_t=beta_2_t,
                  beta_2_power=beta_2_power,
                  one_minus_beta_2_t=1 - beta_2_t))
  
      def set_weights(self, weights):
          params = self.weights
          # If the weights are generated by Keras V1 optimizer, it includes vhats
          # even without amsgrad, i.e, V1 optimizer has 3x + 1 variables, while V2
          # optimizer has 2x + 1 variables. Filter vhats out for compatibility.
          num_vars = int((len(params) - 1) / 2)
          if len(weights) == 3 * num_vars + 1:
              weights = weights[:len(params)]
          super(AdaHessian, self).set_weights(weights)
  
  
      def get_gradients_hessian(self, loss, params):
          """Returns gradients and Hessian of `loss` with respect to `params`.
          Arguments:
              loss: Loss tensor.
              params: List of variables.
          Returns:
              List of gradient and Hessian tensors.
          Raises:
              ValueError: In case any gradient cannot be computed (e.g. if gradient
              function not implemented).
          """
          params = nest.flatten(params)
          with backend.get_graph().as_default(), backend.name_scope(self._name +
                                                                      "/gradients"):
              grads = gradients.gradients(loss, params)
              for grad, param in zip(grads, params):
                  if grad is None:
                      raise ValueError("Variable {} has `None` for gradient. "
                                      "Please make sure that all of your ops have a "
                                      "gradient defined (i.e. are differentiable). "
                                      "Common ops without gradient: "
                                      "K.argmax, K.round, K.eval.".format(param))
  
              # WARNING: for now we do not support gradient clip
              # grads = self._clip_gradients(grads)
  
              v = [np.random.uniform(0, 1, size = p.shape) for p in params]
              for vi in v:
                  vi[ vi < 0.5] =  -1 
                  vi[ vi >= 0.5] =  1 
              v = [tf.convert_to_tensor(vi, dtype = tf.dtypes.float32) for vi in v]
  
              vprod = tf.reduce_sum([ tf.reduce_sum(vi * grad) for vi, grad in zip(v, grads)])
  
              Hv = gradients.gradients(vprod, params)
  
              Hd = [ tf.abs(Hvi * vi) for Hvi, vi in zip(Hv, v)]
  
          return grads, Hd
  
      def _filter_grads_hessian(self, grads_hessian_and_vars):
          """Filter out iterable with grad equal to None."""
          grads_hessian_and_vars = tuple(grads_hessian_and_vars)
          if not grads_hessian_and_vars:
              return grads_hessian_and_vars
          filtered = []
          vars_with_empty_grads = []
          for grad, hessian, var in grads_hessian_and_vars:
              if grad is None:
                  vars_with_empty_grads.append(var)
              else:
                  filtered.append((grad, hessian, var))
          filtered = tuple(filtered)
  
          if not filtered:
              raise ValueError("No gradients provided for any variable: %s." %
                              ([v.name for _, v in grads_and_vars],))
          if vars_with_empty_grads:
              logging.warning(
                  ("Gradients do not exist for variables %s when minimizing the loss."),
                  ([v.name for v in vars_with_empty_grads]))
          return filtered
  
      def apply_gradients_hessian(self,
                        grads_hessian_and_vars,
                        name=None,
                        experimental_aggregate_gradients=True):
          grads_hessian_and_vars = self._filter_grads_hessian(grads_hessian_and_vars)
          var_list = [v for (_, _, v) in grads_hessian_and_vars]
  
          with backend.name_scope(self._name):
              # Create iteration if necessary.
              with ops.init_scope():
                  self._create_all_weights(var_list)
  
          if not grads_hessian_and_vars:
              # Distribution strategy does not support reducing an empty list of
              # gradients
              return control_flow_ops.no_op()
  
          if distribute_ctx.in_cross_replica_context():
              raise RuntimeError(
                  "`apply_gradients() cannot be called in cross-replica context. "
                  "Use `tf.distribute.Strategy.run` to enter replica "
                  "context.")
  
          strategy = distribute_ctx.get_strategy()
          if (not experimental_aggregate_gradients and strategy and isinstance(
              strategy.extended,
              parameter_server_strategy.ParameterServerStrategyExtended)):
              raise NotImplementedError(
                  "`experimental_aggregate_gradients=False is not supported for "
                  "ParameterServerStrategy and CentralStorageStrategy")
  
          apply_state = self._prepare(var_list)
          if experimental_aggregate_gradients:
              reduced_grads, reduced_hessian = self._aggregate_gradients_hessian(grads_hessian_and_vars)
              var_list = [v for _, _, v in grads_hessian_and_vars]
              grads_hessian_and_vars = list(zip(reduced_grads, reduced_hessian, var_list))
  
  
          return distribute_ctx.get_replica_context().merge_call(
              functools.partial(self._distributed_apply, apply_state=apply_state),
              args=(grads_hessian_and_vars,),
              kwargs={
                  "name": name,
              })
  
      def _aggregate_gradients_hessian(self, grads_hessian_and_vars):
          """Returns all-reduced gradients.
          Args:
          grads_and_vars: List of (gradient, hessian, variable) pairs.
          Returns:
          Two lists of all-reduced gradients and Hessian.
          """
          grads_hessian_and_vars = list(grads_hessian_and_vars)
          filtered_grads_hessian_and_vars = self._filter_grads_hessian(grads_hessian_and_vars)
  
          # split the list so that we can use the all_recude_fn
          filtered_grads_and_vars = tuple([(g, v) for (g, h, v) in filtered_grads_hessian_and_vars])
          filtered_hessian_and_vars = tuple([(h, v) for (g, h, v) in filtered_grads_hessian_and_vars])
  
  
          def all_reduce_fn(distribution, grads_hessian_and_vars):
              # WARNING: this ReduceOp.SUM can only support two entries, for now we have three.
              # So far now, we do it for two steps to make life easier.
              return distribution.extended.batch_reduce_to(
                  ds_reduce_util.ReduceOp.SUM, grads_hessian_and_vars)
  
          if filtered_grads_hessian_and_vars:
              reduced_part1 = distribute_ctx.get_replica_context().merge_call(
                  all_reduce_fn, args=(filtered_grads_and_vars,))
              reduced_part2 = distribute_ctx.get_replica_context().merge_call(
                  all_reduce_fn, args=(filtered_hessian_and_vars,))
          else:
              reduced = []
  
          # Copy 'reduced' but add None gradients back in
          reduced_with_nones_grads = []
          reduced_with_nones_hessian = []
  
          reduced_pos = 0
          for g, h, _ in grads_hessian_and_vars:
              if g is None:
                  reduced_with_nones_grads.append( None )
                  reduced_with_nones_hessian.append( None )
              else:
                  reduced_with_nones_grads.append(reduced_part1[reduced_pos])
                  reduced_with_nones_hessian.append(reduced_part2[reduced_pos])
                  reduced_pos += 1
  
          return reduced_with_nones_grads, reduced_with_nones_hessian
  
  
      @def_function.function(experimental_compile=True)
      def _resource_apply_dense(self, grad, hess, var, apply_state=None):
          var_device, var_dtype = var.device, var.dtype.base_dtype
          coefficients = ((apply_state or {}).get((var_device, var_dtype)) or
                          self._fallback_apply_state(var_device, var_dtype))
  
          m = self.get_slot(var, 'm')
          v = self.get_slot(var, 'v')
  
          m.assign_add((grad - m) * (1 - coefficients['beta_1_t']))
          # this part need to be changed for spatial averaging
  
          if len(v.shape) == 1:
              resize = self.average_size_1d
          elif len(v.shape) == 2:
              resize = self.average_size_2d
          elif len(v.shape) == 3:
              resize = self.average_size_3d
          elif len(v.shape) == 4:
              resize = self.average_size_4d
          else:
              raise Exception('You need to define the spatial average size by yourself!')
  
          if resize == None:
              v.assign_add((math_ops.square(hess) - v) * (1 - coefficients['beta_2_t']))
          elif resize == -1:
              if len(v.shape) == 1:
                  v.assign_add((math_ops.square(hess) - v) * (1 - coefficients['beta_2_t']))
              elif len(v.shape) == 2:
                  hess_average = tf.reduce_mean(hess, [0], keepdims=True)
                  v.assign_add((math_ops.square(hess_average) - v) * (1 - coefficients['beta_2_t'])) 
              elif len(v.shape) == 3:
                  hess_average = tf.reduce_mean(hess, [0], keepdims=True)
                  v.assign_add((math_ops.square(hess_average) - v) * (1 - coefficients['beta_2_t'])) 
              elif len(v.shape) == 4:
                  hess_average = tf.reduce_mean(hess, [0, 1], keepdims=True)
                  v.assign_add((math_ops.square(hess_average) - v) * (1 - coefficients['beta_2_t'])) 
          else:
              if resize <= 0:
                  raise Exception('You need to define the spatial average size >= 1!')
              hess_average = tf.reshape(hess, [resize, -1])
              hess_average = tf.reduce_mean(hess_average, [0])
              hess_average = tf.repeat(hess_average, resize)
              hess_average = tf.reshape(hess_average, v.shape)
              v.assign_add((math_ops.square(hess_average) - v) * (1 - coefficients['beta_2_t']))
          
          bias_correct1 = 1 - coefficients['beta_1_power']
          bias_correct2 = 1 - coefficients['beta_2_power']
  
          if self.weight_decay != 0:
              var.assign_sub(coefficients['lr_t'] * self.weight_decay * var)
  
          denom = np.power(math_ops.sqrt(v / bias_correct2), self.hessian_power) + coefficients['epsilon']
  
          var.assign_sub( coefficients['lr_t'] * m / bias_correct1 / denom  )
  
      @def_function.function(experimental_compile=True)
      def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
          raise Exception('For now, we do not support sparse update yet.')
  
      def get_config(self):
          config = super(AdaHessian, self).get_config()
          config.update({
              'learning_rate': self._serialize_hyperparameter('learning_rate'),
              'decay': self._serialize_hyperparameter('decay'),
              'beta_1': self._serialize_hyperparameter('beta_1'),
              'beta_2': self._serialize_hyperparameter('beta_2'),
              'epsilon': self.epsilon,
              'weight_decay': self.weight_decay
          })
          return config
  
      def _distributed_apply(self, distribution, grads_hessian_and_vars, name, apply_state):
          """`apply_gradients` using a `DistributionStrategy`."""
  
          def apply_grad_to_update_var(var, grad, hess):
              """Apply gradient to variable."""
              if isinstance(var, ops.Tensor):
                  raise NotImplementedError("Trying to update a Tensor ", var)
  
              apply_kwargs = {}
  
              if "apply_state" in self._dense_apply_args:
                  apply_kwargs["apply_state"] = apply_state
              update_op = self._resource_apply_dense(grad, hess, var, **apply_kwargs)
  
              if var.constraint is not None:
                  with ops.control_dependencies([update_op]):
                      return var.assign(var.constraint(var))
              else:
                  return update_op
  
          eagerly_outside_functions = ops.executing_eagerly_outside_functions()
          update_ops = []
          with ops.name_scope(name or self._name, skip_on_eager=True):
              for grad, hess, var in grads_hessian_and_vars:
  
                  def _assume_mirrored(grad, hess):
                      if isinstance(grad, ds_values.PerReplica):
                          return ds_values.Mirrored(grad.values), ds_values.Mirrored(hess.values)
                      return grad, hess
  
                  grad, hess = nest.map_structure(_assume_mirrored, grad, hess)
                  # Colocate the update with variables to avoid unnecessary communication
                  # delays. See b/136304694.
                  with distribution.extended.colocate_vars_with(var):
                      with ops.name_scope("update" if eagerly_outside_functions else
                                    "update_" + var.op.name, skip_on_eager=True):
                          update_ops.extend(distribution.extended.update(
                                  var, apply_grad_to_update_var, args=(grad, hess), group=False))
  
              any_symbolic = any(isinstance(i, ops.Operation) or
                               tf_utils.is_symbolic_tensor(i) for i in update_ops)
              if not context.executing_eagerly() or any_symbolic:
                  # If the current context is graph mode or any of the update ops are
                  # symbolic then the step update should be carried out under a graph
                  # context. (eager updates execute immediately)
                  with ops._get_graph_from_inputs(update_ops).as_default():  # pylint: disable=protected-access
                      with ops.control_dependencies(update_ops):
                          return self._iterations.assign_add(1, read_value=False)
  
              return self._iterations.assign_add(1)
  