def split_coupling(
    x, x_mask, split_dim, identity_first, decoder_self_attention_bias):
  """Split function used in coupling flows."""
  n_channels = common_layers.shape_list(x)[-1]
  if split_dim == "c":
    n_transform = n_identity = n_channels // 2
    x_id = x[..., :n_identity] if identity_first else x[..., n_transform:]
    x_tr = x[..., n_identity:] if identity_first else x[..., :n_transform]
    bias, mask = decoder_self_attention_bias, x_mask

  elif split_dim == "a":
    n_transform = n_identity = n_channels // 2
    x_id = x[..., 0::2] if identity_first else x[..., 1::2]
    x_tr = x[..., 1::2] if identity_first else x[..., 0::2]
    bias, mask = decoder_self_attention_bias, x_mask

  elif split_dim == "t":
    n_transform = n_identity = n_channels
    x_id = x[:, 0::2, :] if identity_first else x[:, 1::2, :]
    x_tr = x[:, 1::2, :] if identity_first else x[:, 0::2, :]
    bias, mask = decoder_self_attention_bias[..., 0::2], x_mask[..., 0::2]

  return x_id, x_tr, n_identity, n_transform, bias, mask
Split function used in coupling flows.
tensor2tensor/layers/transformer_glow_layers_ops.py