# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
# def get_config():
# def setup_main():
# def setup_tf():
# def smart_shape(x):
# def ilog2(x):
# def find_latest_checkpoint(dir, glob_term='model.ckpt-*.meta'):
# def get_latest_global_step(dir):
# def get_latest_global_step_in_subdir(dir):
# def getter_ema(ema, getter, name, *args, **kwargs):
# def model_vars(scope=None):
# def gpu(x):
# def get_available_gpus():
# def get_gpu():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# def interleave(x, batch):
# def de_interleave(x, batch):
# def combine_dicts(*args):
# 
# Path: libml/data.py
# PAIR_DATASETS = functools.partial(create_datasets, augment_module.pair_augment_function)


import functools
import os
import numpy as np
import tensorflow as tf
from absl import app
from absl import flags
from tqdm import trange
from libml import models, utils
from libml.data import PAIR_DATASETS
        return tf.stop_gradient(logits)

    @staticmethod
    def kl_divergence_from_logits(p_logits, q_logits):
        p = tf.nn.softmax(p_logits)
        log_p = tf.nn.log_softmax(p_logits)
        log_q = tf.nn.log_softmax(q_logits)
        kl = tf.reduce_sum(p * (log_p - log_q), -1)
        return kl

    @staticmethod
    def entropy_from_logits(logits):
        log_prob = tf.nn.log_softmax(logits, axis=-1)
        prob = tf.exp(log_prob)
        ent = tf.reduce_sum(-prob * log_prob, axis=-1)
        return ent

    def train(self, train_nimg, report_nimg):
        if FLAGS.eval_ckpt:
            self.eval_checkpoint(FLAGS.eval_ckpt)
            return
        batch = FLAGS.batch
        train_labeled = self.dataset.train_labeled.repeat().shuffle(FLAGS.shuffle).parse().augment()
        train_labeled = train_labeled.batch(batch).prefetch(16).make_one_shot_iterator().get_next()
        train_unlabeled = self.dataset.train_unlabeled.repeat().shuffle(FLAGS.shuffle).parse().augment()
        train_unlabeled = train_unlabeled.batch(batch * self.params['uratio']).prefetch(16)
        train_unlabeled = train_unlabeled.make_one_shot_iterator().get_next()
        scaffold = tf.train.Scaffold(saver=tf.train.Saver(max_to_keep=FLAGS.keep_ckpt,
                                                          pad_step_number=10))


# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
# def get_config():
# def setup_main():
# def setup_tf():
# def smart_shape(x):
# def ilog2(x):
# def find_latest_checkpoint(dir, glob_term='model.ckpt-*.meta'):
# def get_latest_global_step(dir):
# def get_latest_global_step_in_subdir(dir):
# def getter_ema(ema, getter, name, *args, **kwargs):
# def model_vars(scope=None):
# def gpu(x):
# def get_available_gpus():
# def get_gpu():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# def interleave(x, batch):
# def de_interleave(x, batch):
# def combine_dicts(*args):
# 
# Path: libml/data.py
# PAIR_DATASETS = functools.partial(create_datasets, augment_module.pair_augment_function)


# Path: uda.py
import functools
import os
import numpy as np
import tensorflow as tf
from absl import app
from absl import flags
from tqdm import trange
from libml import models, utils
from libml.data import PAIR_DATASETS

        return tf.stop_gradient(logits)

    @staticmethod
    def kl_divergence_from_logits(p_logits, q_logits):
        p = tf.nn.softmax(p_logits)
        log_p = tf.nn.log_softmax(p_logits)
        log_q = tf.nn.log_softmax(q_logits)
        kl = tf.reduce_sum(p * (log_p - log_q), -1)
        return kl

    @staticmethod
    def entropy_from_logits(logits):
        log_prob = tf.nn.log_softmax(logits, axis=-1)
        prob = tf.exp(log_prob)
        ent = tf.reduce_sum(-prob * log_prob, axis=-1)
        return ent

    def train(self, train_nimg, report_nimg):
        if FLAGS.eval_ckpt:
            self.eval_checkpoint(FLAGS.eval_ckpt)
            return
        batch = FLAGS.batch
        train_labeled = self.dataset.train_labeled.repeat().shuffle(FLAGS.shuffle).parse().augment()
        train_labeled = train_labeled.batch(batch).prefetch(16).make_one_shot_iterator().get_next()
        train_unlabeled = self.dataset.train_unlabeled.repeat().shuffle(FLAGS.shuffle).parse().augment()
        train_unlabeled = train_unlabeled.batch(batch * self.params['uratio']).prefetch(16)
        train_unlabeled = train_unlabeled.make_one_shot_iterator().get_next()
        scaffold = tf.train.Scaffold(saver=tf.train.Saver(max_to_keep=FLAGS.keep_ckpt,
                                                          pad_step_number=10))


        with tf.Session(config=utils.get_config()) as sess: