class_docstr: Reset all weights of some layers.
  
  This class implements the primacy bias paper, in which all weights of the
  last layers of the model are randomly reinitalized.
  
  Attributes:
class_docstr_tok: ["Reset","all","weights","of","some","layers",".","This","class","implements","the","primacy","bias","paper",",","in","which","all","weights","of","the","last","layers","of","the","model","are","randomly","reinitalized",".","Attributes",":"]
code: def update_weights(self, intermediates, params, key, opt_state):
      param_dict = flax.traverse_util.flatten_dict(params, sep="/")
      mask_dict = {
          k: jnp.zeros_like(p) if p.ndim != 1 else None
          for k, p in param_dict.items()
      }
      random_keys_dict = {k: None for k in param_dict}
      # set mask = 1 for the weights that will be recycled
      for param_key, param in param_dict.items():
          param_layer_name = param_key.split("/")[1]
          if param.ndim != 1 and param_layer_name in self.reset_layers:
              mask_dict[param_key] = jnp.ones_like(param)
              key, subkey = random.split(key)
              random_keys_dict[param_key] = subkey
              # reset bias
              bias_key = "params/" + param_layer_name + "/bias"
              new_bias = jnp.zeros_like(param_dict[bias_key])
              param_dict[bias_key] = new_bias
  
      params = flax.core.freeze(
          flax.traverse_util.unflatten_dict(param_dict, sep="/")
      )
      masks = flax.core.freeze(
          flax.traverse_util.unflatten_dict(mask_dict, sep="/")
      )
      random_keys = flax.core.freeze(
          flax.traverse_util.unflatten_dict(random_keys_dict, sep="/")
      )
      # reset weights
      weight_random_reset_fn = jax.jit(
          functools.partial(jax.tree_map, weight_reinit_random)
      )
      params = weight_random_reset_fn(params, masks, random_keys)
  
      # 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], masks)
      new_nu = reset_momentum_fn(opt_state[0][2], masks)
      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 2 for calling the LayerReset obj's underlying member method code with example usage: obj.update_weights(intermediates, params, key, opt_state) and returns: params,  opt_state