reject sampling with tf dataset api

test.py

(tf.data.Dataset.apply(tf.contrib.data.rejection_resample(class_func=lambda x: tf.py_func(self.get_class_label, [x],
                                                                                          [tf.int64],
                                                                                          name='class_label_fn'),
                                                          target_dist=tf.constant(self.uniform_dist, dtype=tf.float32),
                                                          initial_dist=tf.constant(self.init_dist, dtype=tf.float32)))
                .map(lambda _, x: x)
                .shuffle(buffer_size=1000)
                .batch(self._args.batch_size)
                .map(lambda x: tf.py_func(self._make_batch, [x], tf.string, name='train_mkbatch_fn'))
                .prefetch(self._args.batch_size * self._args.prefetch_factor))

def get_class_label(self, sample):
    return json.loads(sample)['label']