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","."]
code: def __init__(
      self,
      all_layers_names,
      dead_neurons_threshold=0.0,
      reset_start_layer_idx=0,
      reset_period=200_000,
      reset_start_step=0,
      reset_end_step=100_000_000,
      logging_period=20_000,
      sub_mean_score=False,
  ):
      self.all_layers_names = all_layers_names
      self.dead_neurons_threshold = dead_neurons_threshold
      self.reset_layers = all_layers_names[reset_start_layer_idx:]
      self.reset_period = reset_period
      self.reset_start_step = reset_start_step
      self.reset_end_step = reset_end_step
      self.logging_period = logging_period
      self.prev_neuron_score = None
      self.sub_mean_score = sub_mean_score
  
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 1 to create a new BaseRecycler object  with example: obj = BaseRecycler(all_layers_names, dead_neurons_threshold, reset_start_layer_idx, reset_period, reset_start_step, reset_end_step, logging_period, sub_mean_score)