def test_run_reversible_same_as_default_terraformer(self):
    """Runs the reversible trainer, check results are the same as default."""
    inputs_batch = np.arange(8).reshape((2, 4)) + 1
    targets_batch = 2 * inputs_batch
    labeled_batch = (inputs_batch, targets_batch, np.ones_like(targets_batch))
    int_sig = shapes.ShapeDtype((2, 4), dtype=np.int32)
    input_sig = (int_sig, int_sig, int_sig)
    # We want to test rng propagation too, so adding some dropout layers.
    model = terraformer.ConfigurableTerraformer(
        20, d_model=8, d_ff=32, n_heads=1, dropout=0.0, n_encoder_layers=2,
        n_decoder_layers=2, ff_sparsity=(4, 8, 0.0, 1.0), pos_type=None,
        reversible_encoder=True)
    loss = tl.Serial(tl.LogSoftmax(), tl.CrossEntropyLoss())
    optimizer_fn = optimizers.Adafactor
    blocks, loss_layer = optimizers.trainer.extract_reversible_blocks(
        [model, loss], loss_chunk_size=4)
    blocks_serial = [(tl.Serial(std), rev) for (std, rev) in blocks]
    model_with_loss = tl.Serial(model, loss)
    rng_init = fastmath.random.get_prng(12)
    model_with_loss.init(input_sig, rng=rng_init)

    # Make 3 steps with the original trainer.
    optimizer = optimizer_fn()
    optimizer.tree_init(model_with_loss.weights)
    trainer = optimizers.Trainer(model_with_loss, 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_weights = blocks_serial[0][0].weights
    first_rev_weights = blocks[0][1][0].weights
    loss_weights = loss_layer.weights

    # Now make 3 steps with reversible trainer.
    model_with_loss.init(input_sig, rng=rng_init)
    trainer = optimizers.ReversibleSerialTrainer(
        blocks, loss_layer, optimizer_fn)
    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_weights, loss_layer.weights)
    self._assert_all_equal(first_rev_weights, blocks[0][1][0].weights)
    self._assert_all_equal(first_weights, blocks_serial[0][0].weights)
Runs the reversible trainer, check results are the same as default.
trax/optimizers/trainer_test.py