class_docstr: Mixture of Experts Layer.
class_docstr_tok: ["Mixture","of","Experts","Layer","."]
docstr: Call the layer.
docstr_tok: ["Call","the","layer","."]
code: def call(self, context, x, losses=None):
      """Call the layer."""
      if context.model.ensemble_dim:
          raise NotImplementedError(
              "MoE not yet implemented with ensembles"
          )
  
      has_length_dim = context.length_dim in x.shape.dims
      if not has_length_dim:
          x_shape = x.shape
          shape_with_length = mtf.Shape(
              x_shape.dims[:-1]
              + [mtf.Dimension("length", 1)]
              + x_shape.dims[-1:]
          )
          x = mtf.reshape(x, shape_with_length)
  
      # Extract the MoE output dimension
      if self._hparams.moe_output_dim is not None:
          output_dim = self._hparams.moe_output_dim
      else:
          output_dim = context.model.model_dim
      y, loss = transformer_moe_layer_v1(
          x,
          output_dim,
          self._hparams,
          context.train,
          context.variable_dtype,
          layout=context.model.layout,
          mesh_shape=context.model.mesh_shape,
          nonpadding=context.nonpadding,
          activation=self._activation,
          num_microbatches=context.num_microbatches,
          token_embeddings=context.input_embeddings,
      )
      if context.losses is not None:
          context.losses.append(loss)
      if not has_length_dim:
          if self._hparams.moe_use_experts_attention:
              y_reshape = [mtf.reshape(y_out, x_shape) for y_out in y]
              y = y_reshape
          else:
              y = mtf.reshape(y, x_shape)
      return y
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling the MoE1D obj's underlying member method code with example usage: obj.call(context, x, losses) and returns: y