class_docstr: A stack of layers for Universal Transformer.
  
  This implementation is largely adapted from t2t universal transformer
  implementation. Reference:
  third_party/py/tensor2tensor/models/research
class_docstr_tok: ["A","stack","of","layers","for","Universal","Transformer",".","This","implementation","is","largely","adapted","from","t2t","universal","transformer","implementation",".","Reference",":","third_party\/py\/tensor2tensor\/models\/research"]
docstr: Build a Universal Transformer ACT layer.
docstr_tok: ["Build","a","Universal","Transformer","ACT","layer","."]
code: def act_layer(self, context, x, mask):
      """Build a Universal Transformer ACT layer."""
      state = x
      act_max_steps = self.act_max_steps
      threshold = 1.0 - self.act_epsilon
      state_shape_static = state.shape.dims
  
      state_slice = slice(0, 3)
      if self.act_type == "global":
          state_slice = slice(0, 2)
  
      # Dynamic shape for update tensors below
      update_shape = state_shape_static[state_slice]
  
      # Halting probabilities (p_t^n in the paper)
      halting_probability = mtf.zeros(
          context.mesh, update_shape, dtype=context.activation_dtype
      )
  
      # Remainders (R(t) in the paper)
      remainders = mtf.zeros(
          context.mesh, update_shape, dtype=context.activation_dtype
      )
  
      # Number of updates performed (N(t) in the paper)
      n_updates = mtf.zeros(
          context.mesh, update_shape, dtype=context.activation_dtype
      )
  
      # Previous cell states (s_t in the paper)
      previous_state = mtf.zeros_like(state)
      step = mtf.constant(context.mesh, 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 = self.step_preprocess(context, state, step)
  
          if self.act_type == "random":
              # random as halting probability
              p = mtf.random_uniform(
                  context.mesh,
                  shape=halting_probability.shape.dims,
                  dtype=context.variable_dtype,
              )
          else:
              last_dim_name = state.shape.dimension_names[-1]
              new_dims = [mtf.Dimension(last_dim_name, 1)]
              with tf.variable_scope(
                  "sigmoid_activation_for_pondering",
                  reuse=tf.AUTO_REUSE,
              ):
                  p = mtf.layers.dense(
                      state,
                      variable_dtype=context.variable_dtype,
                      reduced_dims=[state.shape.dims[-1]],
                      new_dims=new_dims,
                      activation=mtf.sigmoid,
                      use_bias=True,
                  )
                  if self.act_type == "global":
                      # average over all positions (as a global halting prob)
                      p = mtf.reduce_mean(
                          p, reduced_dim=p.shape.dims[1]
                      )
                      p = mtf.squeeze(p)
                  else:
                      # maintain position-wise probabilities
                      new_shape = p.shape.dims[:-1]
                      p = mtf.reshape(p, new_shape)
          # Mask for inputs which have not halted yet
          still_running = mtf.cast(
              mtf.less(halting_probability, 1.0),
              context.activation_dtype,
          )
  
          # Mask of inputs which halted at this step
          new_halted = (
              mtf.cast(
                  mtf.greater(
                      halting_probability + p * still_running, threshold
                  ),
                  context.activation_dtype,
              )
              * still_running
          )
          # Mask of inputs which haven't halted, and didn't halt this step
          still_running = (
              mtf.cast(
                  mtf.less_equal(
                      halting_probability + p * still_running, threshold
                  ),
                  context.activation_dtype,
              )
              * 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
          input_tensor = p * still_running + new_halted * remainders
          update_weights = input_tensor
  
          # apply transformation on the state
          transformed_state = state
  
          for _ in range(self.num_inrecurrence_layers):
              transformed_state = self.vanilla_transformer_layer(
                  context, transformed_state, mask
              )
  
          # update running part in the weighted state and keep the rest
          new_state = (transformed_state * update_weights) + (
              previous_state * (1 - update_weights)
          )
  
          if self.act_type == "accumulated":
              # Add in the weighted state
              new_state = (
                  transformed_state * update_weights
              ) + previous_state
  
          step += 1
  
          return (
              transformed_state,
              step,
              halting_probability,
              remainders,
              n_updates,
              new_state,
          )
  
      for _ in range(act_max_steps + 1):
          (
              state,
              step,
              halting_probability,
              remainders,
              n_updates,
              previous_state,
          ) = ut_function(
              state,
              step,
              halting_probability,
              remainders,
              n_updates,
              previous_state,
          )
      ponder_times = n_updates
  
      mtf.scalar_summary("ponder_times", mtf.reduce_mean(ponder_times))
      return previous_state
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import math
  import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  import tensorflow.compat.v1
label_desc: Use image node_id 10 for calling the UTLayerStack obj's underlying member method code with example usage: obj.act_layer(context, x, mask) and returns: previous_state, transformed_state,  step,  halting_probability,  remainders,  n_updates,  new_state