class_docstr: Base class for weight update methods.
  
  Attributes:
    all_layers_names: list of layer names in a model.
    recycle_type: neuron, layer based.
    dead_neurons_threshold: below this threshold a neuron is considered dead.
    reset_layers: list of layer names to be recycled.
    reset_start_layer_idx: index of the layer from which we start recycling.
    reset_period: int represents the period of weight update.
    reset_start_step: start recycle from start step
    reset_end_step:  end recycle from end step
    logging_period:  the period of statistics logging e.g., dead neurons.
    prev_neuron_score: score at last reset step or log step in case of no reset.
    sub_mean_score: if True the average activation will be subtracted
                    for each neuron when we calculate the score.
class_docstr_tok: ["Base","class","for","weight","update","methods",".","Attributes",":","all_layers_names",":","list","of","layer","names","in","a","model",".","recycle_type",":","neuron",",","layer","based",".","dead_neurons_threshold",":","below","this","threshold","a","neuron","is","considered","dead",".","reset_layers",":","list","of","layer","names","to","be","recycled",".","reset_start_layer_idx",":","index","of","the","layer","from","which","we","start","recycling",".","reset_period",":","int","represents","the","period","of","weight","update",".","reset_start_step",":","start","recycle","from","start","step","reset_end_step",":","end","recycle","from","end","step","logging_period",":","the","period","of","statistics","logging","e.g.",",","dead","neurons",".","prev_neuron_score",":","score","at","last","reset","step","or","log","step","in","case","of","no","reset",".","sub_mean_score",":","if","True","the","average","activation","will","be","subtracted","for","each","neuron","when","we","calculate","the","score","."]
docstr: log dead neurons in each layer.
  
  For conv layer we also log dead elements in the spatial dimension.
  
  Args:
    intermediates: intermidate activation in each layer.
  
  Returns:
    log_dict_elements_per_neuron
    log_dict_neurons
docstr_tok: ["log","dead","neurons","in","each","layer",".","For","conv","layer","we","also","log","dead","elements","in","the","spatial","dimension",".","Args",":","intermediates",":","intermidate","activation","in","each","layer",".","Returns",":","log_dict_elements_per_neuron","log_dict_neurons"]
code: def log_dead_neurons_count(self, intermediates):
      """log dead neurons in each layer.
  
      For conv layer we also log dead elements in the spatial dimension.
  
      Args:
        intermediates: intermidate activation in each layer.
  
      Returns:
        log_dict_elements_per_neuron
        log_dict_neurons
      """
  
      def log_dict(score, score_type):
          total_neurons, total_deadneurons = 0.0, 0.0
          score_dict = flax.traverse_util.flatten_dict(score, sep="/")
  
          log_dict = {}
          for k, m in score_dict.items():
              if "final_layer" in k:
                  continue
              m = m[0]
              layer_size = float(jnp.size(m))
              deadneurons_count = jnp.count_nonzero(
                  m <= self.dead_neurons_threshold
              )
              total_neurons += layer_size
              total_deadneurons += deadneurons_count
              log_dict[f"dead_{score_type}_percentage/{k[:-9]}"] = (
                  float(deadneurons_count) / layer_size
              ) * 100.0
              log_dict[f"dead_{score_type}_count/{k[:-9]}"] = float(
                  deadneurons_count
              )
          log_dict[f"{score_type}/total"] = total_neurons
          log_dict[f"{score_type}/deadcount"] = float(total_deadneurons)
          log_dict[f"dead_{score_type}_percentage"] = (
              float(total_deadneurons) / total_neurons
          ) * 100.0
          return log_dict
  
      neuron_score = jax.tree_map(
          self.estimate_neuron_score, intermediates
      )
      log_dict_neurons = log_dict(neuron_score, "feature")
  
      return log_dict_neurons
  
imports: import functools
  import logging
  import flax
  from flax import linen
  import gin
  import jax
  from jax import random
  import jax.numpy
  import optax
label_desc: Use image node_id 12 for calling the BaseRecycler obj's underlying member method code with example usage: obj.log_dead_neurons_count(intermediates) and returns: log_dict_neurons, log_dict