text: import re
  from functools import reduce
  
  import numpy as np
  import tensorflow as tf
  from stable_baselines.acktr.kfac_utils import (detect_min_val, factor_reshape,
                                                 gmatmul)
  
  KFAC_OPS = ['MatMul', 'Conv2D', 'BiasAdd']
  KFAC_DEBUG = False
  
  
  class KfacOptimizer:
      def __init__(
          self,
          learning_rate=0.01,
          momentum=0.9,
          clip_kl=0.01,
          kfac_update=2,
          stats_accum_iter=60,
          full_stats_init=False,
          cold_iter=100,
          cold_lr=None,
          async_eigen_decomp=False,
          async_stats=False,
          epsilon=1e-2,
          stats_decay=0.95,
          blockdiag_bias=False,
          channel_fac=False,
          factored_damping=False,
          approx_t2=False,
          use_float64=False,
          weight_decay_dict=None,
          max_grad_norm=0.5,
          verbose=1,
      ):
          self.max_grad_norm = max_grad_norm
          self._lr = learning_rate
          self._momentum = momentum
          self._clip_kl = clip_kl
          self._channel_fac = channel_fac
          self._kfac_update = kfac_update
          self._async_eigen_decomp = async_eigen_decomp
          self._async_stats = async_stats
          self._epsilon = epsilon
          self._stats_decay = stats_decay
          self._blockdiag_bias = blockdiag_bias
          self._approx_t2 = approx_t2
          self._use_float64 = use_float64
          self._factored_damping = factored_damping
          self._cold_iter = cold_iter
          self.verbose = verbose
          if cold_lr is None:
              self._cold_lr = self._lr
          else:
              self._cold_lr = cold_lr
          self._stats_accum_iter = stats_accum_iter
          if weight_decay_dict is None:
              weight_decay_dict = {}
          self._weight_decay_dict = weight_decay_dict
          self._diag_init_coeff = 0.0
          self._full_stats_init = full_stats_init
          if not self._full_stats_init:
              self._stats_accum_iter = self._cold_iter
          self.sgd_step = tf.Variable(0, name='KFAC/sgd_step', trainable=False)
          self.global_step = tf.Variable(0, name='KFAC/global_step', trainable=False)
          self.cold_step = tf.Variable(0, name='KFAC/cold_step', trainable=False)
          self.factor_step = tf.Variable(0, name='KFAC/factor_step', trainable=False)
          self.stats_step = tf.Variable(0, name='KFAC/stats_step', trainable=False)
          self.v_f_v = tf.Variable(0.0, name='KFAC/vFv', trainable=False)
          self.factors = {}
          self.param_vars = []
          self.stats = {}
          self.stats_eigen = {}
          self._update_stats_op = None
  
      def get_factors(self, gradients, varlist):
          default_graph = tf.compat.v1.get_default_graph()
          factor_tensors = {}
          fprop_tensors = []
          bprop_tensors = []
          op_types = []
  
          def _search_factors(gradient, graph):
              bprop_op = gradient.op
              bprop_op_name = bprop_op.name
              b_tensors = []
              f_tensors = []
              if 'AddN' in bprop_op_name:
                  factors = []
                  for grad in gradient.op.inputs:
                      factors.append(_search_factors(grad, graph))
                  op_names = [_item['opName'] for _item in factors]
                  if self.verbose > 1:
                      print(gradient.name)
                      print(op_names)
                      print(len(np.unique(op_names)))
                  assert (
                      len(np.unique(op_names)) == 1
                  ), 'Error: {} is shared among different computation OPs'.format(
                      gradient.name
                  )
                  b_tensors = reduce(
                      lambda x, y: x + y, [_item['bpropFactors'] for _item in factors]
                  )
                  if len(factors[0]['fpropFactors']) > 0:
                      f_tensors = reduce(
                          lambda x, y: x + y, [_item['fpropFactors'] for _item in factors]
                      )
                  fprop_op_name = op_names[0]
                  fprop_op = factors[0]['op']
              else:
                  fprop_op_match = re.search(
                      'gradientsSampled(_[0-9]+|)/(.+?)_grad', bprop_op_name
                  )
                  assert fprop_op_match is not None
                  fprop_op_name = fprop_op_match.group(2)
                  fprop_op = graph.get_operation_by_name(fprop_op_name)
                  if fprop_op.op_def.name in KFAC_OPS:
                      b_tensor = [
                          _i for _i in bprop_op.inputs if 'gradientsSampled' in _i.name
                      ][-1]
                      b_tensor_shape = fprop_op.outputs[0].get_shape()
                      if b_tensor.get_shape()[0].value is None:
                          b_tensor.set_shape(b_tensor_shape)
                      b_tensors.append(b_tensor)
                      if fprop_op.op_def.name == 'BiasAdd':
                          f_tensors = []
                      else:
                          f_tensors.append(
                              [
                                  _i
                                  for _i in fprop_op.inputs
                                  if param.op.name not in _i.name
                              ][0]
                          )
                      fprop_op_name = fprop_op.op_def.name
                  else:
                      b_inputs_list = [
                          _i
                          for _i in bprop_op.inputs[0].op.inputs
                          if 'gradientsSampled' in _i.name
                          if 'Shape' not in _i.name
                      ]
                      if len(b_inputs_list) > 0:
                          b_tensor = b_inputs_list[0]
                          if b_tensor.get_shape():
                              b_tensor_shape = fprop_op.outputs[0].get_shape()
                              if (
                                  len(b_tensor.get_shape()) > 0
                                  and b_tensor.get_shape()[0].value is None
                              ):
                                  b_tensor.set_shape(b_tensor_shape)
                              b_tensors.append(b_tensor)
                      fprop_op_name = 'UNK-' + fprop_op.op_def.name
                      op_types.append(fprop_op_name)
              return {
                  'opName': fprop_op_name,
                  'op': fprop_op,
                  'fpropFactors': f_tensors,
                  'bpropFactors': b_tensors,
              }
  
          for _grad, param in zip(gradients, varlist):
              if KFAC_DEBUG:
                  print(('get factor for ' + param.name))
              found_factors = _search_factors(_grad, default_graph)
              factor_tensors[param] = found_factors
          for param in varlist:
              factor_tensors[param]['assnWeights'] = None
              factor_tensors[param]['assnBias'] = None
          for param in varlist:
              if factor_tensors[param]['opName'] == 'BiasAdd':
                  factor_tensors[param]['assnWeights'] = None
                  for item in varlist:
                      if len(factor_tensors[item]['bpropFactors']) > 0:
                          if (
                              set(factor_tensors[item]['bpropFactors'])
                              == set(factor_tensors[param]['bpropFactors'])
                          ) and (len(factor_tensors[item]['fpropFactors']) > 0):
                              factor_tensors[param]['assnWeights'] = item
                              factor_tensors[item]['assnBias'] = param
                              factor_tensors[param]['bpropFactors'] = factor_tensors[
                                  item
                              ]['bpropFactors']
          for key in ['fpropFactors', 'bpropFactors']:
              for i, param in enumerate(varlist):
                  if len(factor_tensors[param][key]) > 0:
                      if (key + '_concat') not in factor_tensors[param]:
                          tensor = factor_tensors[param][key][0]
                          name_scope = tensor.name.split(':')[0]
                          with tf.compat.v1.name_scope(name_scope):
                              factor_tensors[param][key + '_concat'] = tf.concat(
                                  factor_tensors[param][key], 0
                              )
                  else:
                      factor_tensors[param][key + '_concat'] = None
                  for _, param2 in enumerate(varlist[(i + 1) :]):
                      if (len(factor_tensors[param][key]) > 0) and (
                          set(factor_tensors[param2][key])
                          == set(factor_tensors[param][key])
                      ):
                          factor_tensors[param2][key] = factor_tensors[param][key]
                          factor_tensors[param2][key + '_concat'] = factor_tensors[param][
                              key + '_concat'
                          ]
          if KFAC_DEBUG:
              for items in zip(varlist, fprop_tensors, bprop_tensors, op_types):
                  print((items[0].name, factor_tensors[item]))
          self.factors = factor_tensors
          return factor_tensors
  
      def get_stats(self, factors, varlist):
          if len(self.stats) == 0:
              with tf.device('/cpu'):
                  tmp_stats_cache = {}
                  for var in varlist:
                      bprop_factor = factors[var]['bpropFactors_concat']
                      op_type = factors[var]['opName']
                      if op_type == 'Conv2D':
                          operator_height = bprop_factor.get_shape()[1]
                          operator_width = bprop_factor.get_shape()[2]
                          if (
                              operator_height == 1
                              and operator_width == 1
                              and self._channel_fac
                          ):
                              var_assn_bias = factors[var]['assnBias']
                              if var_assn_bias:
                                  factors[var]['assnBias'] = None
                                  factors[var_assn_bias]['assnWeights'] = None
                  for var in varlist:
                      fprop_factor = factors[var]['fpropFactors_concat']
                      bprop_factor = factors[var]['bpropFactors_concat']
                      op_type = factors[var]['opName']
                      self.stats[var] = {
                          'opName': op_type,
                          'fprop_concat_stats': [],
                          'bprop_concat_stats': [],
                          'assnWeights': factors[var]['assnWeights'],
                          'assnBias': factors[var]['assnBias'],
                      }
                      if fprop_factor is not None:
                          if fprop_factor not in tmp_stats_cache:
                              if op_type == 'Conv2D':
                                  kernel_height = var.get_shape()[0]
                                  kernel_width = var.get_shape()[1]
                                  n_channels = fprop_factor.get_shape()[-1]
                                  operator_height = bprop_factor.get_shape()[1]
                                  operator_width = bprop_factor.get_shape()[2]
                                  if (
                                      operator_height == 1
                                      and operator_width == 1
                                      and self._channel_fac
                                  ):
                                      fprop_factor2_size = kernel_height * kernel_width
                                      slot_fprop_factor_stats2 = tf.Variable(
                                          tf.linalg.tensor_diag(
                                              tf.ones([fprop_factor2_size])
                                          )
                                          * self._diag_init_coeff,
                                          name='KFAC_STATS/' + fprop_factor.op.name,
                                          trainable=False,
                                      )
                                      self.stats[var]['fprop_concat_stats'].append(
                                          slot_fprop_factor_stats2
                                      )
                                      fprop_factor_size = n_channels
                                  else:
                                      fprop_factor_size = (
                                          kernel_height * kernel_width * n_channels
                                      )
                              else:
                                  fprop_factor_size = fprop_factor.get_shape()[-1]
                              if not self._blockdiag_bias and self.stats[var]['assnBias']:
                                  fprop_factor_size += 1
                              slot_fprop_factor_stats = tf.Variable(
                                  tf.linalg.tensor_diag(tf.ones([fprop_factor_size]))
                                  * self._diag_init_coeff,
                                  name='KFAC_STATS/' + fprop_factor.op.name,
                                  trainable=False,
                              )
                              self.stats[var]['fprop_concat_stats'].append(
                                  slot_fprop_factor_stats
                              )
                              if op_type != 'Conv2D':
                                  tmp_stats_cache[fprop_factor] = self.stats[var][
                                      'fprop_concat_stats'
                                  ]
                          else:
                              self.stats[var]['fprop_concat_stats'] = tmp_stats_cache[
                                  fprop_factor
                              ]
                      if bprop_factor is not None:
                          if not (
                              (not self._blockdiag_bias)
                              and self.stats[var]['assnWeights']
                          ):
                              if bprop_factor not in tmp_stats_cache:
                                  slot_bprop_factor_stats = tf.Variable(
                                      tf.linalg.tensor_diag(
                                          tf.ones([bprop_factor.get_shape()[-1]])
                                      )
                                      * self._diag_init_coeff,
                                      name='KFAC_STATS/' + bprop_factor.op.name,
                                      trainable=False,
                                  )
                                  self.stats[var]['bprop_concat_stats'].append(
                                      slot_bprop_factor_stats
                                  )
                                  tmp_stats_cache[bprop_factor] = self.stats[var][
                                      'bprop_concat_stats'
                                  ]
                              else:
                                  self.stats[var]['bprop_concat_stats'] = tmp_stats_cache[
                                      bprop_factor
                                  ]
          return self.stats
  
      def compute_and_apply_stats(self, loss_sampled, var_list=None):
          varlist = var_list
          if varlist is None:
              varlist = tf.compat.v1.trainable_variables()
          stats = self.compute_stats(loss_sampled, var_list=varlist)
          return self.apply_stats(stats)
  
      def compute_stats(self, loss_sampled, var_list=None):
          varlist = var_list
          if varlist is None:
              varlist = tf.compat.v1.trainable_variables()
          gradient_sampled = tf.gradients(
              ys=loss_sampled, xs=varlist, name='gradientsSampled'
          )
          self.gradient_sampled = gradient_sampled
          gradient_sampled, varlist = zip(
              *[
                  (grad, var)
                  for (grad, var) in zip(gradient_sampled, varlist)
                  if grad is not None
              ]
          )
          factors = self.get_factors(gradient_sampled, varlist)
          stats = self.get_stats(factors, varlist)
          update_ops = []
          stats_updates = {}
          stats_updates_cache = {}
          for var in varlist:
              op_type = factors[var]['opName']
              fops = factors[var]['op']
              fprop_factor = factors[var]['fpropFactors_concat']
              fprop_stats_vars = stats[var]['fprop_concat_stats']
              bprop_factor = factors[var]['bpropFactors_concat']
              bprop_stats_vars = stats[var]['bprop_concat_stats']
              svd_factors = {}
              for stats_var in fprop_stats_vars:
                  stats_var_dim = int(stats_var.get_shape()[0])
                  if stats_var not in stats_updates_cache:
                      batch_size = tf.shape(input=fprop_factor)[0]
                      if op_type == 'Conv2D':
                          strides = fops.get_attr("strides")
                          padding = fops.get_attr("padding")
                          convkernel_size = var.get_shape()[0:3]
                          kernel_height = int(convkernel_size[0])
                          kernel_width = int(convkernel_size[1])
                          chan = int(convkernel_size[2])
                          flatten_size = int(kernel_height * kernel_width * chan)
                          operator_height = int(bprop_factor.get_shape()[1])
                          operator_width = int(bprop_factor.get_shape()[2])
                          if (
                              operator_height == 1
                              and operator_width == 1
                              and self._channel_fac
                          ):
                              if len(svd_factors) == 0:
                                  if KFAC_DEBUG:
                                      print(
                                          (
                                              'approx %s act factor with rank-1 SVD factors'
                                              % var.name
                                          )
                                      )
                                  S, U, V = tf.batch_svd(
                                      tf.reshape(
                                          fprop_factor,
                                          [-1, kernel_height * kernel_width, chan],
                                      )
                                  )
                                  sqrt_s1 = tf.expand_dims(tf.sqrt(S[:, 0, 0]), 1)
                                  patches_k = U[:, :, 0] * sqrt_s1
                                  full_factor_shape = fprop_factor.get_shape()
                                  patches_k.set_shape(
                                      [full_factor_shape[0], kernel_height * kernel_width]
                                  )
                                  patches_c = V[:, :, 0] * sqrt_s1
                                  patches_c.set_shape([full_factor_shape[0], chan])
                                  svd_factors[chan] = patches_c
                                  svd_factors[kernel_height * kernel_width] = patches_k
                              fprop_factor = svd_factors[stats_var_dim]
                          else:
                              patches = tf.image.extract_patches(
                                  fprop_factor,
                                  sizes=[1, convkernel_size[0], convkernel_size[1], 1],
                                  strides=strides,
                                  rates=[1, 1, 1, 1],
                                  padding=padding,
                              )
                              if self._approx_t2:
                                  if KFAC_DEBUG:
                                      print(('approxT2 act fisher for %s' % var.name))
                                  fprop_factor = tf.reduce_mean(
                                      input_tensor=patches, axis=[1, 2]
                                  )
                              else:
                                  fprop_factor = (
                                      tf.reshape(patches, [-1, flatten_size])
                                      / operator_height
                                      / operator_width
                                  )
                      fprop_factor_size = int(fprop_factor.get_shape()[-1])
                      if (
                          stats_var_dim == (fprop_factor_size + 1)
                          and not self._blockdiag_bias
                      ):
                          if op_type == 'Conv2D' and not self._approx_t2:
                              fprop_factor = tf.concat(
                                  [
                                      fprop_factor,
                                      tf.ones([tf.shape(input=fprop_factor)[0], 1])
                                      / operator_height
                                      / operator_width,
                                  ],
                                  1,
                              )
                          else:
                              fprop_factor = tf.concat(
                                  [
                                      fprop_factor,
                                      tf.ones([tf.shape(input=fprop_factor)[0], 1]),
                                  ],
                                  1,
                              )
                      cov = tf.matmul(
                          fprop_factor, fprop_factor, transpose_a=True
                      ) / tf.cast(batch_size, tf.float32)
                      update_ops.append(cov)
                      stats_updates[stats_var] = cov
                      if op_type != 'Conv2D':
                          stats_updates_cache[stats_var] = cov
              for stats_var in bprop_stats_vars:
                  if stats_var not in stats_updates_cache:
                      bprop_factor_shape = bprop_factor.get_shape()
                      batch_size = tf.shape(input=bprop_factor)[0]
                      chan = int(bprop_factor_shape[-1])
                      if op_type == 'Conv2D' or len(bprop_factor_shape) == 4:
                          if fprop_factor is not None:
                              if self._approx_t2:
                                  if KFAC_DEBUG:
                                      print(('approxT2 grad fisher for %s' % var.name))
                                  bprop_factor = tf.reduce_sum(
                                      input_tensor=bprop_factor, axis=[1, 2]
                                  )
                              else:
                                  bprop_factor = (
                                      tf.reshape(bprop_factor, [-1, chan])
                                      * operator_height
                                      * operator_width
                                  )
                          else:
                              if KFAC_DEBUG:
                                  print(('block diag approx fisher for %s' % var.name))
                              bprop_factor = tf.reduce_sum(
                                  input_tensor=bprop_factor, axis=[1, 2]
                              )
                      bprop_factor *= tf.cast(batch_size, tf.float32)
                      cov_b = tf.matmul(
                          bprop_factor, bprop_factor, transpose_a=True
                      ) / tf.cast(tf.shape(input=bprop_factor)[0], tf.float32)
                      update_ops.append(cov_b)
                      stats_updates[stats_var] = cov_b
                      stats_updates_cache[stats_var] = cov_b
          if KFAC_DEBUG:
              a_key = list(stats_updates.keys())[0]
              stats_updates[a_key] = tf.compat.v1.Print(
                  stats_updates[a_key],
                  [
                      tf.convert_to_tensor(value='step:'),
                      self.global_step,
                      tf.convert_to_tensor(value='computing stats'),
                  ],
              )
          self.stats_updates = stats_updates
          return stats_updates
  
      def apply_stats(self, stats_updates):
          def _update_accum_stats():
              if self._full_stats_init:
                  return tf.cond(
                      pred=tf.greater(self.sgd_step, self._cold_iter),
                      true_fn=lambda: tf.group(
                          *self._apply_stats(
                              stats_updates,
                              accumulate=True,
                              accumulate_coeff=1.0 / self._stats_accum_iter,
                          )
                      ),
                      false_fn=tf.no_op,
                  )
              else:
                  return tf.group(
                      *self._apply_stats(
                          stats_updates,
                          accumulate=True,
                          accumulate_coeff=1.0 / self._stats_accum_iter,
                      )
                  )
  
          def _update_running_avg_stats(stats_updates):
              return tf.group(*self._apply_stats(stats_updates))
  
          if self._async_stats:
              update_stats = self._apply_stats(stats_updates)
              queue = tf.queue.FIFOQueue(
                  1,
                  [item.dtype for item in update_stats],
                  shapes=[item.get_shape() for item in update_stats],
              )
              enqueue_op = queue.enqueue(update_stats)
  
              def dequeue_stats_op():
                  return queue.dequeue()
  
              self.qr_stats = tf.compat.v1.train.QueueRunner(queue, [enqueue_op])
              update_stats_op = tf.cond(
                  pred=tf.equal(queue.size(), tf.convert_to_tensor(value=0)),
                  true_fn=tf.no_op,
                  false_fn=lambda: tf.group(
                      *[
                          dequeue_stats_op(),
                      ]
                  ),
              )
          else:
              update_stats_op = tf.cond(
                  pred=tf.greater_equal(self.stats_step, self._stats_accum_iter),
                  true_fn=lambda: _update_running_avg_stats(stats_updates),
                  false_fn=_update_accum_stats,
              )
          self._update_stats_op = update_stats_op
          return update_stats_op
  
      def _apply_stats(self, stats_updates, accumulate=False, accumulate_coeff=0.0):
          update_ops = []
          for stats_var in stats_updates:
              stats_new = stats_updates[stats_var]
              if accumulate:
                  update_op = tf.compat.v1.assign_add(
                      stats_var, accumulate_coeff * stats_new, use_locking=True
                  )
              else:
                  update_op = tf.compat.v1.assign(
                      stats_var, stats_var * self._stats_decay, use_locking=True
                  )
                  update_op = tf.compat.v1.assign_add(
                      update_op, (1.0 - self._stats_decay) * stats_new, use_locking=True
                  )
              update_ops.append(update_op)
          with tf.control_dependencies(update_ops):
              stats_step_op = tf.compat.v1.assign_add(self.stats_step, 1)
          if KFAC_DEBUG:
              stats_step_op = tf.compat.v1.Print(
                  stats_step_op,
                  [
                      tf.convert_to_tensor(value='step:'),
                      self.global_step,
                      tf.convert_to_tensor(value='fac step:'),
                      self.factor_step,
                      tf.convert_to_tensor(value='sgd step:'),
                      self.sgd_step,
                      tf.convert_to_tensor(value='Accum:'),
                      tf.convert_to_tensor(value=accumulate),
                      tf.convert_to_tensor(value='Accum coeff:'),
                      tf.convert_to_tensor(value=accumulate_coeff),
                      tf.convert_to_tensor(value='stat step:'),
                      self.stats_step,
                      update_ops[0],
                      update_ops[1],
                  ],
              )
          return [
              stats_step_op,
          ]
  
      def get_stats_eigen(self, stats=None):
          if len(self.stats_eigen) == 0:
              stats_eigen = {}
              if stats is None:
                  stats = self.stats
              tmp_eigen_cache = {}
              with tf.device('/cpu:0'):
                  for var in stats:
                      for key in ['fprop_concat_stats', 'bprop_concat_stats']:
                          for stats_var in stats[var][key]:
                              if stats_var not in tmp_eigen_cache:
                                  stats_dim = stats_var.get_shape()[1].value
                                  eigen_values = tf.Variable(
                                      tf.ones([stats_dim]),
                                      name='KFAC_FAC/'
                                      + stats_var.name.split(':')[0]
                                      + '/e',
                                      trainable=False,
                                  )
                                  eigen_vectors = tf.Variable(
                                      tf.linalg.tensor_diag(tf.ones([stats_dim])),
                                      name='KFAC_FAC/'
                                      + stats_var.name.split(':')[0]
                                      + '/Q',
                                      trainable=False,
                                  )
                                  stats_eigen[stats_var] = {
                                      'e': eigen_values,
                                      'Q': eigen_vectors,
                                  }
                                  tmp_eigen_cache[stats_var] = stats_eigen[stats_var]
                              else:
                                  stats_eigen[stats_var] = tmp_eigen_cache[stats_var]
              self.stats_eigen = stats_eigen
          return self.stats_eigen
  
      def compute_stats_eigen(self):
          with tf.device('/cpu:0'):
              stats_eigen = self.stats_eigen
              computed_eigen = {}
              eigen_reverse_lookup = {}
              update_ops = []
              with tf.control_dependencies([]):
                  for stats_var in stats_eigen:
                      if stats_var not in computed_eigen:
                          eigen_decomposition = tf.linalg.eigh(stats_var)
                          eigen_values = eigen_decomposition[0]
                          eigen_vectors = eigen_decomposition[1]
                          if self._use_float64:
                              eigen_values = tf.cast(eigen_values, tf.float64)
                              eigen_vectors = tf.cast(eigen_vectors, tf.float64)
                          update_ops.append(eigen_values)
                          update_ops.append(eigen_vectors)
                          computed_eigen[stats_var] = {
                              'e': eigen_values,
                              'Q': eigen_vectors,
                          }
                          eigen_reverse_lookup[eigen_values] = stats_eigen[stats_var]['e']
                          eigen_reverse_lookup[eigen_vectors] = stats_eigen[stats_var][
                              'Q'
                          ]
              self.eigen_reverse_lookup = eigen_reverse_lookup
              self.eigen_update_list = update_ops
              if KFAC_DEBUG:
                  self.eigen_update_list = [item for item in update_ops]
                  with tf.control_dependencies(update_ops):
                      update_ops.append(
                          tf.compat.v1.Print(
                              tf.constant(0.0),
                              [tf.convert_to_tensor(value='computed factor eigen')],
                          )
                      )
          return update_ops
  
      def apply_stats_eigen(self, eigen_list):
          update_ops = []
          if self.verbose > 1:
              print(('updating %d eigenvalue/vectors' % len(eigen_list)))
          for _, (tensor, mark) in enumerate(zip(eigen_list, self.eigen_update_list)):
              stats_eigen_var = self.eigen_reverse_lookup[mark]
              update_ops.append(
                  tf.compat.v1.assign(stats_eigen_var, tensor, use_locking=True)
              )
          with tf.control_dependencies(update_ops):
              factor_step_op = tf.compat.v1.assign_add(self.factor_step, 1)
              update_ops.append(factor_step_op)
              if KFAC_DEBUG:
                  update_ops.append(
                      tf.compat.v1.Print(
                          tf.constant(0.0),
                          [tf.convert_to_tensor(value='updated kfac factors')],
                      )
                  )
          return update_ops
  
      def get_kfac_precond_updates(self, gradlist, varlist):
          v_g = 0.0
          assert len(self.stats) > 0
          assert len(self.stats_eigen) > 0
          assert len(self.factors) > 0
          counter = 0
          grad_dict = {var: grad for grad, var in zip(gradlist, varlist)}
          for grad, var in zip(gradlist, varlist):
              grad_reshape = False
              fprop_factored_fishers = self.stats[var]['fprop_concat_stats']
              bprop_factored_fishers = self.stats[var]['bprop_concat_stats']
              if (len(fprop_factored_fishers) + len(bprop_factored_fishers)) > 0:
                  counter += 1
                  grad_shape = grad.get_shape()
                  if len(grad.get_shape()) > 2:
                      kernel_width = int(grad.get_shape()[0])
                      kernel_height = int(grad.get_shape()[1])
                      n_channels = int(grad.get_shape()[2])
                      depth = int(grad.get_shape()[3])
                      if len(fprop_factored_fishers) > 1 and self._channel_fac:
                          grad = tf.reshape(
                              grad, [kernel_width * kernel_height, n_channels, depth]
                          )
                      else:
                          grad = tf.reshape(grad, [-1, depth])
                      grad_reshape = True
                  elif len(grad.get_shape()) == 1:
                      grad = tf.expand_dims(grad, 0)
                      grad_reshape = True
                  if (
                      self.stats[var]['assnBias'] is not None
                  ) and not self._blockdiag_bias:
                      var_assn_bias = self.stats[var]['assnBias']
                      grad = tf.concat(
                          [grad, tf.expand_dims(grad_dict[var_assn_bias], 0)], 0
                      )
                  eig_vals = []
                  for idx, stats in enumerate(self.stats[var]['fprop_concat_stats']):
                      eigen_vectors = self.stats_eigen[stats]['Q']
                      eigen_values = detect_min_val(
                          self.stats_eigen[stats]['e'], var, name='act', debug=KFAC_DEBUG
                      )
                      eigen_vectors, eigen_values = factor_reshape(
                          eigen_vectors, eigen_values, grad, fac_idx=idx, f_type='act'
                      )
                      eig_vals.append(eigen_values)
                      grad = gmatmul(
                          eigen_vectors, grad, transpose_a=True, reduce_dim=idx
                      )
                  for idx, stats in enumerate(self.stats[var]['bprop_concat_stats']):
                      eigen_vectors = self.stats_eigen[stats]['Q']
                      eigen_values = detect_min_val(
                          self.stats_eigen[stats]['e'], var, name='grad', debug=KFAC_DEBUG
                      )
                      eigen_vectors, eigen_values = factor_reshape(
                          eigen_vectors, eigen_values, grad, fac_idx=idx, f_type='grad'
                      )
                      eig_vals.append(eigen_values)
                      grad = gmatmul(
                          grad, eigen_vectors, transpose_b=False, reduce_dim=idx
                      )
                  weight_decay_coeff = 0.0
                  if var in self._weight_decay_dict:
                      weight_decay_coeff = self._weight_decay_dict[var]
                      if KFAC_DEBUG:
                          print(
                              (
                                  'weight decay coeff for %s is %f'
                                  % (var.name, weight_decay_coeff)
                              )
                          )
                  if self._factored_damping:
                      if KFAC_DEBUG:
                          print(('use factored damping for %s' % var.name))
                      coeffs = 1.0
                      num_factors = len(eig_vals)
                      if len(eig_vals) == 1:
                          damping = self._epsilon + weight_decay_coeff
                      else:
                          damping = tf.pow(
                              self._epsilon + weight_decay_coeff, 1.0 / num_factors
                          )
                      eig_vals_tnorm_avg = [
                          tf.reduce_mean(input_tensor=tf.abs(e)) for e in eig_vals
                      ]
                      for eigen_val, e_tnorm in zip(eig_vals, eig_vals_tnorm_avg):
                          eig_tnorm_neg_list = [
                              item for item in eig_vals_tnorm_avg if item != e_tnorm
                          ]
                          if len(eig_vals) == 1:
                              adjustment = 1.0
                          elif len(eig_vals) == 2:
                              adjustment = tf.sqrt(e_tnorm / eig_tnorm_neg_list[0])
                          else:
                              eig_tnorm_neg_list_prod = reduce(
                                  lambda x, y: x * y, eig_tnorm_neg_list
                              )
                              adjustment = tf.pow(
                                  tf.pow(e_tnorm, num_factors - 1.0)
                                  / eig_tnorm_neg_list_prod,
                                  1.0 / num_factors,
                              )
                          coeffs *= eigen_val + adjustment * damping
                  else:
                      coeffs = 1.0
                      damping = self._epsilon + weight_decay_coeff
                      for eigen_val in eig_vals:
                          coeffs *= eigen_val
                      coeffs += damping
                  grad /= coeffs
                  for idx, stats in enumerate(self.stats[var]['fprop_concat_stats']):
                      eigen_vectors = self.stats_eigen[stats]['Q']
                      grad = gmatmul(
                          eigen_vectors, grad, transpose_a=False, reduce_dim=idx
                      )
                  for idx, stats in enumerate(self.stats[var]['bprop_concat_stats']):
                      eigen_vectors = self.stats_eigen[stats]['Q']
                      grad = gmatmul(
                          grad, eigen_vectors, transpose_b=True, reduce_dim=idx
                      )
                  if (
                      self.stats[var]['assnBias'] is not None
                  ) and not self._blockdiag_bias:
                      var_assn_bias = self.stats[var]['assnBias']
                      c_plus_one = int(grad.get_shape()[0])
                      grad_assn_bias = tf.reshape(
                          tf.slice(grad, begin=[c_plus_one - 1, 0], size=[1, -1]),
                          var_assn_bias.get_shape(),
                      )
                      grad_assn_weights = tf.slice(
                          grad, begin=[0, 0], size=[c_plus_one - 1, -1]
                      )
                      grad_dict[var_assn_bias] = grad_assn_bias
                      grad = grad_assn_weights
                  if grad_reshape:
                      grad = tf.reshape(grad, grad_shape)
                  grad_dict[var] = grad
          if self.verbose > 1:
              print(('projecting %d gradient matrices' % counter))
          for grad_1, var in zip(gradlist, varlist):
              grad = grad_dict[var]
              if KFAC_DEBUG:
                  print(('apply clipping to %s' % var.name))
                  tf.compat.v1.Print(
                      grad,
                      [tf.sqrt(tf.reduce_sum(input_tensor=tf.pow(grad, 2)))],
                      "Euclidean norm of new grad",
                  )
              local_vg = tf.reduce_sum(input_tensor=grad * grad_1 * (self._lr * self._lr))
              v_g += local_vg
          if KFAC_DEBUG:
              print('apply vFv clipping')
          scaling = tf.minimum(1.0, tf.sqrt(self._clip_kl / v_g))
          if KFAC_DEBUG:
              scaling = tf.compat.v1.Print(
                  scaling,
                  [
                      tf.convert_to_tensor(value='clip: '),
                      scaling,
                      tf.convert_to_tensor(value=' vFv: '),
                      v_g,
                  ],
              )
          with tf.control_dependencies([tf.compat.v1.assign(self.v_f_v, v_g)]):
              updatelist = [grad_dict[var] for var in varlist]
              for i, item in enumerate(updatelist):
                  updatelist[i] = scaling * item
          return updatelist
  
      @classmethod
      def compute_gradients(cls, loss, var_list=None):
          varlist = var_list
          if varlist is None:
              varlist = tf.compat.v1.trainable_variables()
          gradients = tf.gradients(ys=loss, xs=varlist)
          return [(a, b) for a, b in zip(gradients, varlist)]
  
      def apply_gradients_kfac(self, grads):
          grad, varlist = list(zip(*grads))
          if len(self.stats_eigen) == 0:
              self.get_stats_eigen()
          queue_runner = None
          if self._async_eigen_decomp:
              if self.verbose > 1:
                  print('Using async eigen decomposition')
              factor_ops_dummy = self.compute_stats_eigen()
              queue = tf.queue.FIFOQueue(
                  1,
                  [item.dtype for item in factor_ops_dummy],
                  shapes=[item.get_shape() for item in factor_ops_dummy],
              )
              enqueue_op = tf.cond(
                  pred=tf.logical_and(
                      tf.equal(
                          tf.math.floormod(self.stats_step, self._kfac_update),
                          tf.convert_to_tensor(value=0),
                      ),
                      tf.greater_equal(self.stats_step, self._stats_accum_iter),
                  ),
                  true_fn=lambda: queue.enqueue(self.compute_stats_eigen()),
                  false_fn=tf.no_op,
              )
  
              def dequeue_op():
                  return queue.dequeue()
  
              queue_runner = tf.compat.v1.train.QueueRunner(queue, [enqueue_op])
          update_ops = []
          global_step_op = tf.compat.v1.assign_add(self.global_step, 1)
          update_ops.append(global_step_op)
          with tf.control_dependencies([global_step_op]):
              assert self._update_stats_op is not None
              update_ops.append(self._update_stats_op)
              dependency_list = []
              if not self._async_eigen_decomp:
                  dependency_list.append(self._update_stats_op)
              with tf.control_dependencies(dependency_list):
  
                  def no_op_wrapper():
                      return tf.group(*[tf.compat.v1.assign_add(self.cold_step, 1)])
  
                  if not self._async_eigen_decomp:
                      update_factor_ops = tf.cond(
                          pred=tf.logical_and(
                              tf.equal(
                                  tf.math.floormod(self.stats_step, self._kfac_update),
                                  tf.convert_to_tensor(value=0),
                              ),
                              tf.greater_equal(self.stats_step, self._stats_accum_iter),
                          ),
                          true_fn=lambda: tf.group(
                              *self.apply_stats_eigen(self.compute_stats_eigen())
                          ),
                          false_fn=no_op_wrapper,
                      )
                  else:
                      update_factor_ops = tf.cond(
                          pred=tf.greater_equal(self.stats_step, self._stats_accum_iter),
                          true_fn=lambda: tf.cond(
                              pred=tf.equal(queue.size(), tf.convert_to_tensor(value=0)),
                              true_fn=tf.no_op,
                              false_fn=lambda: tf.group(
                                  *self.apply_stats_eigen(dequeue_op())
                              ),
                          ),
                          false_fn=no_op_wrapper,
                      )
                  update_ops.append(update_factor_ops)
                  with tf.control_dependencies([update_factor_ops]):
  
                      def grad_op():
                          return list(grad)
  
                      def get_kfac_grad_op():
                          return self.get_kfac_precond_updates(grad, varlist)
  
                      u = tf.cond(
                          pred=tf.greater(
                              self.factor_step, tf.convert_to_tensor(value=0)
                          ),
                          true_fn=get_kfac_grad_op,
                          false_fn=grad_op,
                      )
                      optim = tf.compat.v1.train.MomentumOptimizer(
                          self._lr * (1.0 - self._momentum), self._momentum
                      )
  
                      def optim_op():
                          def update_optim_op():
                              if self._full_stats_init:
                                  return tf.cond(
                                      pred=tf.greater(
                                          self.factor_step, tf.convert_to_tensor(value=0)
                                      ),
                                      true_fn=lambda: optim.apply_gradients(
                                          list(zip(u, varlist))
                                      ),
                                      false_fn=tf.no_op,
                                  )
                              else:
                                  return optim.apply_gradients(list(zip(u, varlist)))
  
                          if self._full_stats_init:
                              return tf.cond(
                                  pred=tf.greater_equal(
                                      self.stats_step, self._stats_accum_iter
                                  ),
                                  true_fn=update_optim_op,
                                  false_fn=tf.no_op,
                              )
                          else:
                              return tf.cond(
                                  pred=tf.greater_equal(self.sgd_step, self._cold_iter),
                                  true_fn=update_optim_op,
                                  false_fn=tf.no_op,
                              )
  
                      update_ops.append(optim_op())
          return tf.group(*update_ops), queue_runner
  
      def apply_gradients(self, grads):
          cold_optim = tf.compat.v1.train.MomentumOptimizer(self._cold_lr, self._momentum)
  
          def _cold_sgd_start():
              sgd_grads, sgd_var = zip(*grads)
              if self.max_grad_norm is not None:
                  sgd_grads, _ = tf.clip_by_global_norm(sgd_grads, self.max_grad_norm)
              sgd_grads = list(zip(sgd_grads, sgd_var))
              sgd_step_op = tf.compat.v1.assign_add(self.sgd_step, 1)
              cold_optim_op = cold_optim.apply_gradients(sgd_grads)
              if KFAC_DEBUG:
                  with tf.control_dependencies([sgd_step_op, cold_optim_op]):
                      sgd_step_op = tf.compat.v1.Print(
                          sgd_step_op,
                          [
                              self.sgd_step,
                              tf.convert_to_tensor(value='doing cold sgd step'),
                          ],
                      )
              return tf.group(*[sgd_step_op, cold_optim_op])
  
          grads = [(grad, var) for (grad, var) in grads if grad is not None]
          kfac_optim_op, queue_runner = self.apply_gradients_kfac(grads)
  
          def _warm_kfac_start():
              return kfac_optim_op
  
          return (
              tf.cond(
                  pred=tf.greater(self.sgd_step, self._cold_iter),
                  true_fn=_warm_kfac_start,
                  false_fn=_cold_sgd_start,
              ),
              queue_runner,
          )
  
      def minimize(self, loss, loss_sampled, var_list=None):
          grads = self.compute_gradients(loss, var_list=var_list)
          self.compute_and_apply_stats(loss_sampled, var_list=var_list)
          return self.apply_gradients(grads)
  