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 __call__(self, x, _) -> Any:
      # We only consider take_a_step here, since for PonderNet, the skip only
      # happens during inference.
      if self.is_mutable_collection("params"):  # Init-mode
          out = self.take_a_step(x)
      else:
          decision = self.stop_fn(x) * self.deterministic
          out = nn.cond(
              decision, self.skip_a_step, self.take_a_step, self, x
          )
      return out, None
  
file_path: scenic/scenic/projects/baselines/pondernet/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 4 for calling the PonderFunction obj's underlying member method code with example usage: obj.__call__(x, _) and returns: out,  None