def slots(self):
    """Returns the slots of all optimizers."""
    optimizers = list(self._optimizers) + [self._loss_opt]
    return fastmath.nested_map(lambda opt: opt.slots, optimizers)
Returns the slots of all optimizers.
trax/optimizers/trainer.py