class_docstr: Adaptive Computation Time module, based on: arxiv.org/abs/1807.03819.
class_docstr_tok: ["Adaptive","Computation","Time","module",",","based","on",":","arxiv.org\/abs\/1807.03819","."]
code: def __call__(self, x: jnp.ndarray, *layer_call_args):
      threshold = 1.0 - self.ac_config.act_epsilon
      max_steps = self.ac_config.act_max_steps
  
      state = x
      original_state_shape = state.shape
  
      if self.ac_config.act_level == "per_example":
          state_slice = slice(0, 1)
      elif self.ac_config.act_level == "per_token":
          state_slice = slice(0, 2)
      else:
          raise ValueError(
              f"Unknown act_level {self.ac_config.act_level}"
          )
  
      # Dynamic shape for update tensors below.
      update_shape = state.shape[state_slice]
      # Halting probabilities (p_t^n in the paper).
      halting_probability = jnp.zeros(update_shape)
      # Remainders (R(t) in the paper).
      remainders = jnp.zeros(update_shape)
      # Number of updates performed (N(t) in the paper).
      n_updates = jnp.zeros(update_shape)
      # Previous cell states (s_t in the paper).
      previous_state = jnp.zeros_like(state)
  
      # Define one stop function to decide the routing result.
      def stop_fn(inputs: Any) -> jnp.ndarray:
          # Returns True if all of halting probability >= 1-eps.
          _, halting_probability, _, _, _, *_ = inputs
          return jnp.all(
              jnp.greater_equal(halting_probability, threshold)
          )
  
      # Run max_steps, for each sample/token, when the decision is True,
      # go to the shunt_layer.
      intermedia_output = (
          state,
          halting_probability,
          remainders,
          n_updates,
          previous_state,
          *layer_call_args,
      )
  
      if self.share_parameters:
          # Scan over `ACTFunction` while broadcasing (sharing) the params.
          act_fn = nn.scan(
              ACTFunction,
              variable_broadcast="params",
              split_rngs={"params": False, "dropout": True},
              length=max_steps,
          )
  
      else:
          # When we want to have different parameters for different layers, if we
          # use simple nn.scan and set "variable_broadcast=None", we also get
          # different parameters for the haulting mechaniems (the dense layer the
          # predicts the halting probs) for different layers, however we want to
          # have the same mdoule make the halting decision across all layers.
          # To do that we need to map variables to two collections: `params` and
          # `shared_params` before sending it to the scan, then set
          # variable_broadcast='shared_params', and then map them back to a single
          # collection.
          def trans_in_fn(target):
              return {
                  "params": dict(
                      target.get("params", {}),
                      **target.get("shared_params", {}),
                  )
              }
  
          def trans_out_fn(target):
              params = target.get("params", {})
              shared_params = {}
              if "act_step" in params:
                  shared_params["act_step"] = params.pop("act_step")
              return {"params": params, "shared_params": shared_params}
  
          act_fn_two_collections = nn.scan(
              # Map  params to a two collections.
              nn.map_variables(
                  ACTFunction,
                  ["params", "shared_params"],
                  trans_in_fn=trans_in_fn,
                  trans_out_fn=trans_out_fn,
                  mutable=True,
              ),
              variable_broadcast="shared_params",
              variable_axes={"params": 0},
              split_rngs={"params": False, "dropout": True},
              length=max_steps,
          )
  
          # Map all params back to a single collection.
          act_fn = nn.map_variables(
              act_fn_two_collections,
              ["params", "shared_params"],
              trans_in_fn=trans_out_fn,
              trans_out_fn=trans_in_fn,
              mutable=True,
          )
  
      output, _ = act_fn(self.ac_config, self.layer, stop_fn)(
          intermedia_output, None
      )
  
      (
          output_state,
          halting_probability,
          remainders,
          ponder_times,
          new_state,
          *layer_call_args,
      ) = output
  
      # Check some shapes
      assert (
          output_state.shape == new_state.shape == original_state_shape
      )
      for x in [halting_probability, remainders, n_updates]:
          assert x.shape == original_state_shape[state_slice]
      return new_state, (ponder_times, remainders)
  
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 AdaptiveComputationTime obj's underlying member method code with example usage: obj.__call__(x) and returns: new_state,  ponder_times,  remainders, jnp, dict, dict