class_docstr: Adaptive Computation Time Function to help we use nn.scan on ACT.
code: def act_step(self, x) -> Any:
      # Unpack the carry input.
      (query, halting_prob, remainders, n_updates, score_mask) = x
  
      # Init the still running mask.
      still_running = jnp.less(halting_prob, self.threshold).astype(
          jnp.float32
      )
  
      # Init the tape token keys.
      # Compute similarity for tape tokens and generate the topK index.
      if self.ac_config.bank_type == "learn":
          tape_token_keys = self.tape_tokens[
              :, : self.features // self.ac_config.split_tt
          ]
          scores = jnp.dot(query, tape_token_keys.T)
      elif self.ac_config.bank_type == "input":
          tape_token_keys = self.tape_tokens[
              :, :, : self.features // self.ac_config.split_tt
          ]
          tape_token_keys = jnp.transpose(tape_token_keys, [0, 2, 1])
          scores = jnp.matmul(
              jnp.expand_dims(query, 1), tape_token_keys
          )
          scores = jnp.squeeze(scores, axis=1)
      else:
          raise NotImplementedError
  
      topk_idndex_inpool = jax.lax.top_k(
          scores - score_mask * 1e9,
          self.ac_config.num_tape_tokens // self.threshold,
      )[1]
      # Select the weights from the scores and softmax.
      weights = jnp.take(scores, topk_idndex_inpool)
      weights = nn.softmax(weights / query.shape[-1] ** 0.5)
      # Compute the entropy or max value for loss function.
      if self.ac_config.dynamic_tape_length.act_loss_type == "entropy":
          entropy = 1.0 - jnp.sum(weights**2, axis=-1)
      else:
          entropy = 1.0 - jnp.max(weights, axis=-1)
  
      # Init the new halted mask.
      sum_weights = jnp.sum(
          weights[:, : self.num_token_per_step], axis=-1
      )
      new_halted = (
          jnp.greater_equal(
              halting_prob + sum_weights, self.threshold
          ).astype(jnp.float32)
          * still_running
      )
  
      # Update still running.
      still_running = still_running - new_halted
  
      # Update remainder.
      remainders = remainders + (new_halted + still_running) * entropy
  
      # Update halting_prob.
      halting_prob = halting_prob + sum_weights * still_running
      halting_prob += new_halted * (self.threshold - halting_prob)
  
      # Increment n_updates for all inputs which are still running.
      n_updates += still_running + new_halted
  
      # Take the new selected tokens from the token bank,
      # and merge them into single tape token.
      if self.ac_config.bank_type == "learn":
          token_selected_wo_merge = jnp.take(
              self.tape_tokens, topk_idndex_inpool, axis=0
          )
      elif self.ac_config.bank_type == "input":
          token_selected_wo_merge = jnp.take_along_axis(
              self.tape_tokens,
              jnp.expand_dims(topk_idndex_inpool, -1),
              axis=1,
          )
      else:
          raise NotImplementedError
  
      token_selected = token_selected_wo_merge * jnp.expand_dims(
          weights, -1
      )
      token_selected = jnp.sum(token_selected, axis=-2, keepdims=True)
  
      # Update score_mask according to the new selected tokens.
      if self.ac_config.bank_type == "learn":
          score_mask += jnp.sum(
              jax.nn.one_hot(
                  topk_idndex_inpool, self.tape_tokens.shape[0]
              ),
              axis=1,
          )
      elif self.ac_config.bank_type == "input":
          score_mask += jnp.sum(
              jax.nn.one_hot(
                  topk_idndex_inpool, self.tape_tokens.shape[1]
              ),
              axis=1,
          )
      else:
          raise NotImplementedError
  
      # Update the query.
      token_selected_keys = token_selected[
          :, :, : self.features // self.ac_config.split_tt
      ]
  
      # Different mode to update query.
      # If True, replace the query by avg of old query and tape keys;
      # If False, replace the query by the tape keys directly.
      if self.ac_config.dynamic_tape_length.complex_query:
          query = (query + jnp.mean(token_selected_keys, axis=1)) / 2.0
      else:
          query = jnp.mean(token_selected_keys, axis=1)
  
      return (
          query,
          halting_prob,
          remainders,
          n_updates,
          score_mask,
      ), token_selected
  
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 ACTFunction obj's underlying member method code with example usage: obj.act_step(x) and returns: query,  halting_prob,  remainders,  n_updates,  score_mask,  token_selected