code: def mnist_dataset(num_classes):
      # Load dataset
      dataset = tf.keras.datasets.mnist
      (train_images, train_labels), (
          test_images,
          test_labels,
      ) = dataset.load_data()
  
      # Convert from (28, 28) images to (28, 28, 1)
      train_images = train_images[..., tf.newaxis].astype("float32")
      test_images = test_images[..., tf.newaxis].astype("float32")
  
      # One hot encore labels 0, 1, .., 9 to [0, 0, .., 1, 0, 0]
      train_labels = tf.keras.utils.to_categorical(
          train_labels, num_classes=num_classes
      )
      test_labels = tf.keras.utils.to_categorical(
          test_labels, num_classes=num_classes
      )
  
      return (
          train_images[0:500],
          train_labels[0:500],
          test_images[0:300],
          test_labels[0:300],
      )
  
imports: import numpy
  import pytest
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: mnist_dataset(num_classes) and returns: train_images,  train_labels,  test_images,  test_labels