input: Write a function to construct a neural network model using Keras to classify the MNIST dataset correctly. The function should return the compiled model.
  
  def create_mnist_model():
      import keras
  
      # 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 isinstance(create_mnist_model(), keras.Model)
  assert create_mnist_model().optimizer.get_config()['name'] == 'Adam'
  assert create_mnist_model().loss == 'sparse_categorical_crossentropy'
  assert create_mnist_model().metrics[0] == 'accuracy'
output: def create_mnist_model():
      import keras
  
      # 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