class_docstr: Token Turing Machine main model encoder.
  
  It implements https://arxiv.org/abs/2211.09119. It essentially repeats
  TokenTuringMachineUnit for the number of steps (of the input tensor).
  
  This version is for the training and inference with a fixed shaped, static
  input tensor. One will need to modify/extend this module together with the
  data pipeline for the streaming inference implementation.
  
  Attributes:
    process_size: Number of tokens for the Transformer to process.
    memory_size: The number of memory tokens in the TTM.
    memory_mode: Specifies the token summarization method to use. Supports
      'TL', 'TL-MHA', or 'TL-AddErase'.
    processing_unit: Specifies which processing unit module to use. Supports
      'transformer', 'mixer', or 'mlp'.
    num_layers: Number of layers in the processing unit.
docstr: Applies Token Turing Machine model.
  
  Args:
    input_tokens: Inputs of shape `[bs, num_steps, n_tokens, c]`.
    train: Weather we are in the training mode.
  
  Returns:
    Tensor of shape `[bs, num_steps, process_size, c]`.
docstr_tok: ["Applies","Token","Turing","Machine","model",".","Args",":","input_tokens",":","Inputs","of","shape","`","[","bs",",","num_steps",",","n_tokens",",","c","]","`",".","train",":","Weather","we","are","in","the","training","mode",".","Returns",":","Tensor","of","shape","`","[","bs",",","num_steps",",","process_size",",","c","]","`","."]
code: def __call__(
      self, input_tokens: jnp.ndarray, train: bool = False
  ) -> jnp.ndarray:
      """Applies Token Turing Machine model.
  
      Args:
        input_tokens: Inputs of shape `[bs, num_steps, n_tokens, c]`.
        train: Weather we are in the training mode.
  
      Returns:
        Tensor of shape `[bs, num_steps, process_size, c]`.
      """
      bs, ns, _, c = input_tokens.shape
  
      output_tokens_list = []
      memory_tokens = jnp.zeros([bs, self.memory_size, c])
  
      for i in range(ns):
          step_tokens = input_tokens[:, i, :, :]
  
          logging.info("step tokens: %s", step_tokens.shape)
          memory_tokens, output_tokens = self.unit_function(
              memory_tokens, step_tokens, train=train
          )
  
          logging.info("output_tokens: %s", output_tokens.shape)
          logging.info("memory_tokens: %s", memory_tokens.shape)
  
          output_tokens = jnp.expand_dims(output_tokens, axis=1)
          output_tokens_list.append(output_tokens)
  
      output_tokens = jnp.concatenate(output_tokens_list, axis=1)
  
      return output_tokens
  
imports: from typing import Tuple
  from absl import logging
  import flax.linen
  import jax
  import jax.numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.projects.baselines import mixer
  from scenic.projects.baselines import vit
  from scenic.projects.token_learner import model
label_desc: Use image node_id 2 for calling the TokenTuringMachineEncoder obj's underlying member method code with example usage: obj.__call__(input_tokens, train) and returns: output_tokens