code: def universal_transformer_depthwise_attention(
      layer_inputs, step, hparams, ffn_unit, attention_unit
  ):
      """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.
  
      """
      _, 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
  
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 12 for calling a global function with example usage: universal_transformer_depthwise_attention(layer_inputs, step, hparams, ffn_unit, attention_unit) and returns: new_state,  inputs,  memory