class_docstr: One Token Turing Machine unit.
  
  This implements the operations in a TTM (https://arxiv.org/abs/2211.09119) at
  each step.
  
  Attributes:
    process_size: Number of tokens for the processing unit to process.
    memory_size: Number of memory tokens.
    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.
    mlp_dim: MLP dim size in the processing unit.
    num_heads: Number of heads in the processing unit.
    use_positional_embedding: Whether to use positional embeddings for the
      memory read/write.
    dropout_rate: Dropout rate.
docstr: Applies Token Turing Machine unit.
  
  Args:
    memory_tokens: Inputs of shape `[bs, memory_size, c]`.
    input_tokens: Inputs of shape `[bs, n_token, c]`.
    train: Weather we are in the training mode.
  
  Returns:
    Tuple of shape `([bs, memory_size, c], [bs, process_size, c])`.
docstr_tok: ["Applies","Token","Turing","Machine","unit",".","Args",":","memory_tokens",":","Inputs","of","shape","`","[","bs",",","memory_size",",","c","]","`",".","input_tokens",":","Inputs","of","shape","`","[","bs",",","n_token",",","c","]","`",".","train",":","Weather","we","are","in","the","training","mode",".","Returns",":","Tuple","of","shape","`","(","[","bs",",","memory_size",",","c","]",",","[","bs",",","process_size",",","c","]",")","`","."]
code: def __call__(
      self,
      memory_tokens: jnp.ndarray,
      input_tokens: jnp.ndarray,
      train: bool = False,
  ) -> Tuple[jnp.ndarray, jnp.ndarray]:
      """Applies Token Turing Machine unit.
  
      Args:
        memory_tokens: Inputs of shape `[bs, memory_size, c]`.
        input_tokens: Inputs of shape `[bs, n_token, c]`.
        train: Weather we are in the training mode.
  
      Returns:
        Tuple of shape `([bs, memory_size, c], [bs, process_size, c])`.
      """
      all_tokens = jnp.concatenate(
          [memory_tokens, input_tokens], axis=1
      )
  
      if self.use_positional_embedding:
          all_tokens = vit.AddPositionEmbs(
              posemb_init=nn.initializers.normal(
                  stddev=0.02
              ),  # from BERT.
              name="read_pos_embed",
          )(all_tokens)
  
      if self.memory_mode == "TL" or self.memory_mode == "TL-AddErase":
          all_tokens = token_learner_model.TokenLearnerModuleV11(
              self.process_size, dropout_rate=self.dropout_rate
          )(all_tokens, deterministic=not train)
      elif self.memory_mode == "TL-MHA":
          all_tokens = TokenLearnerMHA(self.process_size)(all_tokens)
  
      if self.processing_unit == "transformer":
          output_tokens = vit.Encoder(
              num_layers=self.num_layers,
              mlp_dim=self.mlp_dim,
              num_heads=self.num_heads,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.dropout_rate,
          )(all_tokens, train=train)
      elif self.processing_unit == "mixer":
          output_tokens = all_tokens
          for _ in range(self.num_layers):
              output_tokens = mixer.MixerBlock(
                  channels_mlp_dim=self.mlp_dim,
                  sequence_mlp_dim=128,
                  dropout_rate=self.dropout_rate,
                  stochastic_depth=0.0,
                  layer_scale=None,
              )(output_tokens, deterministic=not train)
          output_tokens = nn.LayerNorm()(output_tokens)
      elif self.processing_unit == "mlp":
          output_tokens = all_tokens
          for _ in range(self.num_layers):
              output_tokens = nn.LayerNorm()(output_tokens)
              output_tokens = attention_layers.MlpBlock(
                  mlp_dim=self.mlp_dim,
                  dropout_rate=self.dropout_rate,
                  activation_fn=nn.gelu,
              )(output_tokens, deterministic=not train)
          output_tokens = nn.LayerNorm()(output_tokens)
      else:
          raise ValueError(
              f"Unknown processing unit {self.processing_unit}."
          )
      mem_out_tokens = jnp.concatenate(
          [memory_tokens, input_tokens, output_tokens], axis=1
      )
  
      if self.use_positional_embedding:
          mem_out_tokens = vit.AddPositionEmbs(
              posemb_init=nn.initializers.normal(
                  stddev=0.02
              ),  # from BERT.
              name="write_pos_embed",
          )(mem_out_tokens)
  
      if self.memory_mode == "TL":
          mem_out_tokens = token_learner_model.TokenLearnerModuleV11(
              self.memory_size, dropout_rate=self.dropout_rate
          )(mem_out_tokens, deterministic=not train)
      elif self.memory_mode == "TL-MHA":
          mem_out_tokens = TokenLearnerMHA(self.memory_size)(
              mem_out_tokens
          )
      elif self.memory_mode == "TL-AddErase":
          mem_out_tokens = TokenAddEraseWrite()(
              memory_tokens, output_tokens, train
          )
  
      return (mem_out_tokens, 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 1 for calling the TokenTuringMachineUnit obj's underlying member method code with example usage: obj.__call__(memory_tokens, input_tokens, train) and returns: mem_out_tokens,  output_tokens