def join_coupling(z_id, z_tr, split_dim, identity_first):
  """Reverse split function used in coupling flows."""
  assert z_id.shape.rank == 3 and z_tr.shape.rank == 3
  result = [z_id, z_tr] if identity_first else [z_tr, z_id]
  if split_dim == "c":
    result = tf.concat(result, axis=2)  # concat in the channel dimension
  elif split_dim == "a":
    result = tf.stack(result, axis=3)  # stack in the channel dimension
  elif split_dim == "t":
    result = tf.stack(result, axis=2)  # stack in the time dimension
  return result
Reverse split function used in coupling flows.
tensor2tensor/layers/transformer_glow_layers_ops.py