class_docstr: Takes an ACT step.
class_docstr_tok: ["Takes","an","ACT","step","."]
docstr: An pondernet act step.
  
  Args:
    carry_inputs: A tuple of: - state: An array of shape `[batch_size, length,
      channel]`. - unhalting_probability: An array containing the unhalting
      probabilities. - halted: An array containing halted flag of each token.
      - layer_id: An int denotes the layer id - layer_call_args: Arguments to
      be passed to the self.layer.
  
  Returns:
    A tupe of (output_state, unhalting_probability, halted).
docstr_tok: ["An","pondernet","act","step",".","Args",":","carry_inputs",":","A","tuple","of",":","-","state",":","An","array","of","shape","`","[","batch_size",",","length",",","channel","]","`",".","-","unhalting_probability",":","An","array","containing","the","unhalting","probabilities",".","-","halted",":","An","array","containing","halted","flag","of","each","token",".","-","layer_id",":","An","int","denotes","the","layer","id","-","layer_call_args",":","Arguments","to","be","passed","to","the","self.layer",".","Returns",":","A","tupe","of","(","output_state",",","unhalting_probability",",","halted",")","."]
code: def __call__(self, carry_inputs: Any) -> Any:
      """An pondernet act step.
  
      Args:
        carry_inputs: A tuple of: - state: An array of shape `[batch_size, length,
          channel]`. - unhalting_probability: An array containing the unhalting
          probabilities. - halted: An array containing halted flag of each token.
          - layer_id: An int denotes the layer id - layer_call_args: Arguments to
          be passed to the self.layer.
  
      Returns:
        A tupe of (output_state, unhalting_probability, halted).
      """
      halting_bias_init = self.ac_config.act_halting_bias_init
  
      # Unpack the inputs.
      (
          output_state,
          unhalting_probability,
          halted,
          layer_id,
          _,
          current_state,
          all_states,
          all_p,
          n_updates,
          *layer_call_args,
      ) = carry_inputs
  
      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",
          )(current_state)
      )
  
      # Average over all tokens:
      p = jnp.mean(p, axis=1)
      p = jnp.squeeze(p, axis=-1)
  
      # Update p if this is the last layer.
      # if layer_id == self.act_config.act_max_steps - 1:
      #  p = 1 - jnp.sum(all_p, axis=0)
      p = jax.lax.cond(
          layer_id == self.ac_config.act_max_steps - 1,
          lambda: 1 - jnp.sum(all_p, axis=0),
          lambda: p,
      )
  
      # Init the prob of halting here.
      prob_halt_here = unhalting_probability * p
  
      # Update unhalting_probability according to the new p.
      unhalting_probability = unhalting_probability * (1 - p)
  
      # Init the halting decision by sampling from bernoulli distribution.
      rng = self.make_rng("ponder")
      halt_decision = (1 - halted) * jax.random.bernoulli(
          rng, p=p, shape=p.shape
      )
  
      # Apply the layer on the state. And update output_state by the halted mask.
      new_state = self.layer(current_state, *layer_call_args)
  
      update_halted_decision = jnp.expand_dims(halt_decision, -1)
      update_halted_decision = jnp.expand_dims(
          update_halted_decision, -1
      )
      output_state = output_state + new_state * update_halted_decision
  
      # Update the all states and all p.
      all_states = all_states.at[layer_id].set(
          all_states[layer_id] + new_state
      )
      all_p = all_p.at[layer_id].set(all_p[layer_id] + prob_halt_here)
  
      # Update n_updates.
      n_updates += 1 - halted
  
      # Update the halted.
      halted = halted + halt_decision
  
      return (
          output_state,
          unhalting_probability,
          halted,
          layer_id + 1,
          prob_halt_here,
          new_state,
          all_states,
          all_p,
          n_updates,
          *layer_call_args,
      )
  
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 PonderStep obj's underlying member method code with example usage: obj.__call__(carry_inputs) and returns: output_state,  unhalting_probability,  halted,  unknown,  prob_halt_here,  new_state,  all_states,  all_p,  n_updates,  unknown