code: def universal_transformer_act(x, hparams, ffn_unit, attention_unit):
      """ACT based models.
  
      Implementations of all act models are based on craffel@'s cl/160711592.
  
      (1) Basic AUT based on remainder-distribution ACT (position-wise).
      (2) AUT with global halting probability (not position-wise).
      (3) AUT with random halting probability (not position-wise).
      (4) AUT with final state as accumulation of all states.
  
      Args:
        x: input
        hparams: model hyper-parameters
        ffn_unit: feed-forward unit
        attention_unit: multi-head attention unit
  
      Returns:
        the output tensor,  (ponder_times, remainders)
  
      Raises:
        ValueError: Unknown act type
      """
      if hparams.act_type not in [
          "basic",
          "global",
          "random",
          "accumulated",
      ]:
          raise ValueError("Unknown act type: %s" % hparams.act_type)
  
      state = x
      act_max_steps = hparams.act_max_steps
      threshold = 1.0 - hparams.act_epsilon
      state_shape_static = state.get_shape()
  
      state_slice = slice(0, 2)
      if hparams.act_type == "global":
          state_slice = slice(0, 1)
  
      # Dynamic shape for update tensors below
      update_shape = tf.shape(state)[state_slice]
  
      # Halting probabilities (p_t^n in the paper)
      halting_probability = tf.zeros(
          update_shape, name="halting_probability"
      )
  
      # Remainders (R(t) in the paper)
      remainders = tf.zeros(update_shape, name="remainder")
  
      # Number of updates performed (N(t) in the paper)
      n_updates = tf.zeros(update_shape, name="n_updates")
  
      # Previous cell states (s_t in the paper)
      previous_state = tf.zeros_like(state, name="previous_state")
      step = tf.constant(0, dtype=tf.int32)
  
      def ut_function(
          state,
          step,
          halting_probability,
          remainders,
          n_updates,
          previous_state,
      ):
          """implements act (position-wise halting).
  
          Args:
            state: 3-D Tensor: [batch_size, length, channel]
            step: indicates number of steps taken so far
            halting_probability: halting probability
            remainders: act remainders
            n_updates: act n_updates
            previous_state: previous state
  
          Returns:
            transformed_state: transformed state
            step: step+1
            halting_probability: halting probability
            remainders: act remainders
            n_updates: act n_updates
            new_state: new state
          """
          state = step_preprocess(state, step, hparams)
  
          if hparams.act_type == "random":
              # random as halting probability
              p = tf.random_uniform(
                  shape=common_layers.shape_list(halting_probability)
              )
          else:
              with tf.variable_scope(
                  "sigmoid_activation_for_pondering"
              ):
                  p = common_layers.dense(
                      state,
                      1,
                      activation=tf.nn.sigmoid,
                      use_bias=True,
                      bias_initializer=tf.constant_initializer(
                          hparams.act_halting_bias_init
                      ),
                  )
  
                  if hparams.act_type == "global":
                      # average over all positions (as a global halting prob)
                      p = tf.reduce_mean(p, axis=1)
                      p = tf.squeeze(p)
                  else:
                      # maintain position-wise probabilities
                      p = tf.squeeze(p, axis=-1)
  
          # Mask for inputs which have not halted yet
          still_running = tf.cast(
              tf.less(halting_probability, 1.0), tf.float32
          )
  
          # Mask of inputs which halted at this step
          new_halted = (
              tf.cast(
                  tf.greater(
                      halting_probability + p * still_running, threshold
                  ),
                  tf.float32,
              )
              * still_running
          )
  
          # Mask of inputs which haven't halted, and didn't halt this step
          still_running = (
              tf.cast(
                  tf.less_equal(
                      halting_probability + p * still_running, threshold
                  ),
                  tf.float32,
              )
              * still_running
          )
  
          # Add the halting probability for this step to the halting
          # probabilities for those input 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
          # the remainders when it halted this step
          update_weights = tf.expand_dims(
              p * still_running + new_halted * remainders, -1
          )
          if hparams.act_type == "global":
              update_weights = tf.expand_dims(update_weights, -1)
  
          # apply transformation on the state
          transformed_state = state
          for i in range(hparams.num_inrecurrence_layers):
              with tf.variable_scope("rec_layer_%d" % i):
                  transformed_state = ffn_unit(
                      attention_unit(transformed_state)
                  )
  
          # update running part in the weighted state and keep the rest
          new_state = (transformed_state * update_weights) + (
              previous_state * (1 - update_weights)
          )
  
          if hparams.act_type == "accumulated":
              # Add in the weighted state
              new_state = (
                  transformed_state * update_weights
              ) + previous_state
  
          # remind TensorFlow of everything's shape
          transformed_state.set_shape(state_shape_static)
          for x in [halting_probability, remainders, n_updates]:
              x.set_shape(state_shape_static[state_slice])
          new_state.set_shape(state_shape_static)
          step += 1
          return (
              transformed_state,
              step,
              halting_probability,
              remainders,
              n_updates,
              new_state,
          )
  
      # While loop stops when this predicate is FALSE.
      # Ie all (probability < 1-eps AND counter < N) are false.
      def should_continue(
          u0, u1, halting_probability, u2, n_updates, u3
      ):
          del u0, u1, u2, u3
          return tf.reduce_any(
              tf.logical_and(
                  tf.less(halting_probability, threshold),
                  tf.less(n_updates, act_max_steps),
              )
          )
  
      # Do while loop iterations until predicate above is false.
      (_, _, _, remainder, n_updates, new_state) = tf.while_loop(
          should_continue,
          ut_function,
          (
              state,
              step,
              halting_probability,
              remainders,
              n_updates,
              previous_state,
          ),
          maximum_iterations=act_max_steps + 1,
      )
  
      ponder_times = n_updates
      remainders = remainder
  
      contrib.summary().scalar(
          "ponder_times", tf.reduce_mean(ponder_times)
      )
  
      return new_state, (ponder_times, remainders)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import functools
  from six.moves import range
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models import transformer
  from tensor2tensor.utils import contrib
  from tensor2tensor.utils import expert_utils
  import tensorflow.compat.v1
label_desc: Use image node_id 15 for calling a global function with example usage: universal_transformer_act(x, hparams, ffn_unit, attention_unit) and returns: new_state,  ponder_times,  remainders, transformed_state,  step,  halting_probability,  remainders,  n_updates,  new_state, tf