class_docstr: Adafactor optimizer, as described in https://arxiv.org/abs/1804.04235.
code: def update(self, step, grads, weights, slots, opt_params):
      updates = []
      learning_rate = opt_params["learning_rate"]
      beta1 = opt_params["beta1"]
      decay_rate = opt_params["decay_rate"]
      clipping_threshold = opt_params["clipping_threshold"]
      weight_decay_rate = opt_params["weight_decay_rate"]
      weight_decay_n_steps = opt_params["weight_decay_n_steps"]
      weight_decay_rate = jnp.where(
          weight_decay_n_steps
          < 1,  # if weight_decay_n_steps == 0, ignore it
          weight_decay_rate,
          (
              weight_decay_rate
              * jnp.maximum(weight_decay_n_steps - step, 0.0)
              / jnp.maximum(weight_decay_n_steps, 0.0)
          ),
      )
      epsilon1 = opt_params["epsilon1"]
      epsilon2 = opt_params["epsilon2"]
      decay_rate = self._decay_rate_pow(step, exponent=decay_rate)
      update_scale = learning_rate
      if self._multiply_by_parameter_scale:
          update_scale *= jnp.maximum(
              jnp.sqrt(jnp.mean(weights * weights)), epsilon2
          )
      mixing_rate = 1.0 - decay_rate
  
      grads_sqr = grads * grads
      if self._factored and len(weights.shape) >= 2:
          v_row = slots[
              0
          ]  # In this case, the slots are (v_row, v_col, ...).
          v_col = slots[1]
          new_v_row = decay_rate * v_row + mixing_rate * jnp.mean(
              grads_sqr, axis=-1
          )
          new_v_col = decay_rate * v_col + mixing_rate * jnp.mean(
              grads_sqr, axis=-2
          )
          updates.extend([new_v_row, new_v_col])
          row_mean = jnp.mean(new_v_row, axis=-1, keepdims=True)
          row_factor = (row_mean / (new_v_row + epsilon1)) ** 0.5
          col_factor = (new_v_col + epsilon1) ** -0.5
          y = (
              grads
              * jnp.expand_dims(row_factor, axis=-1)
              * jnp.expand_dims(col_factor, axis=-2)
          )
      else:
          v = slots[0]  # In this case, the slots are (v, ...)
          new_v = decay_rate * v + mixing_rate * grads_sqr
          updates.append(new_v)
          y = grads * (new_v + epsilon1) ** -0.5
  
      if self._do_clipping:
          clipping_denom = jnp.maximum(
              1.0, jnp.sqrt(jnp.mean(y * y)) / clipping_threshold
          )
          y /= clipping_denom
  
      subtrahend = update_scale * y
      if self._do_momentum:
          m = slots[-1]  # Momentum is always the last slot (if used).
          m = m.astype(
              subtrahend.dtype
          )  # Accumulate in subtrahend dtype.
          new_m = beta1 * m + (1.0 - beta1) * subtrahend
          subtrahend = new_m
          updates.append(
              new_m.astype(slots[-1].dtype)
          )  # Back to bfloat if needed.
  
      new_weights = (1 - weight_decay_rate) * weights - subtrahend
      # TODO(lukaszkaiser): why is the astype needed here? Check and correct.
      return new_weights.astype(weights.dtype), updates
  
imports: from trax.fastmath import numpy
  from trax.optimizers import base
label_desc: Use image node_id 4 for calling the Adafactor obj's underlying member method code with example usage: obj.update(step, grads, weights, slots, opt_params) and returns: new_weights,  updates