code: def __init__(self, batch_size: int, shuffle_buffer_size=None, rank=2):
      self.batch_size = batch_size
      self.shuffle_buffer_size = shuffle_buffer_size
      if rank == 2:
          self.tensor_types = (tf.int32, tf.int32)
          self.tensor_shapes = (
              tf.TensorShape([None]),
              tf.TensorShape([None]),
          )
      elif rank == 3:
          self.tensor_types = (tf.int32, tf.int32, tf.int32)
          self.tensor_shapes = (
              tf.TensorShape([None]),
              tf.TensorShape([None]),
              tf.TensorShape([None]),
          )
      else:
          raise ValueError(
              "Rank must be either 2 or 3, but was: {}".format(rank)
          )
  
imports: from typing import Iterable, Callable
  import tensorflow
label_desc: Use image node_id 1 to create a new DatasetGenerator object  with example: obj = DatasetGenerator(batch_size, shuffle_buffer_size, rank)