code: def test_mock_transformer_lm_rewinds(self):
      model = test_utils.MockTransformerLM(
          sequence_fn=arithmetic_sequence,
          vocab_size=10,
          mode="predict",
      )
      sample_3 = functools.partial(
          decoding.autoregressive_sample,
          max_length=3,
          eos_id=-1,
          accelerate=False,
      )
  
      # Generate the 3 initial symbols.
      init_output = sample_3(model, start_id=0)
      np.testing.assert_array_equal(init_output, [[1, 2, 3]])
      state = model.state
  
      # Generate the next 3 symbols.
      next_output = sample_3(model, start_id=init_output[0, -1])
      np.testing.assert_array_equal(next_output, [[4, 5, 6]])
  
      # Rewind and generate the last 3 symbols again.
      model.state = state
      next_output = sample_3(model, start_id=init_output[0, -1])
      np.testing.assert_array_equal(next_output, [[4, 5, 6]])
  
      # Check the buffers.
      model.assert_prediction_buffers_equal([[0, 1, 2, 3, 4, 5]])
  
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 3 for calling the TestUtilsTest obj's underlying member method code with example usage: obj.test_mock_transformer_lm_rewinds() without return types