docstr: Runs the reversible trainer with a lot of weights to test memory use.
docstr_tok: ["Runs","the","reversible","trainer","with","a","lot","of","weights","to","test","memory","use","."]
code: def test_run_reversible_large_weights(self):
      """Runs the reversible trainer with a lot of weights to test memory use."""
      # This test requires > 18GB RAM, only run on TPUs. It does pass on GPU
      # and CPU when you run it locally, but it's too big for unit-testing.
      ram_limited = True  # Set to False to run this test locally.
      if fastmath.global_device_count() == 1 and ram_limited:
          return
  
      # Create inputs and rngs.
      inputs_batch = np.arange(8).reshape((2, 4))
      targets_batch = inputs_batch
      labeled_batch = (
          inputs_batch,
          targets_batch,
          np.ones_like(targets_batch),
      )
      first_layer = tl.Serial(tl.Embedding(9, 16 * 1024), tl.Dup())
      rng_init = fastmath.random.get_prng(12)
      rng_step = fastmath.random.get_prng(13)
  
      # Initialize layers.
      first_layer.init(labeled_batch, rng=rng_init)
      n_layers = (
          18  # 18 layers each 16K x 16K = 256M weights ~= 1GB, 18GB ram
      )
      rev_layers = []
      int_shape = shapes.ShapeDtype((2, 4), dtype=np.int32)
      shape = shapes.ShapeDtype((2, 4, 16 * 1024))
      sig = (shape, shape)
      for _ in range(n_layers):
          layer = tl.ReversibleHalfResidual(tl.Dense(16 * 1024))
          layer.init(sig, rng=rng_init)
          layer.weights = tl.on_cpu(
              layer.weights
          )  # store weights in cpu memory
          rev_layers.append(layer)
          rev_layers.append(tl.ReversibleSwap())
      loss_layer = tl.Serial(
          tl.Concatenate(),
          tl.Dense(9),
          tl.LogSoftmax(),
          tl.CrossEntropyLoss(),
      )
      loss_layer.init((shape, shape, int_shape, int_shape))
      optimizer_fn = optimizers.Adafactor
  
      # Make a step with reversible trainer.
      trainer = optimizers.ReversibleSerialTrainer(
          [(first_layer, rev_layers)], loss_layer, optimizer_fn
      )
      loss, _ = trainer.one_step(labeled_batch, rng_step)
      self.assertLess(
          float(loss.sum()), 10000.0
      )  # Just to get the loss.
      # Set to true to run again, e.g., for profiling.
      run_twice = False
      if run_twice:
          t = time.time()
          loss, _ = trainer.one_step(labeled_batch, rng_step)
          self.assertLess(
              float(loss.sum()), 10000.0
          )  # Just to get the loss.
          print(
              "Took %.3f seconds to run, loss %s"
              % (time.time() - t, loss)
          )
  
imports: import time
  from absl.testing import absltest
  from jax import config
  import numpy
  from trax import fastmath
  from trax import layers
  from trax import optimizers
  from trax import shapes
  from trax.layers import base
  from trax.models.research import terraformer
label_desc: Use image node_id 8 for calling the TrainerTest obj's underlying member method code with example usage: obj.test_run_reversible_large_weights() without return types