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","."]
docstr: Recycle dead neurons by reinitalizie incoming and outgoing connections.
  
  Incoming connections are randomly initalized and outgoing connections
  are zero initalized.
  A featuremap is considered dead when its score is below or equal
  dead neuron threshold.
  Args:
    intermedieates: pytree contains the activations over a batch.
    params: current weights of the model.
    key: used to generate random keys.
    opt_state: state of optimizer.
  
  Returns:
    new model params after recycling dead neurons.
    opt_state: new state for the optimizer
  
  Raises: raise error if init_method_outgoing is not one of the following
  (random, zero).
docstr_tok: ["Recycle","dead","neurons","by","reinitalizie","incoming","and","outgoing","connections",".","Incoming","connections","are","randomly","initalized","and","outgoing","connections","are","zero","initalized",".","A","featuremap","is","considered","dead","when","its","score","is","below","or","equal","dead","neuron","threshold",".","Args",":","intermedieates",":","pytree","contains","the","activations","over","a","batch",".","params",":","current","weights","of","the","model",".","key",":","used","to","generate","random","keys",".","opt_state",":","state","of","optimizer",".","Returns",":","new","model","params","after","recycling","dead","neurons",".","opt_state",":","new","state","for","the","optimizer","Raises",":","raise","error","if","init_method_outgoing","is","not","one","of","the","following","(","random",",","zero",")","."]
code: def recycle_dead_neurons(
      self, intermedieates, params, key, opt_state
  ):
      """Recycle dead neurons by reinitalizie incoming and outgoing connections.
  
      Incoming connections are randomly initalized and outgoing connections
      are zero initalized.
      A featuremap is considered dead when its score is below or equal
      dead neuron threshold.
      Args:
        intermedieates: pytree contains the activations over a batch.
        params: current weights of the model.
        key: used to generate random keys.
        opt_state: state of optimizer.
  
      Returns:
        new model params after recycling dead neurons.
        opt_state: new state for the optimizer
  
      Raises: raise error if init_method_outgoing is not one of the following
      (random, zero).
      """
      activations_score_dict = flax.traverse_util.flatten_dict(
          intermedieates, sep="/"
      )
      param_dict = flax.traverse_util.flatten_dict(params, sep="/")
  
      # create incoming and outgoing masks and reset bias of dead neurons.
      (
          incoming_mask_dict,
          outgoing_mask_dict,
          incoming_random_keys_dict,
          outgoing_random_keys_dict,
          param_dict,
      ) = self.create_masks(param_dict, activations_score_dict, key)
  
      params = flax.core.freeze(
          flax.traverse_util.unflatten_dict(param_dict, sep="/")
      )
      incoming_random_keys = flax.core.freeze(
          flax.traverse_util.unflatten_dict(
              incoming_random_keys_dict, sep="/"
          )
      )
      if self.init_method_outgoing == "random":
          outgoing_random_keys = flax.core.freeze(
              flax.traverse_util.unflatten_dict(
                  outgoing_random_keys_dict, sep="/"
              )
          )
      # reset incoming weights
      incoming_mask = flax.core.freeze(
          flax.traverse_util.unflatten_dict(incoming_mask_dict, sep="/")
      )
      reinit_fn = functools.partial(
          weight_reinit_random,
          weight_scaling=self.weight_scaling,
          scale=self.incoming_scale,
          weights_type="incoming",
      )
      weight_random_reset_fn = jax.jit(
          functools.partial(jax.tree_map, reinit_fn)
      )
      params = weight_random_reset_fn(
          params, incoming_mask, incoming_random_keys
      )
  
      # reset outgoing weights
      outgoing_mask = flax.core.freeze(
          flax.traverse_util.unflatten_dict(outgoing_mask_dict, sep="/")
      )
      if self.init_method_outgoing == "random":
          reinit_fn = functools.partial(
              weight_reinit_random,
              weight_scaling=self.weight_scaling,
              scale=self.outgoing_scale,
              weights_type="outgoing",
          )
          weight_random_reset_fn = jax.jit(
              functools.partial(jax.tree_map, reinit_fn)
          )
          params = weight_random_reset_fn(
              params, outgoing_mask, outgoing_random_keys
          )
      elif self.init_method_outgoing == "zero":
          weight_zero_reset_fn = jax.jit(
              functools.partial(jax.tree_map, weight_reinit_zero)
          )
          params = weight_zero_reset_fn(params, outgoing_mask)
      else:
          raise ValueError(
              f"Invalid init method: {self.init_method_outgoing}"
          )
      # reset mu, nu of adam optimizer for recycled weights.
      reset_momentum_fn = jax.jit(
          functools.partial(jax.tree_map, reset_momentum)
      )
      new_mu = reset_momentum_fn(opt_state[0][1], incoming_mask)
      new_mu = reset_momentum_fn(new_mu, outgoing_mask)
      new_nu = reset_momentum_fn(opt_state[0][2], incoming_mask)
      new_nu = reset_momentum_fn(new_nu, outgoing_mask)
      opt_state_list = list(opt_state)
      opt_state_list[0] = optax.ScaleByAdamState(
          opt_state[0].count, mu=new_mu, nu=new_nu
      )
      opt_state = tuple(opt_state_list)
      return params, opt_state
  
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 7 for calling the NeuronRecycler obj's underlying member method code with example usage: obj.recycle_dead_neurons(intermedieates, params, key, opt_state) and returns: params,  opt_state