def test_run_reversible_same_as_default_extended(self):
    """Runs the reversible trainer, check results are the same as default."""
    inputs_batch = np.arange(8).reshape((2, 4))
    targets_batch = 2 * inputs_batch
    labeled_batch = (inputs_batch, targets_batch, np.ones_like(targets_batch))
    # We want to test rng propagation too, so adding some dropout layers.
    first_layer = tl.Serial(tl.Embedding(9, 4), tl.Dropout(0.5), tl.Dup())
    rev_layers1 = [tl.ReversibleHalfResidual(tl.Dense(4), tl.Dropout(0.2)),
                   tl.ReversibleSwap(),
                   tl.ReversibleHalfResidual(tl.Dropout(0.5), tl.Dense(4)),
                   tl.ReversibleSwap()]
    mid_layer = tl.Serial(tl.Add(), tl.Dense(4), tl.Dup())
    rev_layers2 = [tl.ReversibleHalfResidual(tl.Dense(4), tl.Dropout(0.3)),
                   tl.ReversibleSwap()]
    loss_layer = tl.Serial(tl.Concatenate(), tl.Dense(19), tl.Dropout(0.3),
                           tl.LogSoftmax(), tl.CrossEntropyLoss())
    model = tl.Serial([first_layer] + rev_layers1 + [mid_layer] +
                      rev_layers2 + [loss_layer])
    rng_init = fastmath.random.get_prng(12)
    model.init(labeled_batch, rng=rng_init)
    optimizer_fn = optimizers.Adam  # to test slots

    # Make 3 steps with the original trainer.
    optimizer = optimizer_fn()
    optimizer.tree_init(model.weights)
    trainer = optimizers.Trainer(model, optimizer)
    rng_step1 = fastmath.random.get_prng(7)
    rng_step2 = fastmath.random.get_prng(8)
    rng_step3 = fastmath.random.get_prng(9)
    trainer.one_step(labeled_batch, rng_step1)
    trainer.one_step(labeled_batch, rng_step2, learning_rate=0.02)
    trainer.one_step(labeled_batch, rng_step3, learning_rate=0.03)
    first_layer_weights1 = first_layer.weights
    rev_layer12_weights1 = rev_layers1[2].weights
    mid_layer_weights1 = mid_layer.weights
    rev_layer20_weights1 = rev_layers2[0].weights
    loss_layer_weights1 = loss_layer.weights

    # Now make 3 steps with reversible trainer.
    model.init(labeled_batch, rng=rng_init)
    # TODO(lukaszkaiser): this test seems to fail with memoize_jit, why?
    trainer = optimizers.ReversibleSerialTrainer(
        [(first_layer.sublayers, rev_layers1),
         (mid_layer.sublayers, rev_layers2)],
        loss_layer, optimizer_fn, memoize_jit=False)
    trainer.one_step(labeled_batch, rng_step1)
    trainer.one_step(labeled_batch, rng_step2, learning_rate=0.02)
    trainer.one_step(labeled_batch, rng_step3, learning_rate=0.03)

    # Check that weights end up the same.
    self._assert_all_equal(loss_layer_weights1, loss_layer.weights)
    self._assert_all_equal(rev_layer20_weights1, rev_layers2[0].weights)
    self._assert_all_equal(mid_layer_weights1, mid_layer.weights)
    self._assert_all_equal(rev_layer12_weights1, rev_layers1[2].weights)
    self._assert_all_equal(first_layer_weights1, first_layer.weights)
Runs the reversible trainer, check results are the same as default.
trax/optimizers/trainer_test.py