code: def diagonal_neural_gpu(inputs, hparams, name=None):
      """Improved Neural GPU as in https://arxiv.org/abs/1702.08727."""
      with tf.variable_scope(name, "diagonal_neural_gpu"):
  
          def step(state_tup, inp):
              """Single step of the improved Neural GPU."""
              state, _ = state_tup
              x = state
              for layer in range(hparams.num_hidden_layers):
                  x, new_loss = common_layers.diagonal_conv_gru(
                      x,
                      (hparams.kernel_height, hparams.kernel_width),
                      hparams.hidden_size,
                      dropout=hparams.dropout,
                      name="dcgru_%d" % layer,
                  )
              # Padding input is zeroed-out in the modality, we check this by summing.
              padding_inp = tf.less(
                  tf.reduce_sum(tf.abs(inp), axis=[1, 2]), 0.00001
              )
              new_state = tf.where(
                  padding_inp, state, x
              )  # No-op where inp is padding.
              return new_state, new_loss
  
          final_state, losses = tf.scan(
              step,
              tf.transpose(inputs, [1, 0, 2, 3]),
              initializer=(inputs, tf.constant(0.0)),
              parallel_iterations=1,
              swap_memory=True,
          )
          return final_state[0, :, :, :, :], 2.0 * tf.reduce_mean(
              losses
          )
  
file_path: tensor2tensor/tensor2tensor/models/neural_gpu.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from six.moves import range
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_layers
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: diagonal_neural_gpu(inputs, hparams, name) and returns: final_state,  unknown, new_state,  new_loss