code: 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
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import math
  import six
  from six.moves import range
  from six.moves import zip
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers.vq_discrete import DiscreteBottleneck
  import tensorflow.compat.v1
label_desc: Use image node_id 19 for calling a global function with example usage: local_moe_tpu(inputs, hidden_size, output_size, num_experts, loss_coef, overhead) and returns: output,  loss