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

      Returns:
        tf dataset
      """
      # output_types
      if task == '1c':
        output_types = tuple([tf.int32] * 2)
        output_shapes = tuple([tf.TensorShape([])] * 2)
      elif task == '2c':
        output_types = tuple([tf.int32] * 3)
        output_shapes = tuple([tf.TensorShape([])] * 3)
      elif task == '3c':
        output_types = tuple([tf.int32] * 4)
        output_shapes = tuple([tf.TensorShape([])] * 4)
      elif task == '2i' or task == '2u':
        output_types = tuple([tf.int32] * 4)
        output_shapes = tuple([tf.TensorShape([])] * 4)
      elif task == '3i':
        output_types = tuple([tf.int32] * 6)
        output_shapes = tuple([tf.TensorShape([])] * 6)
      elif task == 'ic' or task == 'uc':
        output_types = tuple([tf.int32] * 5)
        output_shapes = tuple([tf.TensorShape([])] * 5)
      elif task == 'ci':
        output_types = tuple([tf.int32] * 5)
        output_shapes = tuple([tf.TensorShape([])] * 5)
      else:
        raise ValueError

      ds = tf.data.Dataset.from_generator(
          generator=query2box_generator,
          output_types=(output_types, tf.int32),
          output_shapes=(output_shapes, tf.TensorShape([])))

      ds = ds.batch(batch_size)
      ds = ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
      return ds
Create tf.dataset for estimator.

      Returns:
        tf dataset
      
language/emql/data_loader.py