transformerMLP TransformerMLP {..} train input =
  residual f (pure . forward ln) input
  where
    f x =
      dropoutForward dropout1 train
        . forward linear1
        =<< dropoutForward dropout0 train
          . relu
          . forward linear0
        =<< pure x

instance
  ( All KnownNat '[embedDim, ffnDim],
    KnownDType dtype,