def slots(self, slots):
    """Sets the slots of all optimizers."""
    for ((s_opt, r_opts), (s_slots, r_slots)) in zip(
        self._optimizers, slots[:-1]):
      for (opt, slot) in zip([s_opt] + r_opts, [s_slots] + r_slots):
        opt.slots = slot
    self._loss_opt.slots = slots[-1]
Sets the slots of all optimizers.
trax/optimizers/trainer.py