input: Write a function to create and train a neural network model for the MNIST dataset using Keras.
  assert isinstance(create_model(), keras.models.Sequential)
  assert create_model().input_shape == (None, 784)
  assert create_model().output_shape == (None, 10)
output: def create_model():
      from keras.models import Sequential
      from keras.layers import Dense, Flatten
      from keras.datasets import mnist
      from keras.utils import np_utils
  
      # Load MNIST dataset
      (X_train, y_train), (X_test, y_test) = mnist.load_data()
  
      # Normalize the input data
      X_train = X_train.reshape(-1, 784).astype('float32') / 255
      X_test = X_test.reshape(-1, 784).astype('float32') / 255
  
      # Convert labels to one-hot encoding
      y_train = np_utils.to_categorical(y_train, 10)
      y_test = np_utils.to_categorical(y_test, 10)
  
      # Create the model
      model = Sequential()
      model.add(Dense(512, activation='relu', input_shape=(784,)))
      model.add(Dense(256, activation='relu'))
      model.add(Dense(10, activation='softmax'))
  
      # Compile the model
      model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
  
      # Train the model
      model.fit(X_train, y_train, epochs=10, batch_size=128, validation_data=(X_test, y_test))
  
      return model