class_docstr: Fixed scheduled version of the NeuronRecycler.
class_docstr_tok: ["Fixed","scheduled","version","of","the","NeuronRecycler","."]
code: def _score2mask(self, activation, param, next_param, key):
      is_cbp = self.score_type == "cbp"
      score = self.estimate_neuron_score(activation, is_cbp=is_cbp)
      if self.score_type == "redo":
          pass
      elif self.score_type == "random":
          new_key = random.fold_in(key, self._last_update_step)
          score = random.shuffle(new_key, score)
      elif self.score_type == "redo_inverted":
          score = -score
      # Metric used in Continual Backprop pape.
      elif self.score_type == "cbp":
          next_axes = list(range(param.ndim))
          del next_axes[-2]
          current_axes = list(range(param.ndim))
          del current_axes[-1]
          if next_param.ndim == 2 and param.ndim == 4:
              new_shape = activation.shape[1:] + (-1,)
              next_param = jnp.reshape(next_param, new_shape)
          score *= jnp.sum(
              jnp.abs(next_param), axis=next_axes
          ) / jnp.sum(jnp.abs(param), axis=current_axes)
      multiplier = max(0, self._last_update_step / self.reset_end_step)
      ones_fraction = float(jnp.cos(jnp.pi * 0.5 * multiplier))
      ones_fraction *= self.recycle_rate
      logging.info(
          "score_type: %s, multiplier: %f, ones_fraction=%f",
          self.score_type,
          multiplier,
          ones_fraction,
      )
      return leastk_mask(score, ones_fraction)
  
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 NeuronRecyclerScheduled obj's underlying member method code with example usage: obj._score2mask(activation, param, next_param, key) and returns: leastk_mask