input: Write a function to generate a neural network model using TensorFlow and Python for classifying hand-written digits from the MNIST dataset. The function should train the model, evaluate its accuracy, and return the trained model. The model should have two hidden layers with 256 neurons each, and it should use the Adam optimizer with a learning rate of 0.001. If the dataset is not applicable, return None.
  assert generate_model()==None
  assert type(generate_model())==tf.Session
output: def generate_model():
      import tensorflow as tf
      from tensorflow.keras import layers, models
      from tensorflow.keras.datasets import mnist
      from tensorflow.keras.utils import to_categorical
  
      (x_train, y_train), (x_test, y_test) = mnist.load_data()
      if x_train is None or x_test is None:
          return None
  
      x_train = x_train.reshape((x_train.shape[0], 28*28)).astype('float32') / 255
      x_test = x_test.reshape((x_test.shape[0], 28*28)).astype('float32') / 255
      y_train = to_categorical(y_train)
      y_test = to_categorical(y_test)
  
      model = models.Sequential()
      model.add(layers.Dense(256, activation='relu', input_shape=(28*28,)))
      model.add(layers.Dense(256, activation='relu'))
      model.add(layers.Dense(10, activation='softmax'))
  
      model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
                    loss='categorical_crossentropy',
                    metrics=['accuracy'])
  
      model.fit(x_train, y_train, epochs=5, batch_size=128, validation_split=0.2)
      test_loss, test_acc = model.evaluate(x_test, y_test)
  
      return model