code: def test_autoregressive_sample_terraformer_lsh_attn_quality(self):
      gin.add_config_file_search_path(_CONFIG_DIR)
      max_len = (
          32  # 32 is the max length we trained the checkpoint for.
      )
      test_lengths = [8, 16, 32]
      vocab_size = 13
      # The checkpoint is correct on ~90% sequences, set random seed to deflake.
      np.random.seed(0)
      for test_len in test_lengths:
          gin.clear_config()
          gin.parse_config_file("terraformer_copy.gin")
          gin.bind_parameter(
              "LSHSelfAttention.predict_mem_len", 2 * max_len
          )
          gin.bind_parameter(
              "LSHSelfAttention.predict_drop_len", 2 * max_len
          )
  
          pred_model = models.ConfigurableTerraformer(mode="predict")
  
          shape11 = shapes.ShapeDtype((1, 1), dtype=np.int32)
          shape1l = shapes.ShapeDtype((1, max_len), dtype=np.int32)
  
          model_path = os.path.join(
              _TESTDATA, "terraformer_copy_lsh_attn.pkl.gz"
          )
          pred_model.init_from_file(
              model_path,
              weights_only=True,
              input_signature=(shape1l, shape11),
          )
          initial_state = pred_model.state
  
          for _ in range(
              2
          ):  # Set low to make the test run reasonably fast.
              # Pick a length in [1, test_len] at random.
              inp_len = np.random.randint(low=1, high=test_len + 1)
              inputs = np.random.randint(
                  low=1, high=vocab_size - 1, size=(1, max_len)
              )
              # TODO(jaszczur): properly fix padding in terraformer predict mode,
              # and add a test here.
              s = decoding.autoregressive_sample(
                  pred_model,
                  inputs=inputs,
                  eos_id=-1,
                  max_length=inp_len,
                  temperature=0.0,
              )
              np.testing.assert_equal(s[0], inputs[0, :inp_len])
              pred_model.state = initial_state
      gin.clear_config()
  
imports: import functools
  import os
  import gin
  from jax import config
  import numpy
  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 13 for calling the DecodingTest obj's underlying member method code with example usage: obj.test_autoregressive_sample_terraformer_lsh_attn_quality() without return types