def input_fn():
      """Create tf.dataset for estimator.

      Returns:
        tf dataset
      """
      tensor_shape = self.num_entities
      ds = tf.data.Dataset.from_generator(
          generator=set_intersection_generator,
          output_types=((tf.float32, tf.float32, tf.float32), tf.int32),
          output_shapes=((tf.TensorShape([tensor_shape]),
                          tf.TensorShape([tensor_shape]),
                          tf.TensorShape([tensor_shape])), tf.TensorShape([])))

      if n_take > 0:
        ds = ds.take(n_take)
      if shuffle:
        ds = ds.shuffle(SEED)
      ds = ds.repeat(epochs)
      ds = ds.batch(batch_size)
      return ds
Create tf.dataset for estimator.

      Returns:
        tf dataset
      
language/emql/data_loader.py