def test_run_reversible_slots(self):
    """Tests that slots can be read and assigned in reversible trainer."""
    layers = [tl.Dense(4), tl.Dup()]
    rev_layers = [tl.ReversibleHalfResidual(tl.Dense(4)),
                  tl.ReversibleSwap()]
    loss_layer = tl.Serial(tl.Concatenate(), tl.Dense(4),
                           tl.LogSoftmax(), tl.CrossEntropyLoss())
    trainer = optimizers.ReversibleSerialTrainer(
        [(layers, rev_layers)], loss_layer, optimizers.Adam)
    slots = trainer.slots
    trainer.slots = slots
    self.assertEqual(slots, trainer.slots)
Tests that slots can be read and assigned in reversible trainer.
trax/optimizers/trainer_test.py