code: def diagonal_normal(outputs, name="normal"):
      """Split outputs into mu and log_sigma and return z."""
      with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
          loc, log_scale = tf.split(outputs, 2, axis=-1)
          scale = tf.exp(log_scale)
          dist = tfp.distributions.Normal(
              loc=loc,
              scale=scale + tf.keras.backend.epsilon(),
              allow_nan_stats=False,
          )
          return dist
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import math
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models.transformer import transformer_decoder_layer
  import tensorflow.compat.v1
  import tensorflow_probability
label_desc: Use image node_id 13 for calling a global function with example usage: diagonal_normal(outputs, name) and returns: dist