input: Write a python function to build and train a neural network model for MNIST dataset, and preprocess the data.
  assert model.summary() is not None
  assert X_train.shape == (60000, 784)
  assert X_test.shape == (10000, 784)
  assert y_train.shape == (60000, 10)
  assert y_test.shape == (10000, 10)
output: def build_and_train_mnist_model():
      import tensorflow as tf
      from tensorflow.keras.datasets import mnist
      from tensorflow.keras.models import Sequential
      from tensorflow.keras.layers import Dense, Flatten
      from tensorflow.keras.utils import to_categorical
  
      # Load and preprocess the data
      (X_train, y_train), (X_test, y_test) = mnist.load_data()
      X_train = X_train.reshape(60000, 784).astype('float32') / 255
      X_test = X_test.reshape(10000, 784).astype('float32') / 255
      y_train = to_categorical(y_train, 10)
      y_test = to_categorical(y_test, 10)
      
      # Build the model
      model = Sequential([
          Flatten(input_shape=(784,)),
          Dense(512, activation='relu'),
          Dense(256, activation='relu'),
          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_split=0.2)
      
      return model, X_train, X_test, y_train, y_test
  
  model, X_train, X_test, y_train, y_test = build_and_train_mnist_model()