class_docstr: Recycle the weights connected to dead neurons.
  
  In convolutional neural networks, we consider a feature map as neuron.
  
  Attributes:
    next_layers: dict key a current layer name, value next layer name.
    init_method_outgoing: method to init outgoing weights (random, zero).
    weight_scaling: if true, scale reinit weights.
    incoming_scale: scalar for incoming weights.
    outgoing_scale: scalar for outgoing weights.
class_docstr_tok: ["Recycle","the","weights","connected","to","dead","neurons",".","In","convolutional","neural","networks",",","we","consider","a","feature","map","as","neuron",".","Attributes",":","next_layers",":","dict","key","a","current","layer","name",",","value","next","layer","name",".","init_method_outgoing",":","method","to","init","outgoing","weights","(","random",",","zero",")",".","weight_scaling",":","if","true",",","scale","reinit","weights",".","incoming_scale",":","scalar","for","incoming","weights",".","outgoing_scale",":","scalar","for","outgoing","weights","."]
code: def __init__(
      self,
      all_layers_names,
      init_method_outgoing="zero",
      weight_scaling=False,
      incoming_scale=1.0,
      outgoing_scale=1.0,
      network="nature",
      **kwargs,
  ):
      super(NeuronRecycler, self).__init__(all_layers_names, **kwargs)
      self.init_method_outgoing = init_method_outgoing
      self.weight_scaling = weight_scaling
      self.incoming_scale = incoming_scale
      self.outgoing_scale = outgoing_scale
      # prepare a dict that has pointer to next layer give a layer name
      # this is needed because neuron recycle reinitalizes both sides
      # (incoming and outgoing weights) of a neuron and needs a point to the
      # outgoing weights.
      self.next_layers = {}
      for current_layer, next_layer in zip(
          all_layers_names[:-1], all_layers_names[1:]
      ):
          self.next_layers[current_layer] = next_layer
  
      # we don't recycle the neurons in the output layer.
      self.reset_layers = self.reset_layers[:-1]
  
      # if network is resnet, recycle only the incoming/outgoing of the first conv
      # layer in each block and final dense layer
      if network == "resnet":
          self.reset_layers = []
          for layer in self.all_layers_names:
              if (
                  "Conv_1" in layer
                  or "Conv_3" in layer
                  or "Dense" in layer
              ):
                  self.reset_layers.append(layer)
  
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 NeuronRecycler object from inherited base classes: BaseRecycler with example: obj = NeuronRecycler(all_layers_names, init_method_outgoing, weight_scaling, incoming_scale, outgoing_scale, network)