code: def load_mnist_data(
      train: bool = True, n_samples: int = None, shuffle: bool = False
  ) -> t.Tuple[np.array, np.array]:
      """Load MNIST dataset.
  
      Parameters
      ----------
      train : bool, default : True
          Train or Test dataset
      n_samples : int, optional
          Number of samples to load.
      shuffle : bool , default : False
          whether to shuffle the data or not.
      Returns
      -------
      Tuple[np.ndarray, np.ndarray]
      """
      # Load the dataset
      if train:
          (images, labels), _ = keras.datasets.mnist.load_data()
      else:
          _, (images, labels) = keras.datasets.mnist.load_data()
  
      if shuffle:
          indices = np.random.permutation(len(images))
          images = images[indices]
          labels = labels[indices]
          del indices
  
      if n_samples is not None:
          images = images[:n_samples]
          labels = labels[:n_samples]
  
      # Normalize the data and expand the data to 4D
      images = images / 255.0
      images = np.expand_dims(images, axis=-1)
  
      return images, labels
  
file_path: deepchecks/deepchecks/vision/datasets/classification/mnist_tensorflow.py
imports: import logging
  import pathlib
  import pickle
  import typing
  import numpy
  from deepchecks.vision.utils.test_utils import hash_image
  from deepchecks.vision.vision_data import VisionData
label_desc: Use image node_id 3 for calling a global function with example usage: load_mnist_data(train, n_samples, shuffle) and returns: images,  labels