class_docstr: Adaptive Computation Time Function to help we use nn.scan on ACT.
class_docstr_tok: ["Adaptive","Computation","Time","Function","to","help","we","use","nn.scan","on","ACT","."]
code: def setup(self):
      self.act_step = ActStep(
          ac_config=self.ac_config, layer=self.layer, name="act_step"
      )
  
file_path: scenic/scenic/projects/baselines/universal_transformer/layers.py
imports: from typing import Any
  from flax import linen
  import jax
  import jax.numpy
  import ml_collections
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