def standard_normal(x, name="normal"):
  """Return standard normal distribution with same shape as x."""
  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
    dist = tfp.distributions.Normal(
        loc=tf.zeros_like(x),
        scale=tf.ones_like(x),
        allow_nan_stats=False)
    return dist
Return standard normal distribution with same shape as x.
tensor2tensor/layers/transformer_glow_layers_ops.py