docstr: compute the stats values
  
  :param loss_sampled: ([TensorFlow Tensor]) the loss function output
  :param var_list: ([TensorFlow Tensor]) The parameters
  :return: ([TensorFlow Tensor]) stats updates
docstr_tok: ["compute","the","stats","values",":","param","loss_sampled",":","(","[","TensorFlow","Tensor","]",")","the","loss","function","output",":","param","var_list",":","(","[","TensorFlow","Tensor","]",")","The","parameters",":","return",":","(","[","TensorFlow","Tensor","]",")","stats","updates"]
code: def compute_stats(self, loss_sampled, var_list=None):
      """
      compute the stats values
  
      :param loss_sampled: ([TensorFlow Tensor]) the loss function output
      :param var_list: ([TensorFlow Tensor]) The parameters
      :return: ([TensorFlow Tensor]) stats updates
      """
      varlist = var_list
      if varlist is None:
          varlist = tf.trainable_variables()
  
      gradient_sampled = tf.gradients(
          loss_sampled, varlist, name="gradientsSampled"
      )
      self.gradient_sampled = gradient_sampled
  
      # remove unused variables
      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(fprop_factor)[0]  # batch size
                  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
                      ):
                          # factorization along the channels
                          # assume independence among input channels
                          # factor = B x 1 x 1 x (KH xKW x C)
                          # patches = B x Oh x Ow x (KH xKW x C)
                          if len(svd_factors) == 0:
                              if KFAC_DEBUG:
                                  print(
                                      (
                                          "approx %s act factor with rank-1 SVD factors"
                                          % var.name
                                      )
                                  )
                              # find closest rank-1 approx to the feature map
                              S, U, V = tf.batch_svd(
                                  tf.reshape(
                                      fprop_factor,
                                      [
                                          -1,
                                          kernel_height * kernel_width,
                                          chan,
                                      ],
                                  )
                              )
                              # get rank-1 approx slides
                              sqrt_s1 = tf.expand_dims(
                                  tf.sqrt(S[:, 0, 0]), 1
                              )
                              patches_k = (
                                  U[:, :, 0] * sqrt_s1
                              )  # B x KH*KW
                              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  # B x C
                              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:
                          # poor mem usage implementation
                          patches = tf.extract_image_patches(
                              fprop_factor,
                              ksizes=[
                                  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
                                      )
                                  )
                              # T^2 terms * 1/T^2, size: B x C
                              fprop_factor = tf.reduce_mean(
                                  patches, [1, 2]
                              )
                          else:
                              # size: (B x Oh x Ow) x C
                              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:
                          # correct padding for numerical stability (we
                          # divided out OhxOw from activations for T1 approx)
                          fprop_factor = tf.concat(
                              [
                                  fprop_factor,
                                  tf.ones(
                                      [tf.shape(fprop_factor)[0], 1]
                                  )
                                  / operator_height
                                  / operator_width,
                              ],
                              1,
                          )
                      else:
                          # use homogeneous coordinates
                          fprop_factor = tf.concat(
                              [
                                  fprop_factor,
                                  tf.ones(
                                      [tf.shape(fprop_factor)[0], 1]
                                  ),
                              ],
                              1,
                          )
  
                  # average over the number of data points in a batch
                  # divided by B
                  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":
                      # HACK: for convolution we recompute fprop stats for
                      # every layer including forking layers
                      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(bprop_factor)[0]  # batch size
                  chan = int(bprop_factor_shape[-1])  # num channels
                  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(
                                  bprop_factor, [1, 2]
                              )  # T^2 terms * 1/T^2
                          else:
                              bprop_factor = (
                                  tf.reshape(bprop_factor, [-1, chan])
                                  * operator_height
                                  * operator_width
                              )  # T * 1/T terms
                      else:
                          # just doing block diag approx. spatial independent
                          # structure does not apply here. summing over
                          # spatial locations
                          if KFAC_DEBUG:
                              print(
                                  (
                                      "block diag approx fisher for %s"
                                      % var.name
                                  )
                              )
                          bprop_factor = tf.reduce_sum(
                              bprop_factor, [1, 2]
                          )
  
                  # assume sampled loss is averaged. TODO:figure out better
                  # way to handle this
                  bprop_factor *= tf.cast(batch_size, tf.float32)
                  ##
  
                  cov_b = tf.matmul(
                      bprop_factor, bprop_factor, transpose_a=True
                  ) / tf.cast(tf.shape(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.Print(
              stats_updates[a_key],
              [
                  tf.convert_to_tensor("step:"),
                  self.global_step,
                  tf.convert_to_tensor("computing stats"),
              ],
          )
      self.stats_updates = stats_updates
      return stats_updates
  
imports: import re
  from functools import reduce
  import tensorflow
  import numpy
  from stable_baselines.acktr.kfac_utils import detect_min_val, factor_reshape, gmatmul
label_desc: Use image node_id 5 for calling the KfacOptimizer obj's underlying member method code with example usage: obj.compute_stats(loss_sampled, var_list) and returns: stats_updates