code: def weight_reinit_random(
      param,
      mask,
      key,
      weight_scaling=False,
      scale=1.0,
      weights_type="incoming",
  ):
      """Randomly reinit recycled weights and may scale its norm.
  
      If scaling applied, the norm of recycled weights equals
      the average norm of non recycled weights per neuron multiplied by a scalar.
  
      Args:
        param: current param
        mask: incoming/outgoing mask for recycled weights
        key: random key to generate new random weights
        weight_scaling: if true scale recycled weights with the norm of non recycled
        scale: scale to multiply the new weights norm.
        weights_type: incoming or outgoing weights
  
      Returns:
      params: new params after weight recycle.
      """
      if mask is None or key is None:
          return param
  
      new_param = nn.initializers.xavier_uniform()(
          key, shape=param.shape
      )
  
      if weight_scaling:
          axes = list(range(param.ndim))
          if weights_type == "outgoing":
              del axes[-2]
          else:
              del axes[-1]
  
          neuron_mask = jnp.mean(mask, axis=axes)
  
          non_dead_count = neuron_mask.shape[0] - jnp.count_nonzero(
              neuron_mask
          )
          norm_per_neuron = _get_norm_per_neuron(param, axes)
          non_recycled_norm = (
              jnp.sum(norm_per_neuron * (1 - neuron_mask))
              / non_dead_count
          )
          non_recycled_norm = non_recycled_norm * scale
  
          normalized_new_param = _weight_normalization_per_neuron_norm(
              new_param, axes
          )
          new_param = normalized_new_param * non_recycled_norm
  
      param = jnp.where(mask == 1, new_param, param)
      return param
  
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 4 for calling a global function with example usage: weight_reinit_random(param, mask, key, weight_scaling, scale, weights_type) and returns: param, param