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).
docstr: Retrieve tape tokens from a bank.
  
  Args:
    query: Array with last dimension equal the feature depth `features` of
      the embedding of tape tokens.
    bank: Array with candidate tape tokens.
    deterministic: bool denotes training or not.
  
  Returns:
    Output which is embedded input data.  The output shape follows the input,
    with an additional `features` dimension appended.
docstr_tok: ["Retrieve","tape","tokens","from","a","bank",".","Args",":","query",":","Array","with","last","dimension","equal","the","feature","depth","`","features","`","of","the","embedding","of","tape","tokens",".","bank",":","Array","with","candidate","tape","tokens",".","deterministic",":","bool","denotes","training","or","not",".","Returns",":","Output","which","is","embedded","input","data",".","The","output","shape","follows","the","input",",","with","an","additional","`","features","`","dimension","appended","."]
code: def __call__(
      self,
      query: jnp.ndarray,
      bank: Optional[jnp.ndarray],
      deterministic: bool,
  ) -> jnp.ndarray:
      """Retrieve tape tokens from a bank.
  
      Args:
        query: Array with last dimension equal the feature depth `features` of
          the embedding of tape tokens.
        bank: Array with candidate tape tokens.
        deterministic: bool denotes training or not.
  
      Returns:
        Output which is embedded input data.  The output shape follows the input,
        with an additional `features` dimension appended.
      """
      # One trick to improve AdaTape with learnable bank.
      # Add noise into to query during training.
      if (
          not deterministic
          and self.dy_config
          and self.dy_config.query_noise > 0.0
      ):
          rng = self.make_rng("dropout")
          _, rng = jax.random.split(rng)
          query += self.dy_config.query_noise * jax.random.normal(
              rng, query.shape
          )
      if self.ac_config.bank_type == "learn":
          tape_tokens = jnp.asarray(self.tape_tokens, self.dtype)
      elif self.ac_config.bank_type == "input":
          tape_tokens = bank
      else:
          raise NotImplementedError
      # Norm query and bank with the same LayerNorm.
      query = self.norm_layer(query)
      tape_tokens = self.norm_layer(tape_tokens)
      # Split the tape token [:,:, feature_dim] into two sub-vectors, i.e. key
      # and value. We also use half of the query as the real query if so.
      feature_dim = query.shape[-1]
      query = query[:, : feature_dim // self.split_tt]
      # We use ATR for dynamic reading.
      if self.dy_config:
          return ATRTapeAppender(
              self.ac_config,
              self.features,
              tape_tokens,
              self.num_tape_tokens,
          )(query, deterministic)
      # `scores` is an array with final dim `tape_bank_size` corresponding to the
      # batched inner-product of the array of query vectors against
      # each tape token in tape_tokens.
  
      # When we do not consider adaptive length, we use the following code.
      if self.ac_config.bank_type == "learn":
          tape_tokens_keys = tape_tokens[
              :, : feature_dim // self.split_tt
          ]
          scores = jnp.dot(query, tape_tokens_keys.T)
          topk_idndex = jax.lax.top_k(scores, self.num_tape_tokens)[1]
          assert jnp.issubdtype(topk_idndex.dtype, jnp.integer)
          return jnp.take(tape_tokens, topk_idndex, axis=0), (
              None,
              None,
          )  # pytype: disable=bad-return-type  # jax-ndarray
      else:
          token_selected_keys = tape_tokens[
              :, :, : self.features // self.ac_config.split_tt
          ]
          token_selected_keys = jnp.transpose(
              token_selected_keys, [0, 2, 1]
          )
          scores = jnp.matmul(
              jnp.expand_dims(query, 1), token_selected_keys
          )
          scores = jnp.squeeze(scores, axis=1)
          topk_idndex = jax.lax.top_k(scores, self.num_tape_tokens)[1]
          assert jnp.issubdtype(topk_idndex.dtype, jnp.integer)
          topk_idndex = jnp.expand_dims(topk_idndex, axis=-1)
          tape_tokens = jnp.take_along_axis(
              tape_tokens, topk_idndex, axis=1
          )
          return tape_tokens, (None, None)
  
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 2 for calling the TapeBank obj's underlying member method code with example usage: obj.__call__(query, bank, deterministic) and returns: ATRTapeAppender, jnp,  None,  None, tape_tokens,  None,  None