code: def transformer_moe_layer_v1(
      inputs,
      output_dim,
      hparams,
      train,
      master_dtype=tf.bfloat16,
      slice_dtype=tf.float32,
  ):
      """Local mixture of experts that works well on TPU.
  
      Adapted from the paper https://arxiv.org/abs/1701.06538
  
      Note: until the algorithm and inferface solidify, we pass in a hyperparameters
      dictionary in order not to complicate the interface in mtf_transformer.py .
      Once this code moves out of "research", we should pass the hyperparameters
      separately.
  
      Hyperparameters used:
        hparams.moe_num_experts: number of experts
        hparams.moe_hidden_size: size of hidden layer in each expert
        hparams.moe_group_size: size of each "group" for gating purposes
        hparams.moe_capacity_factor_train: a float
        hparams.moe_capacity_factor_eval: a float
        hparams.moe_gating: a string
        + all hyperparmeters used by _top_2_gating()
  
      The number of parameters in the gating network is:
        (input_dim.size * hparams.num_experts) +
  
      The number of parameters in the experts themselves is:
        (hparams.num_experts
         * (input_dim.size + output_dim.size)
         * hparams.moe_hidden_size)
  
      The input is n-dimensional: [<batch_and_length_dims>, input_dim], 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,
        but due to our hacked-up gather-by-matmul implementation, we need to divide
        the batch into "groups".  For each group, the same number of elements
        are sent to each expert.
  
      TODO(noam): Factor this code better.  We want to be able to substitute
      different code for the experts themselves.
  
      Args:
        inputs: a mtf.Tensor with shape [<batch_dims...>, length_dim, input_dim]
        output_dim: a mtf.Dimension (for Transformer, this is input_dim)
        hparams: model hyperparameters
        train: a boolean
        master_dtype: a tf.dtype
        slice_dtype: a tf.dtype
  
      Returns:
        outputs: a Tensor with shape [<batch_dims...>, length_dim, output_dim]
        loss: a mtf scalar
  
      Raises:
        ValueError: on unrecognized hparams.moe_gating
      """
      orig_inputs = inputs
      input_dim = inputs.shape.dims[-1]
      hidden_dim = mtf.Dimension(
          "expert_hidden", hparams.moe_hidden_size
      )
      experts_dim = mtf.Dimension("experts", hparams.moe_num_experts)
      group_size_dim = mtf.Dimension("group", hparams.moe_group_size)
      batch_dim = mtf.Dimension(
          orig_inputs.shape[0].name,
          orig_inputs.shape.size
          // (group_size_dim.size * input_dim.size),
      )
      inputs = mtf.reshape(
          inputs, [batch_dim, group_size_dim, input_dim]
      )
  
      # Each sequence sends expert_capacity positions to each expert.
      capacity_factor = (
          hparams.moe_capacity_factor_train
          if train
          else hparams.moe_capacity_factor_eval
      )
      expert_capacity = min(
          group_size_dim.size,
          int(
              (group_size_dim.size * capacity_factor) / experts_dim.size
          ),
      )
      expert_capacity_dim = mtf.Dimension(
          "expert_capacity", expert_capacity
      )
  
      experts_dim_unsplit = mtf.Dimension(
          "expert_unsplit", experts_dim.size
      )
      batch_dim_unsplit = mtf.Dimension("batch_unsplit", batch_dim.size)
  
      if hparams.moe_gating == "top_2":
          dispatch_tensor, combine_tensor, loss = _top_2_gating(
              inputs=inputs,
              outer_expert_dims=None,
              experts_dim=experts_dim_unsplit,
              expert_capacity_dim=expert_capacity_dim,
              hparams=hparams,
              train=train,
          )
      else:
          raise ValueError(
              "unknown hparams.moe_gating=%s" % hparams.moe_gating
          )
  
      # put num_experts dimension first to make split easier in alltoall
      expert_inputs = mtf.einsum(
          [inputs, dispatch_tensor],
          mtf.Shape(
              [
                  experts_dim_unsplit,
                  batch_dim,
                  expert_capacity_dim,
                  input_dim,
              ]
          ),
      )
  
      expert_inputs = mtf.reshape(
          expert_inputs,
          mtf.Shape(
              [
                  experts_dim,
                  batch_dim_unsplit,
                  expert_capacity_dim,
                  input_dim,
              ]
          ),
      )
  
      # Now feed the expert inputs through the experts.
      h = mtf.layers.dense(
          expert_inputs,
          hidden_dim,
          expert_dims=[experts_dim],
          activation=mtf.relu,
          use_bias=False,
          master_dtype=master_dtype,
          slice_dtype=slice_dtype,
          name="x0",
      )
      expert_output = mtf.layers.dense(
          h,
          output_dim,
          expert_dims=[experts_dim],
          use_bias=False,
          master_dtype=master_dtype,
          slice_dtype=slice_dtype,
          name="x1",
      )
  
      expert_output = mtf.reshape(
          expert_output,
          mtf.Shape(
              [
                  experts_dim_unsplit,
                  batch_dim,
                  expert_capacity_dim,
                  input_dim,
              ]
          ),
      )
  
      output = mtf.einsum(
          [expert_output, combine_tensor],
          mtf.Shape([batch_dim, group_size_dim, output_dim]),
      )
  
      output = mtf.reshape(
          output, orig_inputs.shape.dims[:-1] + [output_dim]
      )
  
      return output, loss * hparams.moe_loss_coef
  
file_path: tensor2tensor/tensor2tensor/models/research/moe.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  import tensorflow.compat.v1
label_desc: Use image node_id 1 for calling a global function with example usage: transformer_moe_layer_v1(inputs, output_dim, hparams, train, master_dtype, slice_dtype) and returns: output,  unknown