class_docstr: Takes an ACT step.
class_docstr_tok: ["Takes","an","ACT","step","."]
docstr: An act step (which either adaptivaly appies the layer or skips).
  
  Args:
    inputs: A tuple of: - state: An array of shape `[batch_size, length,
      channel]`. - halting_probability: An array containing the halting probs.
      - remainders: An array containing the act remainders. - n_updates: An
      array containing the act n_updates. - previous_state: An array that has
      the previous state. - layer_call_args: Arguments to be passed to the
      self.layer.
  
  Returns:
    A tupe of (output_state, new halting_probabilities,
      updated remainders, updated n_updates, new_state).
docstr_tok: ["An","act","step","(","which","either","adaptivaly","appies","the","layer","or","skips",")",".","Args",":","inputs",":","A","tuple","of",":","-","state",":","An","array","of","shape","`","[","batch_size",",","length",",","channel","]","`",".","-","halting_probability",":","An","array","containing","the","halting","probs",".","-","remainders",":","An","array","containing","the","act","remainders",".","-","n_updates",":","An","array","containing","the","act","n_updates",".","-","previous_state",":","An","array","that","has","the","previous","state",".","-","layer_call_args",":","Arguments","to","be","passed","to","the","self.layer",".","Returns",":","A","tupe","of","(","output_state",",","new","halting_probabilities",",","updated","remainders",",","updated","n_updates",",","new_state",")","."]
code: def __call__(self, inputs: Any) -> Any:
      """An act step (which either adaptivaly appies the layer or skips).
  
      Args:
        inputs: A tuple of: - state: An array of shape `[batch_size, length,
          channel]`. - halting_probability: An array containing the halting probs.
          - remainders: An array containing the act remainders. - n_updates: An
          array containing the act n_updates. - previous_state: An array that has
          the previous state. - layer_call_args: Arguments to be passed to the
          self.layer.
  
      Returns:
        A tupe of (output_state, new halting_probabilities,
          updated remainders, updated n_updates, new_state).
      """
      threshold = 1.0 - self.ac_config.act_epsilon
      act_type = self.ac_config.act_type
      halting_bias_init = self.ac_config.act_halting_bias_init
      act_level = self.ac_config.act_level
  
      (
          state,
          halting_probability,
          remainders,
          n_updates,
          previous_state,
          *layer_call_args,
      ) = inputs
      if act_type == "random":
          # Random as halting probability, to be used as a baseline.
          rng = jax.random.PRNGKey(0)  # bind rng to step?
          # TODO(dehghani): currently, it gives the error of:
          # ScanACTFunction_0 needs PRNG for "dropout"!
          p = jax.random.uniform(rng, shape=halting_probability.shape)
  
      else:
          p = nn.sigmoid(
              nn.Dense(
                  features=1,
                  use_bias=True,
                  kernel_init=nn.initializers.zeros,
                  bias_init=lambda k, s, *_: jnp.full(
                      s, halting_bias_init
                  ),
                  dtype=jnp.float32,
                  name="step_halting_prob",
              )(state)
          )
  
          if act_level == "per_example":
              # Average over all tokens:
              p = jnp.mean(p, axis=1)
          p = jnp.squeeze(p, axis=-1)
  
      # Create a mask for inputs which have not halted yet.
      still_running = jnp.less(halting_probability, 1.0).astype(
          jnp.float32
      )
  
      # Create a mask for inputs which halted at this step.
      new_halted = (
          jnp.greater(
              halting_probability + p * still_running, threshold
          ).astype(jnp.float32)
          * still_running
      )
  
      # Crerate mask of inputs which haven't halted and didn't halt this step.
      still_running = (
          jnp.less_equal(
              halting_probability + p * still_running, threshold
          ).astype(jnp.float32)
          * still_running
      )
  
      # Add the halting probability for this step to the halting
      # probabilities for those inputs which haven't halted yet.
      halting_probability += p * still_running
  
      # Compute remainders for the inputs which halted at this step.
      remainders += new_halted * (1 - halting_probability)
  
      # Add the remainders to those inputs which halted at this step.
      halting_probability += new_halted * remainders
  
      # Increment n_updates for all inputs which are still running.
      n_updates += still_running + new_halted
  
      # Compute the weight to be applied to the new state and output:
      # 0: when the input has already halted.
      # p: when the input hasn't halted yet.
      # remainders: when it halted this step.
      update_weights = jnp.expand_dims(
          p * still_running + new_halted * remainders, -1
      )
      if act_level == "per_example":
          update_weights = jnp.expand_dims(update_weights, -1)
  
      # Apply the layer on the state.
      output_state = self.layer(state, *layer_call_args)
  
      if act_type in ["basic", "random"]:
          # Update running part in the weighted state and keep the rest
          new_state = (output_state * update_weights) + (
              previous_state * (1 - update_weights)
          )
      elif act_type == "accumulated":
          # Add in the weighted state.
          new_state = (output_state * update_weights) + previous_state
      else:
          raise ValueError(f"Unknown act_type {act_type}!")
  
      return (
          output_state,
          halting_probability,
          remainders,
          n_updates,
          new_state,
          *layer_call_args,
      )
  
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 ActStep obj's underlying member method code with example usage: obj.__call__(inputs) and returns: output_state,  halting_probability,  remainders,  n_updates,  new_state,  unknown