code: def minibatch(x_train, y_train, batch_size, train_epochs):
      """Generate minibatches of data for a set number of epochs."""
      epoch = 0
      start = 0
      key = random.PRNGKey(0)
  
      while epoch < train_epochs:
          end = start + batch_size
  
          if end > x_train.shape[0]:
              key, split = random.split(key)
              permutation = random.permutation(
                  split,
                  np.arange(x_train.shape[0], dtype=np.int64),
                  independent=True,
              )
              x_train = x_train[permutation]
              y_train = y_train[permutation]
              epoch += 1
              start = 0
              continue
  
          yield x_train[start:end], y_train[start:end]
          start = start + batch_size
  
file_path: neural-tangents/examples/datasets.py
imports: import gzip
  import os
  import shutil
  import urllib.request
  from jax import random
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 4 for calling a global function with example usage: minibatch(x_train, y_train, batch_size, train_epochs) without return types