input: Write a function that generates a Machine Learning model in Python to recognize handwritten digits from 0 to 9. The function should include the necessary data loading, model definition, compilation, and fitting steps. Test the function with the MNIST dataset.
  assert (predict_digit(np.array([[[0,0,0,...,0,0,0],[0,0,0,...,0,0,0],...,[0,0,0,...,0,0,0]]]))) in range(10)
  assert isinstance(predict_digit(np.array([[[0,0,0,...,0,0,0],[0,0,0,...,0,0,0],...,[0,0,0,...,0,0,0]]]))[0], int)
  assert isinstance(predict_digit(np.array([[[0,0,0,...,0,0,0],[0,0,0,...,0,0,0],...,[0,0,0,...,0,0,0]]]))[0], int)
  assert predict_digit(np.array([[[0,0,0,...,0,0,0],[0,0,0,...,0,0,0],...,[0,0,0,...,0,0,0]]])).shape == (1,)
output: 
  def predict_digit(image):
      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
      import numpy as np
  
      (x_train, y_train), (x_test, y_test) = mnist.load_data()
      x_train = x_train / 255.0
      x_test = x_test / 255.0
      y_train = to_categorical(y_train, 10)
      y_test = to_categorical(y_test, 10)
  
      model = Sequential([
          Flatten(input_shape=(28, 28)),
          Dense(128, activation='relu'),
          Dense(10, activation='softmax')
      ])
  
      model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
      model.fit(x_train, y_train, epochs=5, batch_size=32, validation_data=(x_test, y_test))
  
      image = image / 255.0
      prediction = model.predict(np.array([image]))
      return np.argmax(prediction, axis=1)
  