def _feedforward_layer(self, x, layer_type, losses=None):
    """Feed-forward layer.

    Args:
      x: a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim]
      layer_type: a string
      losses: a list to be appended-to
    Returns:
      a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim]
    Raises:
      ValueError: if hparams make no sense
    """
    hparams = self._hparams
    mode = getattr(hparams, "mode", tf_estimator.ModeKeys.TRAIN)
    is_training = mode == tf_estimator.ModeKeys.TRAIN
    if layer_type == "drd":
      return mtf.layers.dense_relu_dense(
          x, self.feedforward_dim, is_training, dropout=hparams.relu_dropout,
          dropout_broadcast_dims=[self.length_dim],
          master_dtype=self.master_dtype,
          slice_dtype=self.slice_dtype)
    elif layer_type == "none":
      return x
    elif layer_type == "moe":
      output, loss = moe.transformer_moe_layer_v1(
          x,
          self.model_dim,
          hparams,
          hparams.mode == tf_estimator.ModeKeys.TRAIN,
          master_dtype=self.master_dtype,
          slice_dtype=self.slice_dtype)
      if losses is not None:
        losses.append(loss)
      return output
    elif layer_type == "hmoe":
      output, loss = moe.transformer_moe_layer_v2(
          x,
          self.model_dim,
          hparams,
          hparams.mode == tf_estimator.ModeKeys.TRAIN,
          master_dtype=self.master_dtype,
          slice_dtype=self.slice_dtype)
      if losses is not None:
        losses.append(loss)
      return output
    else:
      raise ValueError("layer_type not recognized %s" % layer_type)
Feed-forward layer.

    Args:
      x: a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim]
      layer_type: a string
      losses: a list to be appended-to
    Returns:
      a mtf.Tensor with shape [<batch_dims>, length_dim, model_dim]
    Raises:
      ValueError: if hparams make no sense
    
tensor2tensor/models/mtf_transformer.py