code: def test_autoregressive_sample_terraformer_timing(self):
      template_to_use = "medium_transformer"
  
      settings_templates = {
          # full model
          # # 54B params
          # 'full_model': {
          #     'encoder_layers': 6, 'decoder_layers': 36, 'vocab': 32000,
          #     'attention_layers': 2,
          #     'd_ff': 64*1024, 'd_model': 96*96, 'n_heads': 96,
          #     'ff_use_sru': (1, 64), 'ff_sparsity': (256, 32),
          #     'loss_sparsity': 8,
          #     'attn': (tl.MultiplicativeConvCausalAttention,
          #              {'length_kernel_size': 3, 'sparsity': 64})},
          # 1/18 of model (1/6 of encoder, 1/18 of decoder, full vocab)
          # 4B params
          # 'big_terraformer': {
          #     'model': 'terraformer',
          #     'encoder_layers': 1, 'decoder_layers': 2, 'vocab': 32000,
          #     'attention_layers': 2,
          #     'd_ff': int(5/8 * 64*1024), 'd_model': 96*96, 'n_heads': 96,
          #     'ff_use_sru': 0, 'ff_sparsity': 0, 'loss_sparsity': 0,
          #     'attn': (tl.CausalAttention, {})},
          # 'big_transformer': {
          #     'model': 'transformer',
          #     'encoder_layers': 1, 'decoder_layers': 2, 'vocab': 32000,
          #     'attention_layers': 2,
          #     'd_ff': int(5/8 * 64*1024), 'd_model': 96*96, 'n_heads': 96,
          #     'ff_use_sru': 0, 'ff_sparsity': 0, 'loss_sparsity': 0,
          #     'attn': (tl.CausalAttention, {})},
          # medium model
          # 275M params (only decoder)
          "medium_transformer": {
              "model": "transformer",
              "encoder_layers": 2,
              "decoder_layers": 24,
              "vocab": 32000,
              "attention_layers": 2,
              "d_ff": 4 * 1024,
              "d_model": 1024,
              "n_heads": 16,
              "ff_use_sru": 0,
              "ff_sparsity": 0,
              "loss_sparsity": 0,
              "attn": (tl.CausalAttention, {}),
          },
          # 'medium_terraformer': {
          #     'model': 'terraformer',
          #     'encoder_layers': 2, 'decoder_layers': 24, 'vocab': 32000,
          #     'attention_layers': 2,
          #     'd_ff': 4*1024, 'd_model': 1024, 'n_heads': 16,
          #     'ff_use_sru': 0, 'ff_sparsity': 0, 'loss_sparsity': 0,
          #     'attn': (tl.CausalAttention, {})},
      }
  
      sweep_settings = {
          # 'big_transformer': [  # for big
          #     dict(), # baseline
          #     {'ff_sparsity': (256, 32)},  # + Sparse FF
          #     {'attn': (  # + Sparse QKV
          #         tl.MultiplicativeConvCausalAttention,
          #         {'length_kernel_size': 3, 'sparsity': 64}),
          #      'd_ff': 64*1024,
          #      },
          #     {'ff_sparsity': (256, 32),
          #      'attn': (  # + Sparse FF+QKV
          #         tl.MultiplicativeConvCausalAttention,
          #         {'length_kernel_size': 3, 'sparsity': 64}),
          #      'd_ff': 64*1024,
          #      },
          # ],
          "medium_transformer": [  # for medium
              dict(),  # baseline
              {
                  "ff_sparsity": 64,
                  "attn": (  # Sparse FF+QKV
                      tl.MultiplicativeConvCausalAttention,
                      {"length_kernel_size": 3, "sparsity": 16},
                  ),
                  "d_ff": 6 * 1024,
              },
              # {'ff_sparsity': 64,  # Sparse FF+QKV + Loss
              #  'attn': (
              #     tl.MultiplicativeConvCausalAttention,
              #     {'length_kernel_size': 3, 'sparsity': 16}),
              #  'd_ff': 6*1024,
              #  'loss_sparsity': 4,
              #  },
              # {'attn': (  # Sparse QKV
              #     tl.MultiplicativeConvCausalAttention,
              #     {'length_kernel_size': 3, 'sparsity': 16}),
              #  'd_ff': 6*1024,
              #  },
              # {'loss_sparsity': 4},  # Sparse Loss
              # {'ff_sparsity': 64},  # Sparse FF
              # {'ff_sparsity': 128},  # + Sparse FF 128
              # APPENDIX below
              # different loss layers
              # {'loss_sparsity': 8},
              # {'loss_sparsity': 2},
              # {'loss_sparsity': 0},
          ],
          #  'big_terraformer': [  # for big terraformer
          #      dict(), # baseline
          #      {'ff_sparsity': 64},  # + Sparse FF  / Sparse FF 64
          #      {'ff_sparsity': 64,
          #       'attn': (  # + Sparse FF+QKV
          #          tl.MultiplicativeConvCausalAttention,
          #          {'length_kernel_size': 3, 'sparsity': 16}),
          #       'd_ff': 6*1024,
          #       },
          #      {'ff_sparsity': 64,  # + Sparse FF+QKV+Loss
          #       'attn': (
          #          tl.MultiplicativeConvCausalAttention,
          #          {'length_kernel_size': 3, 'sparsity': 16}),
          #       'd_ff': 6*1024,
          #       'loss_sparsity': 4,
          #       },
          #         ],
          # 'medium_terraformer': [  # for medium terraformer
          #     {'ff_sparsity': 64,  # + Sparse FF+QKV+Loss
          #      'attn': (
          #         tl.MultiplicativeConvCausalAttention,
          #         {'length_kernel_size': 3, 'sparsity': 16}),
          #      'd_ff': 6*1024,
          #      'loss_sparsity': 4,
          #      },
          # ],
      }
  
      encoding_times = []
      decoding_times = []
      sizes = []
      memories = []
      messages = []
      for override_settings in sweep_settings[template_to_use]:
          settings = copy.deepcopy(settings_templates[template_to_use])
          settings.update(override_settings)
  
          init_memory = _memory_usage()
          (
              size,
              elapsed_times,
              peak_memory,
          ) = self._terraformer_decoding_time(settings)
  
          # TODO(jaszczur): Why is elapsed_times[0] always small?
          encoding_time = elapsed_times[1]
          decoding_time_10 = sum(elapsed_times[2:])
  
          after_memory = _memory_usage()
          model_memory_gigabytes = (
              peak_memory - init_memory
          ) / 1024**3
          decoding_time_diff = (
              max(elapsed_times[2:]) - min(elapsed_times[2:])
          ) / 2
          decoding_time_diff_percent = int(
              decoding_time_diff / np.mean(elapsed_times) * 100
          )
          message = (
              "\n\n"
              "Params: {}\n"
              "Settings: {}\n"
              "Override: {}\n"
              "Init memory: {:.1f} GiB\n"
              "Peak memory: {:.1f} GiB\n"
              "After memory: {:.1f} GiB\n"
              "Estimated model memory: {:.1f} GiB\n"
              "Times for each step: {}\n"
              "Time for encoding: {:.4f} s\n"
              "Time for decoding 10 tokens: {:.4f} s +/- {} %\n"
              "\n\n".format(
                  size,
                  settings,
                  override_settings,
                  init_memory / 1024**3,
                  peak_memory / 1024**3,
                  after_memory / 1024**3,
                  model_memory_gigabytes,
                  elapsed_times,
                  encoding_time,
                  decoding_time_10,
                  decoding_time_diff_percent,
              )
          )
          print(message)
          messages.append(message)
          encoding_times.append(encoding_time)
          decoding_times.append(decoding_time_10)
          sizes.append(size)
          memories.append(model_memory_gigabytes)
  
      print("Final results (recap):")
      for message in messages:
          print(message)
  
imports: import copy
  import functools
  import gc
  import os
  import time
  from jax import config
  import numpy
  import psutil
  from tensorflow.compat.v2 import test
  from trax import fastmath
  from trax import layers
  from trax import models
  from trax import shapes
  from trax.supervised import decoding
label_desc: Use image node_id 2 for calling the DecodingTimingTest obj's underlying member method code with example usage: obj.test_autoregressive_sample_terraformer_timing() without return types