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):
      max_steps = self.ac_config.act_max_steps
      deterministic = layer_call_args[0]
  
      state = x
      original_state_shape = state.shape
  
      state_slice = slice(0, 1)
  
      # Dynamic shape for update tensors below.
      update_shape = state.shape[state_slice]
      # Unhalting probabilities.
      unhalting_probability = jnp.ones(update_shape)
      # Halted mask.
      halted = jnp.zeros(update_shape)
      # All states from different steps.
      all_states = jnp.zeros((max_steps,) + original_state_shape)
      # All p from different steps.
      all_p = jnp.zeros((max_steps,) + update_shape)
      # Count how many updates we did, we use this to log and debug
      n_updates = jnp.zeros(update_shape)
  
      # 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.
          _, _, halted, _, _, _, *_ = inputs
          return jnp.all(halted)
  
      # Create one empty_probability with unhalting_probability
      empty_probability = jnp.zeros_like(unhalting_probability)
      # empty_state = jnp.zeros_like(state)
  
      # Run max_steps, for each sample/token, when the decision is True,
      # go to the shunt_layer.
      scan_carry_input = (
          state,
          unhalting_probability,
          halted,
          0,
          empty_probability,
          state,
          all_states,
          all_p,
          n_updates,
          *layer_call_args,
      )
  
      if self.share_parameters:
          # Scan over `PonderFunction` while broadcasing (sharing) the params.
          act_fn = nn.scan(
              PonderFunction,
              variable_broadcast="params",
              split_rngs={
                  "params": False,
                  "dropout": True,
                  "ponder": True,
              },
              length=max_steps,
          )
      else:
          # Scan over PonderFunction while only broadcasing the shared param.
          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 "ponder_step" in params:
                  shared_params["ponder_step"] = params.pop(
                      "ponder_step"
                  )
              return {"params": params, "shared_params": shared_params}
  
          act_fn_two_collections = nn.scan(
              # Map params to a two collections.
              nn.map_variables(
                  PonderFunction,
                  ["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,
                  "ponder": 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,
          )
  
      scan_carry_output, _ = act_fn(
          self.ac_config,
          self.layer,
          stop_fn,
          deterministic,
      )(scan_carry_input, None)
  
      (
          final_state,
          unhalting_probability,
          halted,
          _,
          _,
          _,
          all_states,
          all_p,
          n_updates,
          *layer_call_args,
      ) = scan_carry_output
  
      # Check some shapes
      assert final_state.shape == original_state_shape
      for x in [unhalting_probability, halted]:
          assert x.shape == original_state_shape[state_slice]
      return final_state, (all_states, all_p, n_updates)
  
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 1 for calling the AdaptiveComputationTime obj's underlying member method code with example usage: obj.__call__(x) and returns: final_state,  all_states,  all_p,  n_updates, jnp, dict, dict