class_docstr: BERT model ("Bidirectional Encoder Representations from Transformers").
class_docstr_tok: ["BERT","model","(","``","Bidirectional","Encoder","Representations","from","Transformers","''",")","."]
docstr: Mixture of experts layer.
  
  TODO(noam): clean up the mixture-of-experts code in Transformer.
  
  Args:
    x: layer input
    layout: a mtf.LayoutRules
    mesh_shape: a mtf.Shape
    input_mask: a mtf.Tensor
    is_training: a boolean
  Returns:
    a mtf.Tensor (the layer output)
code: def moe(self, x, layout, mesh_shape, input_mask, is_training):
      """Mixture of experts layer.
  
      TODO(noam): clean up the mixture-of-experts code in Transformer.
  
      Args:
        x: layer input
        layout: a mtf.LayoutRules
        mesh_shape: a mtf.Shape
        input_mask: a mtf.Tensor
        is_training: a boolean
      Returns:
        a mtf.Tensor (the layer output)
      """
      hparams = moe.HParams(
          moe_gating="top_2",
          moe_num_experts=self.config.moe_num_experts,
          moe_loss_coef=1e-3,
          moe_hidden_size=self.config.moe_intermediate_size,
          moe_group_size=2048,
          moe_capacity_factor_train=1.25,
          moe_capacity_factor_eval=8.0,
          moe_use_second_place_loss=False,
          moe_second_policy_train="random",
          moe_second_policy_eval="random",
          moe_second_threshold_train=0.2,
          moe_second_threshold_eval=0.2,
      )
      layer_output, loss = moe.transformer_moe_layer_v1(
          inputs=x,
          output_dim=self.model_dim,
          hparams=hparams,
          train=is_training,
          variable_dtype=tf.float32,
          layout=layout,
          mesh_shape=mesh_shape,
          nonpadding=(
              mtf.cast(input_mask, tf.float32) if input_mask else None
          ),
          activation=get_activation(
              self.config.feedforward_intermediate_act
          ),
      )
      self._extra_losses.append(loss)
      return layer_output
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import copy
  import json
  import math
  import re
  import mesh_tensorflow
  import mesh_tensorflow.transformer.moe
  import six
  import tensorflow.compat.v1
label_desc: Use image node_id 4 for calling the BertModel obj's underlying member method code with example usage: obj.moe(x, layout, mesh_shape, input_mask, is_training) and returns: layer_output