class_docstr: Mock TransformerLM for testing autoregressive sampling routines.
  
  Mimics the behavior of a perfectly-trained, deterministic TransformerLM.
  Allows to specify the \sigma^* -> \sigma function implemented by the model
  and to make assertions about the input sequence passed to the model.
  
  Supports two modes: stateful "predict" for fast inference, and stateless
  non-"predict" ("train", "eval" etc).
  
  Useful for testing any logic that relies on autoregressive sampling, as it
  removes the additional layer of complexity related to training a model or
  maintaining a pretrained one. Makes the tests run MUCH faster.
  
  Does not support acceleration. Do not wrap in tl.Accelerate().
class_docstr_tok: ["Mock","TransformerLM","for","testing","autoregressive","sampling","routines",".","Mimics","the","behavior","of","a","perfectly-trained",",","deterministic","TransformerLM",".","Allows","to","specify","the","\\sigma^","*","-",">","\\sigma","function","implemented","by","the","model","and","to","make","assertions","about","the","input","sequence","passed","to","the","model",".","Supports","two","modes",":","stateful","``","predict","''","for","fast","inference",",","and","stateless","non-","''","predict","''","(","``","train","''",",","``","eval","''","etc",")",".","Useful","for","testing","any","logic","that","relies","on","autoregressive","sampling",",","as","it","removes","the","additional","layer","of","complexity","related","to","training","a","model","or","maintaining","a","pretrained","one",".","Makes","the","tests","run","MUCH","faster",".","Does","not","support","acceleration",".","Do","not","wrap","in","tl.Accelerate","(",")","."]
code: def assert_prediction_buffers_equal(self, expected_buffers):
      if self._prediction_buffers is None:
          batch_size = expected_buffers.shape[0]
          actual_buffers = np.empty((batch_size, 0))
      else:
          actual_buffers = np.array(self._prediction_buffers)
  
      np.testing.assert_array_equal(actual_buffers, expected_buffers)
  
file_path: trax/trax/layers/test_utils.py
imports: import copy
  import functools
  import numpy
  from trax import fastmath
  from trax import layers
  from trax import shapes
label_desc: Use image node_id 8 for calling the MockTransformerLM obj's underlying member method code with example usage: obj.assert_prediction_buffers_equal(expected_buffers) without return types