class_docstr: TapeBank Module.
  
  Given the TAPE token embedding, returns the top-k tpe tokens from the bank.
  
  Attributes:
    ac_config: Configuration of the adaptive computation.
    features: Number of feature dimensions for each embedding.
    tape_init: Tape initializer.
    dtype: The dtype of the embedding vectors (default: float32).
code: def setup(self):
      self.split_tt = self.ac_config.split_tt
      self.num_tape_tokens = self.ac_config.num_tape_tokens
      self.tape_bank_size = self.ac_config.tape_bank_size
      if self.ac_config.bank_type == "learn":
          self.tape_tokens = self.param(
              "tape_tokens",
              self.tape_init,
              (self.tape_bank_size, self.features),
              self.dtype,
          )
      self.norm_layer = nn.LayerNorm(name="bank_norm", dtype=self.dtype)
      self.dy_config = self.ac_config.dynamic_tape_length
  
imports: from typing import Any, Callable, Optional, Sequence, Tuple
  from absl import logging
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling the TapeBank obj's underlying member method code with example usage: obj.setup() without return types