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: generate incoming and outgoing weight mask given dead neurons mask.
  
  Args:
    neuron_mask: mask of size equals the width of a layer.
    current_param: incoming weights of a layer.
    next_param: outgoing weights of a layer.
  
  Returns:
    incoming_mask
    outgoing_mask
docstr_tok: ["generate","incoming","and","outgoing","weight","mask","given","dead","neurons","mask",".","Args",":","neuron_mask",":","mask","of","size","equals","the","width","of","a","layer",".","current_param",":","incoming","weights","of","a","layer",".","next_param",":","outgoing","weights","of","a","layer",".","Returns",":","incoming_mask","outgoing_mask"]
code: def create_mask_helper(self, neuron_mask, current_param, next_param):
      """generate incoming and outgoing weight mask given dead neurons mask.
  
      Args:
        neuron_mask: mask of size equals the width of a layer.
        current_param: incoming weights of a layer.
        next_param: outgoing weights of a layer.
  
      Returns:
        incoming_mask
        outgoing_mask
      """
  
      def mask_creator(
          expansion_axis, expansion_axes, param, neuron_mask
      ):
          """create a mask of weight matrix given 1D vector of neurons mask.
  
          Args:
            expansion_axis: List contains 1 axis. The dimension to expand the mask
              for dense layers (weight shape 2D).
            expansion_axes: List conrtains 3 axes. The dimensions to expand the
              score for convolutional layers (weight shape 4D).
            param: weight.
            neuron_mask: 1D mask that represents dead neurons(features).
  
          Returns:
            mask: mask of weight.
          """
          if param.ndim == 2:
              axes = expansion_axis
              # flatten layer
              # The size of neuron_mask is the same as the width of last conv layer.
              # This conv layer will be flatten and connected to dense layer.
              # we repeat each value of a feature map to cover the spatial dimension.
              if axes[0] == 1 and (
                  param.shape[0] > neuron_mask.shape[0]
              ):
                  num_repeatition = int(
                      param.shape[0] / neuron_mask.shape[0]
                  )
                  neuron_mask = jnp.repeat(
                      neuron_mask, num_repeatition, axis=0
                  )
          elif param.ndim == 4:
              axes = expansion_axes
          mask = jnp.expand_dims(neuron_mask, axis=tuple(axes))
          for i in range(len(axes)):
              mask = jnp.repeat(
                  mask, param.shape[axes[i]], axis=axes[i]
              )
          return mask
  
      incoming_mask = mask_creator(
          [0], [0, 1, 2], current_param, neuron_mask
      )
      outgoing_mask = mask_creator(
          [1], [0, 1, 3], next_param, neuron_mask
      )
      return incoming_mask, outgoing_mask
  
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 10 for calling the NeuronRecycler obj's underlying member method code with example usage: obj.create_mask_helper(neuron_mask, current_param, next_param) and returns: incoming_mask,  outgoing_mask, mask