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 __call__(self, inputs, rng=None):
      del rng
  
      assert (
          inputs.ndim == 2
      ), "The input sequences should have exactly two axes."
  
      if self._prediction_buffers is None:
          # Initialize the buffer.
          batch_size = inputs.shape[0]
          # [[]] * batch_size would create multiple references to the same
          # list, and we want separate lists.
          self._prediction_buffers = [[] for _ in range(batch_size)]
  
      if self._mode == "predict":
          output_fn = self._output_symbol_predict
      else:
          output_fn = self._output_symbols_eval
  
      # Calculate the output separately for each sequence in the batch.
      output_symbols = np.array(
          [
              output_fn(input_seq, pred_buffer)
              for (input_seq, pred_buffer) in zip(
                  inputs, self._prediction_buffers
              )
          ]
      )
      return self._symbols_to_logits(output_symbols)
  
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 7 for calling the MockTransformerLM obj's underlying member method code with example usage: obj.__call__(inputs, rng) and returns: self