code: def load_dataset(
      train: bool = True,
      with_predictions: bool = True,
      batch_size: t.Optional[int] = None,
      shuffle: bool = False,
      n_samples: int = None,
      object_type="VisionData",
  ) -> VisionData:
      """Return MNIST VisionData, containing prediction produced by a simple fully connected model.
  
      Model and data are taken from https://www.tensorflow.org/tutorials/quickstart/beginner.
  
      Parameters
      ----------
      train : bool, default : True
          Train or Test dataset
      with_predictions : bool, default : True
          Whether the returned VisonData should contain predictions
      batch_size: int, optional
          how many samples per batch to load
      shuffle : bool , default : False
          To reshuffled data at every epoch or not.
      n_samples : int, optional
          Number of samples to load. Return the first n_samples if shuffle is False otherwise selects n_samples at random.
          If None, returns all samples.
      object_type : str, default : 'VisionData'
          Kept for compatibility with torch datasets. Not used.
      Returns
      -------
      :obj:`deepchecks.vision.VisionData`
      """
      if object_type != "VisionData":
          raise ValueError(
              "only VisionData is supported for MNIST dataset"
          )
  
      batch_size = batch_size or (64 if train else 1000)
  
      if with_predictions:
          model = load_model()
      else:
          model = None
  
      return VisionData(
          batch_loader=mnist_generator(
              shuffle, batch_size, train, n_samples, model
          ),
          task_type="classification",
          dataset_name=f'mnist {"train" if train else "test"}',
          reshuffle_data=False,
      )
  
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 1 for calling a global function with example usage: load_dataset(train, with_predictions, batch_size, shuffle, n_samples, object_type) and returns: VisionData