class_docstr: ATRTapeToken Module.
  
  Given the TAPE token embedding, returns the top-k tpe tokens from the bank.
  
  Attributes:
    ac_config: ml_collections.ConfigDict to use adaptive config
    features: Number of feature dimensions for each embedding.
    tape_tokens: Tape initializer.
    dtype: The dtype of the embedding vectors (default: float32).
code: def __call__(
      self, query: jnp.ndarray, deterministic: bool
  ) -> jnp.ndarray:
      bs = query.shape[0]
      # Init the values for scan func.
      # Carry value required: query, halting_probability,
      # loss_atr, n_updates, score_mask.
      # Dynamic shape for update tensors below.
      update_shape = query.shape[:1]
      halting_probability = jnp.zeros(update_shape)
      # ATR loss term for each sample.
      loss_atr = jnp.zeros(update_shape)
      # Number of updates performed (N(t) in the paper).
      n_updates = jnp.zeros(update_shape)
      # Scan value required: None.
      score_mask = jnp.zeros(
          [update_shape[0], self.tape_tokens.shape[-2]]
      )
  
      # One trick to enhance AdaTape with Learnable bank.
      # We mask a subset of the whole bank during training.
      if (
          not deterministic
          and self.ac_config.dynamic_tape_length.bernoulli_p > 0.0
      ):
          # We split the rng, one for longer sequence,
          # another one is for early exit.
          rng = self.make_rng("dropout")
          _, rng = jax.random.split(rng)
          score_mask = score_mask + jax.random.bernoulli(
              rng,
              p=self.ac_config.dynamic_tape_length.bernoulli_p,
              shape=score_mask.shape,
          ).astype(jnp.float32)
  
      # Pack the carry input for nn.scan.
      carry_input = (
          query,
          halting_probability,
          loss_atr,
          n_updates,
          score_mask,
      )
  
      # Init act_fn.
      act_fn = nn.scan(
          ACTFunction,
          variable_broadcast="params",
          split_rngs={"params": False, "dropout": True},
          length=self.max_steps,
      )
  
      # Conduct the act_fn.
      carry_output, scan_out = act_fn(
          self.ac_config, self.features, self.tape_tokens, deterministic
      )(carry_input, None)
  
      # Reshapt the outputs of nn.scan.
      scan_out = jnp.transpose(scan_out, (1, 0, 2, 3))
      scan_out = jnp.reshape(
          scan_out,
          (bs, self.max_steps * self.num_token_per_step, self.features),
      )
      _, halting_probability, remainders, n_updates, _ = carry_output
  
      return scan_out, (remainders, n_updates)
  
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 ATRTapeAppender obj's underlying member method code with example usage: obj.__call__(query, deterministic) and returns: scan_out,  remainders,  n_updates