class_docstr: Adaptive Computation Time Function to help we use nn.scan on ACT.
code: def setup(self):
      # Init how many tokens we append at each time step.
      self.num_token_per_step = (
          self.ac_config.dynamic_tape_length.num_token_per_step
      )
      # Init max steps we need to ponder.
      self.max_steps = (
          self.ac_config.num_tape_tokens // self.num_token_per_step
      )
      # Init the threshold for act.
      self.threshold = self.ac_config.dynamic_tape_length.act_epsilon
  
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 ACTFunction obj's underlying member method code with example usage: obj.setup() without return types