def split(name, x, x_mask, inverse, temp=1.0, dtype=tf.float32, z=None):
  """Splits / concatenates x into x1 and x2 across number of channels.

  x2 is modelled with a standard gaussian distribution.
  Args:
    name: variable scope.
    x: 3-D Tensor, shape=[B, L, C].
    x_mask: 2-D Tensor, shape=[B, L].
    inverse: forward or inverse pass.
    temp: only used for inverse pass. temperature for sampling.
    dtype: dtype
    z: used in inverse pass to check invertibility.

  Returns:
    x: if forward, returns the 1st half of the channel dimensions.
      if inverse, return concat[input, N(0,1)]
    z: second half of the channel dimensions. modelled as standard normal.
    log_p: log p(x2; N(0,1)), shape=[B]
  """
  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
    if not inverse:
      x1, x2 = tf.split(x, 2, axis=-1)
      log_p = gops.standard_normal_density(x2, x_mask)
      return x1, x2, log_p
    else:
      if z is None:
        x2 = tf.random.normal(
            common_layers.shape_list(x), stddev=temp, dtype=dtype)
      else:
        x2 = z
      log_p = gops.standard_normal_density(x2, x_mask)
      return tf.concat([x, x2], 2), None, log_p
Splits / concatenates x into x1 and x2 across number of channels.

  x2 is modelled with a standard gaussian distribution.
  Args:
    name: variable scope.
    x: 3-D Tensor, shape=[B, L, C].
    x_mask: 2-D Tensor, shape=[B, L].
    inverse: forward or inverse pass.
    temp: only used for inverse pass. temperature for sampling.
    dtype: dtype
    z: used in inverse pass to check invertibility.

  Returns:
    x: if forward, returns the 1st half of the channel dimensions.
      if inverse, return concat[input, N(0,1)]
    z: second half of the channel dimensions. modelled as standard normal.
    log_p: log p(x2; N(0,1)), shape=[B]
  
tensor2tensor/layers/transformer_glow_layers.py