class_docstr: Adafactor optimizer, as described in https://arxiv.org/abs/1804.04235.
code: def init(self, weights):
      shape = weights.shape
      slots = []
      if self._factored and len(shape) >= 2:
          v_row = jnp.zeros(shape[:-1], dtype=jnp.float32)
          v_col = jnp.zeros(shape[:-2] + shape[-1:], dtype=jnp.float32)
          slots.extend([v_row, v_col])
      else:
          v = jnp.zeros_like(weights)
          slots.append(v)
      if self._do_momentum:
          m = jnp.zeros_like(weights)
          if self._momentum_in_bfloat16:
              m = m.astype(jnp.bfloat16)
          slots.append(m)
      return slots
  
imports: from trax.fastmath import numpy
  from trax.optimizers import base
label_desc: Use image node_id 3 for calling the Adafactor obj's underlying member method code with example usage: obj.init(weights) and returns: slots