def _FeedForward(d_model, d_ff, dropout, activation, act_dropout,
                 use_bfloat16, mode):
  """Feed-forward block with layer normalization at start."""
  if act_dropout is None:
    act_dropout = dropout
  return [
      tl.Dense(d_ff, use_bfloat16=use_bfloat16),
      tl.Dropout(rate=act_dropout, shared_axes=[-2], mode=mode),
      activation(),
      tl.Dense(d_model, use_bfloat16=use_bfloat16),
  ]
Feed-forward block with layer normalization at start.
trax/models/research/configurable_transformer.py