Loss inspired by the sliced WGAN paper: https://arxiv.org/abs/1804.01947.

  Puts input1 and input2 through the provided discriminator to get logits.
  Then, computes num_vecs random projections of the logits, sorts them on
  the batch dimension and returns the L2 loss between the sorted vectors.
  See the above-mentioned paper for the reasoning behind it.

  Args:
    input1: first discriminator inputs.
    input2: second discriminator inputs.
    discriminator: inputs -> logits function.
    num_vecs: how many random vectors to use for projections.
    do_random_vecs: whether to use random vectors or just tanh of the logits.
    do_tanh: if true (default) we'll also just use tanh of the logits.
    return_logits: Whether or not to return the logits.

  Returns:
    The generator loss, i.e., the sliced approximation of the distance between
    the projected distributions (warning: discriminator should maximize it).

def sliced_gan_loss(input1,
                    input2,
                    discriminator,
                    num_vecs,
                    do_random_vecs=True,
                    do_tanh=True,
                    return_logits=False):
  """Loss inspired by the sliced WGAN paper: https://arxiv.org/abs/1804.01947.

  Puts input1 and input2 through the provided discriminator to get logits.
  Then, computes num_vecs random projections of the logits, sorts them on
  the batch dimension and returns the L2 loss between the sorted vectors.
  See the above-mentioned paper for the reasoning behind it.

  Args:
    input1: first discriminator inputs.
    input2: second discriminator inputs.
    discriminator: inputs -> logits function.
    num_vecs: how many random vectors to use for projections.
    do_random_vecs: whether to use random vectors or just tanh of the logits.
    do_tanh: if true (default) we'll also just use tanh of the logits.
    return_logits: Whether or not to return the logits.

  Returns:
    The generator loss, i.e., the sliced approximation of the distance between
    the projected distributions (warning: discriminator should maximize it).
  """
  with tf.variable_scope("sliced_gan"):
    with tf.variable_scope("discriminator"):
      logits1 = discriminator(input1)
    with tf.variable_scope("discriminator", reuse=True):
      logits2 = discriminator(input2)

    if do_random_vecs:
      random_vecs = tf.nn.l2_normalize(
          tf.random_uniform([shape_list(logits1)[-1], num_vecs]), axis=0)

    def get_sorted_projections(x):
      """Make projections of x and sort them on the batch dimension."""
      x = tf.reshape(x, [-1, shape_list(x)[-1]])
      batch_size = shape_list(x)[0]
      if do_random_vecs and do_tanh:
        n = tf.nn.l2_normalize(x, axis=1)
        proj = tf.concat([tf.matmul(n, random_vecs), tf.tanh(n)], axis=1)
      elif do_random_vecs:
        n = tf.nn.l2_normalize(x, axis=1)
        proj = tf.matmul(n, random_vecs)
      else:
        proj = tf.tanh(x)
      proj = tf.transpose(proj, [1, 0])  # [num_vecs, batch] after this.

      if is_xla_compiled():
        proj_dtype = proj.dtype
        proj = tf.cast(proj, tf.bfloat16)

        # Currently TPU only supports 1-D top_k calls.
        map_fn = lambda x: tf.nn.top_k(x, k=batch_size, sorted=True)[0]
        values = tf.map_fn(map_fn, proj)

        values = tf.cast(values, proj_dtype)
      else:
        values, _ = tf.nn.top_k(proj, k=batch_size, sorted=True)

      return values

    proj1 = get_sorted_projections(logits1)
    proj2 = get_sorted_projections(logits2)
    dist = tf.reduce_mean(tf.squared_difference(proj1, proj2))
    if return_logits:
      return dist, logits1, logits2
    return dist