code: def _terraformer_decoding_time(self, settings):
      # Garbage collection influences the timing, so we turn it off.
      gc.disable()
      max_len = 16
  
      def _self_attention_fn():
          return functools.partial(
              tl.SelfAttention,
              predict_drop_len=2 * max_len,
              predict_mem_len=2 * max_len,
          )
  
      def _causal_attention_fn():
          attn_layer, attn_kwargs = settings["attn"]
          return functools.partial(
              attn_layer,
              max_inference_length=2 * max_len,
              **attn_kwargs,
          )
  
      if settings["model"] == "terraformer":
          pred_model = models.ConfigurableTerraformer(
              mode="predict",
              d_model=settings["d_model"],
              d_ff=settings["d_ff"],
              dropout=0.1,
              max_len=max_len,
              n_heads=settings["n_heads"],
              n_encoder_layers=settings["encoder_layers"],
              n_decoder_layers=settings["decoder_layers"],
              encoder_attention_type=_self_attention_fn(),
              encoder_decoder_attention_type=_causal_attention_fn(),
              input_vocab_size=settings["vocab"],
              ff_sparsity=settings["ff_sparsity"],
              ff_use_sru=settings["ff_use_sru"],
              ff_dropout=0.1,
              # ff_chunk_size=1024,
              # attention_chunk_size=1,
              n_decoder_attention_layers=settings["attention_layers"],
              loss_sparsity=settings["loss_sparsity"],
              pos_axial_shape=None,
              use_bfloat16=True,
          )
      elif settings["model"] == "transformer":
          pred_model = models.ConfigurableTransformer(
              mode="predict",
              d_model=settings["d_model"],
              d_ff=settings["d_ff"],
              dropout=0.1,
              max_len=max_len,
              n_heads=settings["n_heads"],
              n_encoder_layers=settings["encoder_layers"],
              n_decoder_layers=settings["decoder_layers"],
              # encoder_attention_type=_self_attention_fn(),
              encoder_decoder_attention_type=_causal_attention_fn(),
              input_vocab_size=settings["vocab"],
              ff_sparsity=settings["ff_sparsity"],
              ff_use_sru=settings["ff_use_sru"],
              # ff_dropout=0.1,
              # ff_chunk_size=1024,
              # attention_chunk_size=1,
              # n_decoder_attention_layers=settings['attention_layers'],
              loss_sparsity=settings["loss_sparsity"],
              pos_axial_shape=None,
              # enc_dec_attention_sparsity=settings['enc_dec_sparsity'],
              # use_bfloat16=True,
          )
      else:
          assert False
      # We put acceleration outside of autoregressive_sample_stream, because
      # we want to have a separate run (separate input) for model compilation.
      pred_model = tl.Accelerate(pred_model)
  
      shape11 = shapes.ShapeDtype((1, 1), dtype=np.int32)
      shape1l = shapes.ShapeDtype((1, max_len), dtype=np.int32)
      pred_model.init(input_signature=(shape1l, shape11))
      original_state = copy.deepcopy(pred_model.state)
  
      inputs_warmup = np.zeros((1, max_len), dtype=np.int32)
      inputs = np.arange(max_len, dtype=np.int32).reshape(1, max_len)
  
      # This is a warm-up run, for compilation.
      result, current_time = [], time.time()
      elapsed_warmup_times = []
      for index, sample in zip(
          range(0, 4),
          decoding.autoregressive_sample_stream(
              pred_model,
              inputs_warmup,
              temperature=0.0,
              accelerate=False,
          ),
      ):
          del index  # unused
          result.append(
              sample[:, None]
          )  # to be sure that the result is computed
  
          current_time, start_time = time.time(), current_time
          elapsed_warmup_times.append(current_time - start_time)
  
      # This is a real decoding timing run that we measure.
      pred_model.state = original_state
      result, current_time = [], time.time()
      elapsed_times = []
      for index, sample in zip(
          range(12),
          decoding.autoregressive_sample_stream(
              pred_model, inputs, temperature=0.0, accelerate=False
          ),
      ):
          del index  # unused
          result.append(
              sample[:, None]
          )  # to be sure that the result is computed
  
          current_time, start_time = time.time(), current_time
          elapsed_times.append(current_time - start_time)
      peak_memory = _memory_usage()
  
      if min(elapsed_times[2:]) * 2 < max(elapsed_times[2:]):
          print(
              "WARNING! High variance found in elapsed times! Settings: {} ; "
              "elapsed times: {} ; Probably more warm-up steps should be used, "
              "or model size should be increased.".format(
                  settings, elapsed_times
              )
          )
      # Check resulting shapes.
      s = np.concatenate(result, axis=1)
      self.assertEqual(s.shape[0], 1)
      self.assertEqual(s.shape[1], 12)
      model_size = int(_size_of_model(pred_model))
  
      # We delete the model weights, because in some situations they won't be
      # deleted automatically.
      _recurrent_delete(pred_model.weights)
      gc.enable()
      return model_size, elapsed_times, peak_memory
  
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 1 for calling the DecodingTimingTest obj's underlying member method code with example usage: obj._terraformer_decoding_time(settings) and returns: model_size,  elapsed_times,  peak_memory, functools, functools