code: def universal_transformer_depthwise_attention(layer_inputs,
                                                step, hparams,
                                                ffn_unit,
                                                attention_unit):
    
    _, inputs, memory = layer_inputs
    all_states = memory
  
    # add depth signal
    if hparams.depth_embedding:
      all_states = add_depth_embedding(all_states)
  
    # get the states up to the current step (non-zero part of the memory)
    states_so_far = all_states[:step, :, :, :]
  
    states_so_far_weights = tf.nn.softmax(
        common_layers.dense(
            states_so_far, (hparams.hidden_size if hparams.dwa_elements else 1),
            activation=None,
            use_bias=True),
        axis=-1)
  
    # prepare the state tensor that will be transformed
    state_to_be_transformed = tf.reduce_sum(
        (states_so_far * states_so_far_weights), axis=0)
  
    new_state = step_preprocess(state_to_be_transformed, step, hparams)
  
    for i in range(hparams.num_inrecurrence_layers):
      with tf.variable_scope("rec_layer_%d" % i):
        new_state = ffn_unit(attention_unit(new_state))
  
    # add the new state to the memory
    memory = fill_memory_slot(memory, new_state, step + 1)
  
    return new_state, inputs, memory
docs: universal_transformer with depth-wise attention.
  
    It uses an attention mechanism-flipped vertically-
    over all the states from previous steps to generate the new_state.
  
    Args:
      layer_inputs:
        - state: state
        - memory: contains states from all the previous steps.
      step: indicating number of steps take so far
      hparams: model hyper-parameters.
      ffn_unit: feed-forward unit
      attention_unit: multi-head attention unit
  
  
    Returns:
      layer_output:
          new_state: new state
          memory: contains states from all the previous steps.