code: def load_shakespeare(batch_size, num_steps, token_type):
      data_iter = DataLoader(batch_size, num_steps, token_type)
      return data_iter, data_iter.vocab
  
imports: import collections
  import random
  import re
  import jax
  import jax.numpy
  import tensorflow_datasets
  import objax
  from objax.functional import one_hot
  from objax.zoo.rnn import RNN
label_desc: Use image node_id 5 for calling a global function with example usage: load_shakespeare(batch_size, num_steps, token_type) and returns: data_iter,  data_iter