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

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

      ds = ds.batch(batch_size)
      return ds
Create tf.dataset for estimator.

      Returns:
        tf dataset
      
language/emql/data_loader.py