Local mixture of experts that works well on TPU.

  See https://arxiv.org/abs/1701.06538

  There are num_experts expert networks, each containing a relu-activated
  hidden layer of size hidden_size, followed by an output projection.

  The number of parameters is thus:
    num_experts * (input_size * hidden_size + hidden_size * output_size)

  The input is 3d: [batch, length, depth], consisting of the representations
  of all positions in a batch of sequences.

  Each position of each sequence is sent to 0-2 experts.  The expert
  choices and the combination weights are determined by a learned gating
  function.

  This function returns a small auxiliary loss that should be added to the
  training loss of the model.  This loss helps to balance expert usage.
  Without the loss, it is very likely that a few experts will be trained and
  the rest will starve.

  Several hacks are necessary to get around current TPU limitations:

  - To ensure static shapes, we enforce (by truncation/padding)
    that each sequence send the same number of elements to each expert.

    It would make more sense to enforce this equality over the entire batch,
    as opposed to on individual sequences.  This would allow more freedom
    for individual sequences to be unbalanced.  Unfortunately, that would
    slow down our hacked-up gather-by-matmul implementation.

    TODO(noam): There is no real reason for a single sequence to be the unit
      of equal allocation.  Reshaping the inputs would allow us to pick a
      different unit of equal allocation.

  TODO(noam): Factor this code better.  We want to be able to substitute
  different code for the experts themselves.  We also want to integrate this
  gating/dispatching logic into multi-device mixtures-of-experts.

  Args:
    inputs: a Tensor with shape [batch, length, depth]
    hidden_size: an integer
    output_size: an integer
    num_experts: an integer
    loss_coef: a float scalar
    overhead: multiplicative factor of how much spare capacity to assign

  Returns:
    outputs: a Tensor with shape [batch, length, output_size]
    loss: a scalar

