code: def test_mock_transformer_lm_decodes_arithmetic_sequence(self):
      model = test_utils.MockTransformerLM(
          sequence_fn=arithmetic_sequence,
          vocab_size=10,
          mode="predict",
      )
      output = decoding.autoregressive_sample(
          model, max_length=5, start_id=0, eos_id=-1, accelerate=False
      )
  
      # Sequence including the leading 0 and the last predicted symbol.
      full_seq = list(range(6))
      # decoding.autoregressive_sample doesn't return the leading 0.
      np.testing.assert_array_equal(output, [full_seq[1:]])
      # The prediction buffers don't include the last predicted symbol.
      model.assert_prediction_buffers_equal([full_seq[:-1]])
  
file_path: trax/trax/layers/test_utils_test.py
imports: import functools
  from absl.testing import absltest
  import numpy
  from trax.layers import test_utils
  from trax.supervised import decoding
label_desc: Use image node_id 2 for calling the TestUtilsTest obj's underlying member method code with example usage: obj.test_mock_transformer_lm_decodes_arithmetic_sequence() without return types