<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def MLP(num_hidden_layers=2, hidden_size=512, activation_fn=layers.Relu, num_output_classes=10, mode="train"):
 """Multi-layer feed-forward neural network with non-linear activations."""

  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)