input: Write a function to construct a neural network model in Keras that classifies the MNIST data set correctly. The function should return the compiled model.
  
  import keras
  
  def build_mnist_model():
      # Input layer
      input_layer = keras.layers.Input(shape=(784,))
      
      # Hidden layer
      hidden_layer = keras.layers.Dense(128, activation='relu')(input_layer)
      
      # Output layer
      output_layer = keras.layers.Dense(10, activation='softmax')(hidden_layer)
      
      # Model
      model = keras.Model(inputs=input_layer, outputs=output_layer)
      
      # Compile model
      model.compile(optimizer='adam', 
                    loss='sparse_categorical_crossentropy',
                    metrics=['accuracy'])
      
      return model
  
  assert type(build_mnist_model()) == keras.engine.functional.Functional
  model = build_mnist_model()
  assert model.input_shape == (None, 784)
  assert model.output_shape == (None, 10)
  assert model.optimizer.get_config()['name'] == 'Adam'
  assert model.loss == 'sparse_categorical_crossentropy'
output: def build_mnist_model():
      input_layer = keras.layers.Input(shape=(784,))
      hidden_layer = keras.layers.Dense(128, activation='relu')(input_layer)
      output_layer = keras.layers.Dense(10, activation='softmax')(hidden_layer)
      model = keras.Model(inputs=input_layer, outputs=output_layer)
      model.compile(optimizer='adam', 
                    loss='sparse_categorical_crossentropy',
                    metrics=['accuracy'])
      return model