code: def MLP(num_hidden_layers=2,
          hidden_size=512,
          activation_fn=layers.Relu,
          num_output_classes=10,
          mode="train"):
    
    del mode
    cur_layers = [layers.Flatten()]
    for _ in range(num_hidden_layers):
      cur_layers += [layers.Dense(hidden_size), activation_fn()]
    cur_layers += [layers.Dense(num_output_classes), layers.LogSoftmax()]
    return layers.Serial(*cur_layers)
docs: Multi-layer feed-forward neural network with non-linear activations.