def local_moe_tpu(inputs,
                  hidden_size,
                  output_size,
                  num_experts,
                  loss_coef=1e-3,
                  overhead=1.0):
  """Local mixture of experts that works well on TPU.

  See https://arxiv.org/abs/1701.06538

  There are num_experts expert networks, each containing a relu-activated
  hidden layer of size hidden_size, followed by an output projection.

  The number of parameters is thus:
    num_experts * (input_size * hidden_size + hidden_size * output_size)

  The input is 3d: [batch, length, depth], consisting of the representations
  of all positions in a batch of sequences.

  Each position of each sequence is sent to 0-2 experts.  The expert
  choices and the combination weights are determined by a learned gating
  function.

  This function returns a small auxiliary loss that should be added to the
  training loss of the model.  This loss helps to balance expert usage.
  Without the loss, it is very likely that a few experts will be trained and
  the rest will starve.

  Several hacks are necessary to get around current TPU limitations:

  - To ensure static shapes, we enforce (by truncation/padding)
    that each sequence send the same number of elements to each expert.

    It would make more sense to enforce this equality over the entire batch,
    as opposed to on individual sequences.  This would allow more freedom
    for individual sequences to be unbalanced.  Unfortunately, that would
    slow down our hacked-up gather-by-matmul implementation.

    TODO(noam): There is no real reason for a single sequence to be the unit
      of equal allocation.  Reshaping the inputs would allow us to pick a
      different unit of equal allocation.

  TODO(noam): Factor this code better.  We want to be able to substitute
  different code for the experts themselves.  We also want to integrate this
  gating/dispatching logic into multi-device mixtures-of-experts.

  Args:
    inputs: a Tensor with shape [batch, length, depth]
    hidden_size: an integer
    output_size: an integer
    num_experts: an integer
    loss_coef: a float scalar
    overhead: multiplicative factor of how much spare capacity to assign

  Returns:
    outputs: a Tensor with shape [batch, length, output_size]
    loss: a scalar
  """
  batch, length, input_size = common_layers.shape_list(inputs)[:]
  # Each sequence sends expert_capacity positions to each expert.
  if isinstance(length, int):
    expert_capacity = min(
        length, int((length * 2 * overhead) / num_experts))
  else:
    expert_capacity = tf.minimum(
        length, tf.to_int32(
            tf.to_float(length) * 2 * overhead / num_experts))
  expert_capacity_f = tf.to_float(expert_capacity)

  # This is the learned gating function.
  gates = tf.nn.softmax(
      tf.to_float(common_layers.dense(inputs, num_experts, name="logits")))

  # Find the top expert for each position.
  gate_1, index_1 = common_layers.top_1_tpu(gates)
  # [batch, length, num_experts]
  mask_1 = tf.one_hot(index_1, num_experts)
  # [batch, length, num_experts]
  # This is the position within the expert's mini-batch for this sequence
  position_in_expert_1 = common_layers.cumsum(
      mask_1, axis=1, exclusive=True) * mask_1
  # Remove the elements that don't fit.
  mask_1 *= tf.to_float(tf.less(position_in_expert_1, expert_capacity_f))
  # [batch, 1, num_experts]
  # How many examples in this sequence go to this expert
  mask_1_count = tf.reduce_sum(mask_1, axis=1, keepdims=True)
  # [batch, length] - mostly ones, but zeros where something didn't fit
  mask_1_flat = tf.reduce_sum(mask_1, axis=2)
  position_in_expert_1 = tf.reduce_sum(position_in_expert_1, axis=2)
  # Weight assigned to first expert.
  gate_1 *= mask_1_flat

  # Pick a second-place expert for each position.
  # We first mask out the experts that we expect to be over-capacity
  space_remaining = expert_capacity_f - mask_1_count
  use_rate = (mask_1_count + 1.0) / tf.to_float(length)
  # At what point in the sequence do we expect the expert to be full.
  expected_exhaustion_pos = space_remaining / use_rate
  # A Tensor with shape [batch, length, num_experts] representing a boolean
  #   - whether we expect that the expert will already be full.
  expected_exhausted = tf.to_float(tf.greater(
      tf.reshape(tf.to_float(tf.range(length)), [1, length, 1]),
      expected_exhaustion_pos))
  masked_gates = gates - mask_1 - expected_exhausted
  # This section is similar to the section above.
  gate_2, index_2 = common_layers.top_1_tpu(masked_gates)
  # [batch, length, num_experts]
  mask_2 = tf.one_hot(index_2, num_experts)
  position_in_expert_2 = (
      common_layers.cumsum(mask_2, axis=1, exclusive=True) + mask_1_count)
  position_in_expert_2 *= mask_2
  mask_2 *= tf.to_float(tf.less(position_in_expert_2, expert_capacity_f))
  mask_2_count = tf.reduce_sum(mask_2, axis=1, keepdims=True)
  mask_2_flat = tf.reduce_sum(mask_2, axis=2)
  position_in_expert_2 = tf.reduce_sum(position_in_expert_2, axis=2)
  gate_2 *= mask_2_flat

  # What fraction didn't fit - show summaries
  miss_rate_1 = 1.0 - tf.reduce_sum(mask_1_count) / tf.to_float(batch * length)
  miss_rate_2 = 1.0 - tf.reduce_sum(mask_2_count) / tf.to_float(batch * length)
  tf.summary.scalar("miss_rate_1", miss_rate_1)
  tf.summary.scalar("miss_rate_2", miss_rate_2)

  # renormalize the two gate values to add up to 1
  denom = gate_1 + gate_2 + 1e-9
  gate_1 /= denom
  gate_2 /= denom

  # inputs: [batch, length, input_size]
  # forward_assignment: [batch, length, num_experts * expert_capacity]
  # expert_inputs: [batch, num_experts * expert_capacity, input_size]

  segment_ids_forward_1 = (
      (index_1 * expert_capacity) +
      tf.to_int32(position_in_expert_1) +
      tf.to_int32(1.0 - mask_1_flat) * (num_experts * expert_capacity))

  segment_ids_forward_2 = (
      (index_2 * expert_capacity) +
      tf.to_int32(position_in_expert_2) +
      tf.to_int32(1.0 - mask_2_flat) * (num_experts * expert_capacity))

  # Gather and scatter are painfully slow on TPU.
  # We will use one_hot and matmul instead.

  # [batch, length, num_experts * expert_capacity]
  one_hot_1 = tf.one_hot(
      segment_ids_forward_1, num_experts * expert_capacity, dtype=inputs.dtype)
  one_hot_2 = tf.one_hot(
      segment_ids_forward_2, num_experts * expert_capacity, dtype=inputs.dtype)

  forward_assignment = (one_hot_1 + one_hot_2)

  # [batch, num_experts * expert_capacity, input_size]
  expert_inputs = tf.matmul(forward_assignment, inputs, transpose_a=True)

  # [batch, num_experts, expert_capacity, input_size]
  expert_inputs = tf.reshape(
      expert_inputs, [batch, num_experts, expert_capacity, input_size])
  # [num_experts, batch, expert_capacity, input_size]
  expert_inputs = tf.transpose(expert_inputs, [1, 0, 2, 3])

  # [num_experts, batch * expert_capacity, input_size]
  expert_inputs = tf.reshape(
      expert_inputs, [num_experts, batch * expert_capacity, input_size])

  # Now feed the expert inputs through the experts.
  h = common_layers.batch_dense(
      expert_inputs, hidden_size, activation=tf.nn.relu, name="x0")
  expert_output = common_layers.batch_dense(h, output_size, name="x1")
  expert_output = tf.reshape(
      expert_output, [num_experts, batch, expert_capacity, output_size])

  # [batch, num_experts, expert_capacity, output_size]
  expert_output = tf.transpose(expert_output, [1, 0, 2, 3])
  expert_output = tf.reshape(
      expert_output, [batch, num_experts * expert_capacity, output_size])

  # Again, use matmul instead of unsorted_segment_sum.  This time, we need
  # to multiply by the combination weights gate_1 and gate_2.

  # expert_output: [batch, num_experts * expert_capacity, output_size]
  # backward_assigmnent: [batch, length, num_experts * expert_capacity]
  # output: [batch, length, output_size]
  backward_assigmnent = (
      one_hot_1 * tf.cast(tf.expand_dims(gate_1, 2), inputs.dtype) +
      one_hot_2 * tf.cast(tf.expand_dims(gate_2, 2), inputs.dtype))
  output = tf.matmul(backward_assigmnent, expert_output)

  # Compute a loss equal to the coefficient ov variation of the
  # total gate value per expert per sequence.
  # This loss causes the experts to be used about equally used per sequence.
  importance = tf.reduce_sum(gates * (mask_1 + mask_2), 1)
  loss = loss_coef * cv_squared(importance)
  return output, loss