text: import numpy as np
  
  
  #
  # TODO(pkhorrami4) Consider datatype of input (convert to float32, etc.)
  # TODO(pkhorrami4) Decide datatype for conversion from theano.config.floatX
  # TODO(pkhorrami4) Add centering + rescaling of data as option
  # TODO(pkhorrami4) Add verbosity flag option
  #
  
  
  class SupervisedDataset(object):
      # Class to construct a supervised dataset
  
      def __init__(self, X, y):  # , batch_size=None, iter_type=None):
          # print('Dataset loaded with size: {}'.format(X.shape))
          self.X = X
          self.n_samples = X.shape[0]
          self.n_channels = X.shape[1]
          self.height = X.shape[2]
          self.width = X.shape[3]
          self.y = y
  
      def __iter__(self):
          return self.iter  # self.iterator()
  
      def iterator(self,
                   mode='sequential',
                   batch_size=None,
                   num_batches=None,
                   rng_seed=0):
  
          if(batch_size is not None):
              if(batch_size > self.n_samples):
                  raise ValueError("batch_size is larger than number of samples")
  
          if(mode == 'sequential'):
              if batch_size is None:
                  if num_batches is not None:
                      batch_size = int(self.n_samples / num_batches)
                  else:
                      raise ValueError("need one of batch_size, num_batches"
                                       "for sequential batch iteration")
              elif batch_size is not None:
                  if num_batches is not None:
                      max_num_batches = int(np.ceil(self.n_samples / batch_size))
                      if num_batches > max_num_batches:
                          raise ValueError("dataset of %d examples can only "
                                           "provide %d batches with batch_size "
                                           "%d, but %d batches were requested" %
                                           (self.n_samples, max_num_batches,
                                            batch_size, num_batches))
                  else:
                      num_batches = int(self.n_samples / batch_size)
  
              self.iter = DatasetIteratorSequential(self.X, self.y, batch_size,
                                                    num_batches, self.n_samples)
  
          elif(mode == 'random_uniform'):
              if batch_size is None:
                  raise ValueError("batch_size cannot be None for random "
                                   "uniform iteration")
              elif num_batches is None:
                  raise ValueError("num_batches cannot be None for random "
                                   "uniform iteration")
              self.iter = DatasetIteratorRandomUniform(self.X, self.y,
                                                       batch_size, num_batches,
                                                       self.n_samples, rng_seed)
  
          elif(mode == 'random_uniform_no_rep'):
              if batch_size is None:
                  if num_batches is not None:
                      # Floor Function chosen to ensure that uneven segment
                      # (i.e. "runt") is ignored
                      batch_size = int(np.floor(np.float32(self.n_samples) /
                                                num_batches))
                  else:
                      raise ValueError("need one of batch_size, num_batches"
                                       "for sequential batch iteration")
              elif batch_size is not None:
                  if num_batches is not None:
                      max_num_batches = int(np.ceil(self.n_samples / batch_size))
                      if num_batches > max_num_batches:
                          raise ValueError("dataset of %d examples can only "
                                           "provide %d batches with batch_size "
                                           " %d, but %d batches were requested" %
                                           (self.n_samples, max_num_batches,
                                            batch_size, num_batches))
                  else:
                      # Floor Function chosen to ensure that uneven segment
                      # (i.e. "runt") is ignored
                      num_batches = int(np.floor(np.float32(self.n_samples) /
                                                 batch_size))
  
              self.iter = DatasetIteratorRandomUniformNoRep(self.X, self.y,
                                                            batch_size,
                                                            num_batches,
                                                            self.n_samples,
                                                            rng_seed)
  
          else:
              raise ValueError("please specify the iterator mode as either "
                               "('sequential', 'random_uniform', "
                               " 'random_uniform_no_rep')")
  
          return self.iter
  
      def get_num_samples(self):
          return self.n_samples
  
      def get_batch(self):
          return self.iter.next()
  
  
  class BasicIterator(object):
  
      def __init__(self, X, y,
                   batch_size=None,
                   num_batches=None,
                   num_samples=None):
  
          self.X = X
          self.y = y
          self.batch_size = batch_size
          self.num_batches = num_batches
          self.num_samples = num_samples
          # print('Batch Size: {}'.format(batch_size))
          # print('Num Batches: {}'.format(num_batches))
          # print('Num Samples: {}'.format(num_samples))
  
          # Batch Counter
          self.batch_count = 0
          # Sample Counter
          self.sample_count = 0
          # Batch Indices
          self.last = 0
  
      def reset(self):
          self.batch_count = 0
          self.sample_count = 0
          self.last = 0
  
  
  class DatasetIteratorSequential(BasicIterator):
  
      #
      # Iterator that traverses the data by extracting sequential slices
      # of size (batch_size) from the data tensor
      #
  
      def __init__(self, X, y,
                   batch_size=None,
                   num_batches=None,
                   num_samples=None):
          # print('Using Sequential Iterator')
          super(DatasetIteratorSequential, self).__init__(X, y,
                                                          batch_size,
                                                          num_batches,
                                                          num_samples)
  
      def __iter__(self):
          return self
  
      def next(self):
          # print('In Dataset_Iter_Seq_Next')
          if (self.batch_count >= self.num_batches or
             self.sample_count >= self.num_samples):
              raise StopIteration()
          # fixes the problem where (self._num_samples % self._batch_size) != 0
          elif (self.sample_count + self.batch_size) > self.num_samples:
              self.last = slice(self.sample_count, self.num_samples)
              self.sample_count = self.num_samples
              return self.X[self.last, :, :, :]
          else:
              self.last = slice(self.sample_count,
                                self.sample_count + self.batch_size)
              self.sample_count += self.batch_size
              self.batch_count += 1
              return self.X[self.last, :, :, :], self.y[self.last]
  
  
  class DatasetIteratorRandomUniform(BasicIterator):
  
      #
      # Iterator that uniformly samples batches of size (batch_size)
      # from the data tensor (with replacement)
      #
  
      def __init__(self, X, y,
                   batch_size=None,
                   num_batches=None,
                   num_samples=None,
                   rng_seed=0):
          # print('Using Random Uniform Iterator (with replacement)')
          np.random.seed(rng_seed)
          super(DatasetIteratorRandomUniform, self).__init__(X, y, batch_size,
                                                             num_batches,
                                                             num_samples)
  
      def __iter__(self):
          return self
  
      def next(self):
          if self.batch_count >= self.num_batches:
              raise StopIteration()
          else:
              self.last = np.random.random_integers(low=0,
                                                    high=self.num_samples - 1,
                                                    size=(self.batch_size,))
              self.batch_count += 1
              return self.X[self.last, :, :, :], self.y[self.last]
  
  
  class DatasetIteratorRandomUniformNoRep(BasicIterator):
  
      #
      # Iterator that uniformly samples batches of size (batch_size)
      # from the data tensor (without replacement)
      #
  
      def __init__(self, X, y,
                   batch_size=None,
                   num_batches=None,
                   num_samples=None,
                   rng_seed=0):
          # print('Using Random Uniform Iterator (w/o replacement)')
          np.random.seed(rng_seed)
          self.remain_samples = np.arange(num_samples)
          super(DatasetIteratorRandomUniformNoRep, self).__init__(X, y,
                                                                  batch_size,
                                                                  num_batches,
                                                                  num_samples)
  
      def __iter__(self):
          return self
  
      def next(self):
          if self.batch_count >= self.num_batches:
              raise StopIteration()
  
          else:
              idx = np.random.permutation(
                  np.arange(len(self.remain_samples)))[0:self.batch_size]
              self.last = self.remain_samples[idx]
              self.remain_samples = np.setdiff1d(self.remain_samples, self.last)
              self.batch_count += 1
              return self.X[self.last, :, :, :], self.y[self.last]
  
  
  if __name__ == "__main__":
      # print('Making Datset Object')
      X = np.zeros((1024, 3, 5, 5))
      y = np.zeros((1024,))
      d = SupervisedDataset(X, y)
      # d.iterator(batch_size=128)
      # d.iterator(mode='random_uniform', batch_size=128,
      #            num_batches = 1000, rng_seed=10)
      d.iterator(mode='random_uniform_no_rep', batch_size=128)
  
      for i, [b, b_l] in enumerate(d):
          print('Batch {} --- {}'.format(i+1, b.shape))
          print('Batch Labels {}'.format(b_l.shape))
  