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: create the masks for recycled weights based on neurons scores.
  
  Args:
    param_dict: dict of model params.
    activations_dict: dict of the neuron score of each layer.
    key: used seed for random weights.
  
  Returns:
    incoming_mask_dict
    outgoing_mask_dict
    ingoing_random_keys_dict
    outgoing_random_keys_dict
    param_dict
docstr_tok: ["create","the","masks","for","recycled","weights","based","on","neurons","scores",".","Args",":","param_dict",":","dict","of","model","params",".","activations_dict",":","dict","of","the","neuron","score","of","each","layer",".","key",":","used","seed","for","random","weights",".","Returns",":","incoming_mask_dict","outgoing_mask_dict","ingoing_random_keys_dict","outgoing_random_keys_dict","param_dict"]
code: def create_masks(self, param_dict, activations_dict, key):
      """create the masks for recycled weights based on neurons scores.
  
      Args:
        param_dict: dict of model params.
        activations_dict: dict of the neuron score of each layer.
        key: used seed for random weights.
  
      Returns:
        incoming_mask_dict
        outgoing_mask_dict
        ingoing_random_keys_dict
        outgoing_random_keys_dict
        param_dict
      """
      incoming_mask_dict = {
          k: jnp.zeros_like(p) if p.ndim != 1 else None
          for k, p in param_dict.items()
      }
      outgoing_mask_dict = {
          k: jnp.zeros_like(p) if p.ndim != 1 else None
          for k, p in param_dict.items()
      }
      ingoing_random_keys_dict = {k: None for k in param_dict}
      outgoing_random_keys_dict = (
          {k: None for k in param_dict}
          if self.init_method_outgoing == "random"
          else {}
      )
  
      # prepare mask of incoming and outgoing recycled connections
      for k in self.reset_layers:
          param_key = "params/" + k + "/kernel"
          param = param_dict[param_key]
          # This won't work for DRQ, since returned keys can be a list.
          # We don't support that at the moment.
          next_key = self.next_layers[k]
          if isinstance(next_key, list):
              next_key = next_key[0]
          next_param = param_dict["params/" + next_key + "/kernel"]
          activation = activations_dict[k + "_act/__call__"][0]
          # TODO(evcu) Maybe use per_layer random keys here.
          neuron_mask = self._score2mask(
              activation, param, next_param, key
          )
  
          # the for loop handles the case where a layer has multiple next layers
          # like the case in DrQ where the output layer has multihead.
          next_keys = (
              self.next_layers[k]
              if isinstance(self.next_layers[k], list)
              else [self.next_layers[k]]
          )
          for next_k in next_keys:
              next_param_key = "params/" + next_k + "/kernel"
              next_param = param_dict[next_param_key]
              incoming_mask, outgoing_mask = self.create_mask_helper(
                  neuron_mask, param, next_param
              )
              incoming_mask_dict[param_key] = incoming_mask
              outgoing_mask_dict[next_param_key] = outgoing_mask
              key, subkey = random.split(key)
              ingoing_random_keys_dict[param_key] = subkey
              if self.init_method_outgoing == "random":
                  key, subkey = random.split(key)
                  outgoing_random_keys_dict[next_param_key] = subkey
  
          # reset bias
          bias_key = "params/" + k + "/bias"
          new_bias = jnp.zeros_like(param_dict[bias_key])
          param_dict[bias_key] = jnp.where(
              neuron_mask, new_bias, param_dict[bias_key]
          )
  
      return (
          incoming_mask_dict,
          outgoing_mask_dict,
          ingoing_random_keys_dict,
          outgoing_random_keys_dict,
          param_dict,
      )
  
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 9 for calling the NeuronRecycler obj's underlying member method code with example usage: obj.create_masks(param_dict, activations_dict, key) and returns: incoming_mask_dict,  outgoing_mask_dict,  ingoing_random_keys_dict,  outgoing_random_keys_dict,  param_dict