def MLP(
    layer_widths=(128, 64),
    activation_fn=tl.Relu,
    out_activation=False,
    flatten=True,
    mode='train'):
  """A "multilayer perceptron" (MLP) network.

  This is a classic fully connected feedforward network, with one or more
  layers and a (nonlinear) activation function between each layer. For
  historical reasons, such networks are often called multilayer perceptrons;
  but they are more accurately described as multilayer networks, where
  each layer + activation function is a perceptron-like unit (see, e.g.,
  [https://en.wikipedia.org/wiki/Multilayer_perceptron#Terminology]).

  Args:
    layer_widths: Tuple of ints telling the number of layers and the width of
        each layer. For example, setting `layer_widths=(128, 64, 32)` would
        yield 3 layers with successive widths of 128, 64, and 32.
    activation_fn: Type of activation function between pairs of fully connected
        layers; must be an activation-type subclass of `Layer`.
    out_activation: If True, include a copy of the activation function as the
        last layer in the network.
    flatten: If True, insert a layer at the head of the network to flatten the
        input tensor into a matrix of shape (batch_size. -1).
    mode: Ignored.

  Returns:
    An assembled MLP network with the specified layers. This network can either
    be initialized and trained as a full model, or can be used as a building
    block in a larger network.
  """
  del mode

  layers = []
  for width in layer_widths:
    layers.append(tl.Dense(width))
    layers.append(activation_fn())

  if not out_activation:
    # Don't need the last activation.
    layers.pop()

  return tl.Serial(
      [tl.Flatten()] if flatten else [],
      layers,
  )
A "multilayer perceptron" (MLP) network.

  This is a classic fully connected feedforward network, with one or more
  layers and a (nonlinear) activation function between each layer. For
  historical reasons, such networks are often called multilayer perceptrons;
  but they are more accurately described as multilayer networks, where
  each layer + activation function is a perceptron-like unit (see, e.g.,
  [https://en.wikipedia.org/wiki/Multilayer_perceptron#Terminology]).

  Args:
    layer_widths: Tuple of ints telling the number of layers and the width of
        each layer. For example, setting `layer_widths=(128, 64, 32)` would
        yield 3 layers with successive widths of 128, 64, and 32.
    activation_fn: Type of activation function between pairs of fully connected
        layers; must be an activation-type subclass of `Layer`.
    out_activation: If True, include a copy of the activation function as the
        last layer in the network.
    flatten: If True, insert a layer at the head of the network to flatten the
        input tensor into a matrix of shape (batch_size. -1).
    mode: Ignored.

  Returns:
    An assembled MLP network with the specified layers. This network can either
    be initialized and trained as a full model, or can be used as a building
    block in a larger network.
  
trax/models/mlp.py