code: def source_dataset(source, batch_size, seed, dataset_size=None):
      """Returns a `tf.data.Dataset` of samples from `source`."""
      dataset = tf.data.Dataset.random(seed=seed)
      if dataset_size is not None:
          # This rounds up to multiple of batch size.
          batches = (dataset_size - 1) // batch_size + 1
          dataset = dataset.take(batches)
      return dataset.map(
          lambda seed: source.sample(
              batch_size, seed=tf.bitcast(seed, tf.int32)
          )
      )
  
file_path: compression/models/toy_sources/compression_model.py
imports: import abc
  import matplotlib.pyplot
  import numpy
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: source_dataset(source, batch_size, seed, dataset_size) and returns: dataset