code: def get_datasets(
      dir_path,
      classes,
      batch_size=256,
      test_ratio=0.1,
      max_items_per_class=4096,
  ):
      x, y = np.empty([0, 784]), np.empty([0])
      for idx, cls in enumerate(classes):
          cls_path = os.path.join(dir_path, f"{cls}.npy")
          data = np.load(cls_path)[:max_items_per_class, :]
          labels = np.full(data.shape[0], idx)
          x, y = np.concatenate((x, data), axis=0), np.append(y, labels)
  
      assert x.shape[0] % batch_size == 0
  
      x, y = x.astype(jnp.float32) / 255.0, y.astype(jnp.int32)
      # Reshaping to square images
      x = np.reshape(x, (x.shape[0], 28, 28, 1))
      permutation = np.random.permutation(y.shape[0])
      x = x[permutation, :]
      y = y[permutation]
  
      x = np.reshape(
          x, [x.shape[0] // batch_size, batch_size] + list(x.shape[1:])
      )
      y = np.reshape(y, [y.shape[0] // batch_size, batch_size])
  
      nb_test_elements = int(x.shape[0] * test_ratio)
  
      x_test, y_test = x[:nb_test_elements], y[:nb_test_elements]
      x_train, y_train = x[nb_test_elements:], y[nb_test_elements:]
  
      return list(zip(x_train, y_train)), list(zip(x_test, y_test))
  
imports: from jax import numpy
  import numpy
  import os
  import requests
label_desc: Use image node_id 2 for calling a global function with example usage: get_datasets(dir_path, classes, batch_size, test_ratio, max_items_per_class) and returns: list,  list