# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# def get_config():
# 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 average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# 
# Path: libml/data.py
# _DATA_CACHE = None
# DATA_DIR = os.environ['ML_DATA']
# FLAGS = flags.FLAGS
# DATASETS = {}
# def record_parse(serialized_example):
# def default_parse(dataset: tf.data.Dataset, parse_fn=record_parse) -> tf.data.Dataset:
# def dataset(filenames: list) -> tf.data.Dataset:
# def memoize(dataset: tf.data.Dataset) -> tf.data.Dataset:
#     def tf_get(index):
#         def get(index):
# def augment_mirror(x):
# def augment_shift(x, w):
# def augment_noise(x, std):
# def compute_mean_std(data: tf.data.Dataset):
#         def iterator():
#     def __init__(self, name, train_labeled, train_unlabeled, test, valid, eval_labeled, eval_unlabeled,
#                  height=32, width=32, colors=3, nclass=10, mean=0, std=1, p_labeled=None, p_unlabeled=None):
#     def creator(cls, name, seed, label, valid, augment, parse_fn=default_parse, do_memoize=True, colors=3,
#                 nclass=10, height=32, width=32, name_suffix=''):
#         def create():
# class DataSet:
# 
# Path: libml/layers.py
# def smart_shape(x):
# def entropy_from_logits(logits):
# def entropy_penalty(logits, entropy_penalty_multiplier, mask):
# def kl_divergence_from_logits(logits_a, logits_b):
# def mse_from_logits(output_logits, target_logits):
# def interleave_offsets(batch, nu):
# def interleave(xy, batch):
# def renorm(v):
# def shakeshake(a, b, training):
#     def __init__(self, name, nclass, buf_size):
#     def __call__(self):
#     def update(self, entry):
#     def __init__(self, dataset: DataSet):
#     def __call__(self):
#     def update(self, entry, decay=0.999):
#     def __init__(self, mode):
#     def augment_pair(x0, l0, x1, l1, beta, **kwargs):
#     def augment(x, l, beta, **kwargs):
#     def __call__(self, xl: list, ll: list, betal: list):
# class PMovingAverage:
# class PData:
# class MixMode:
#     MODES = 'xx.yy xxy.yxy xx.yxy xx.yx xx. .yy xxy. .yxy .'.split()
# 
# Path: libml/models.py
# class CNN13(ClassifySemi):
# class ConvNet(ClassifySemi):
# class ResNet(ClassifySemi):
# class ShakeNet(ClassifySemi):
# class MultiModel(CNN13, ConvNet, ResNet, ShakeNet):
#     def classifier(self, x, scales, filters, training, getter=None, **kwargs):
#     def classifier(self, x, scales, filters, getter=None, **kwargs):
#     def classifier(self, x, scales, filters, repeat, training, getter=None, **kwargs):
#         def conv_args(k, f):
#         def residual(x0, filters, stride=1, activate_before_residual=False):
#     def classifier(self, x, scales, filters, repeat, training, getter=None, **kwargs):
#         def conv_args(k, f):
#         def residual(x0, filters, stride=1):
#             def branch():
#     def augment(self, x, l, smoothing, **kwargs):
#     def classifier(self, x, arch, **kwargs):
#     MODEL_CNN13, MODEL_CONVNET, MODEL_RESNET, MODEL_SHAKE = 'cnn13 convnet resnet shake'.split()
#     MODELS = MODEL_CNN13, MODEL_CONVNET, MODEL_RESNET, MODEL_SHAKE


import functools
import os
import tensorflow as tf
from absl import app
from absl import flags
from easydict import EasyDict
from libml import utils, data, layers, models
from third_party import vat_utils
        tf.summary.scalar('losses/xe', loss)
        tf.summary.scalar('losses/vat', loss_vat)
        tf.summary.scalar('losses/entmin', loss_entmin)

        ema = tf.train.ExponentialMovingAverage(decay=ema)
        ema_op = ema.apply(utils.model_vars())
        ema_getter = functools.partial(utils.getter_ema, ema)
        post_ops.append(ema_op)
        post_ops.extend([tf.assign(v, v * (1 - wd)) for v in utils.model_vars('classify') if 'kernel' in v.name])

        train_op = tf.train.AdamOptimizer(lr).minimize(loss + loss_vat * warmup * vat + entmin_weight * loss_entmin,
                                                       colocate_gradients_with_ops=True)
        with tf.control_dependencies([train_op]):
            train_op = tf.group(*post_ops)

        # Tuning op: only retrain batch norm.
        skip_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        classifier(x_in, training=True)
        train_bn = tf.group(*[v for v in tf.get_collection(tf.GraphKeys.UPDATE_OPS)
                              if v not in skip_ops])

        return EasyDict(
            x=x_in, y=y_in, label=l_in, train_op=train_op, tune_op=train_bn,
            classify_raw=tf.nn.softmax(classifier(x_in, training=False)),  # No EMA, for debugging.
            classify_op=tf.nn.softmax(classifier(x_in, getter=ema_getter, training=False)))


def main(argv):
    del argv  # Unused.
    dataset = data.DATASETS[FLAGS.dataset]()

# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# def get_config():
# 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 average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# 
# Path: libml/data.py
# _DATA_CACHE = None
# DATA_DIR = os.environ['ML_DATA']
# FLAGS = flags.FLAGS
# DATASETS = {}
# def record_parse(serialized_example):
# def default_parse(dataset: tf.data.Dataset, parse_fn=record_parse) -> tf.data.Dataset:
# def dataset(filenames: list) -> tf.data.Dataset:
# def memoize(dataset: tf.data.Dataset) -> tf.data.Dataset:
#     def tf_get(index):
#         def get(index):
# def augment_mirror(x):
# def augment_shift(x, w):
# def augment_noise(x, std):
# def compute_mean_std(data: tf.data.Dataset):
#         def iterator():
#     def __init__(self, name, train_labeled, train_unlabeled, test, valid, eval_labeled, eval_unlabeled,
#                  height=32, width=32, colors=3, nclass=10, mean=0, std=1, p_labeled=None, p_unlabeled=None):
#     def creator(cls, name, seed, label, valid, augment, parse_fn=default_parse, do_memoize=True, colors=3,
#                 nclass=10, height=32, width=32, name_suffix=''):
#         def create():
# class DataSet:
# 
# Path: libml/layers.py
# def smart_shape(x):
# def entropy_from_logits(logits):
# def entropy_penalty(logits, entropy_penalty_multiplier, mask):
# def kl_divergence_from_logits(logits_a, logits_b):
# def mse_from_logits(output_logits, target_logits):
# def interleave_offsets(batch, nu):
# def interleave(xy, batch):
# def renorm(v):
# def shakeshake(a, b, training):
#     def __init__(self, name, nclass, buf_size):
#     def __call__(self):
#     def update(self, entry):
#     def __init__(self, dataset: DataSet):
#     def __call__(self):
#     def update(self, entry, decay=0.999):
#     def __init__(self, mode):
#     def augment_pair(x0, l0, x1, l1, beta, **kwargs):
#     def augment(x, l, beta, **kwargs):
#     def __call__(self, xl: list, ll: list, betal: list):
# class PMovingAverage:
# class PData:
# class MixMode:
#     MODES = 'xx.yy xxy.yxy xx.yxy xx.yx xx. .yy xxy. .yxy .'.split()
# 
# Path: libml/models.py
# class CNN13(ClassifySemi):
# class ConvNet(ClassifySemi):
# class ResNet(ClassifySemi):
# class ShakeNet(ClassifySemi):
# class MultiModel(CNN13, ConvNet, ResNet, ShakeNet):
#     def classifier(self, x, scales, filters, training, getter=None, **kwargs):
#     def classifier(self, x, scales, filters, getter=None, **kwargs):
#     def classifier(self, x, scales, filters, repeat, training, getter=None, **kwargs):
#         def conv_args(k, f):
#         def residual(x0, filters, stride=1, activate_before_residual=False):
#     def classifier(self, x, scales, filters, repeat, training, getter=None, **kwargs):
#         def conv_args(k, f):
#         def residual(x0, filters, stride=1):
#             def branch():
#     def augment(self, x, l, smoothing, **kwargs):
#     def classifier(self, x, arch, **kwargs):
#     MODEL_CNN13, MODEL_CONVNET, MODEL_RESNET, MODEL_SHAKE = 'cnn13 convnet resnet shake'.split()
#     MODELS = MODEL_CNN13, MODEL_CONVNET, MODEL_RESNET, MODEL_SHAKE


# Path: vat.py
import functools
import os
import tensorflow as tf
from absl import app
from absl import flags
from easydict import EasyDict
from libml import utils, data, layers, models
from third_party import vat_utils

        tf.summary.scalar('losses/xe', loss)
        tf.summary.scalar('losses/vat', loss_vat)
        tf.summary.scalar('losses/entmin', loss_entmin)

        ema = tf.train.ExponentialMovingAverage(decay=ema)
        ema_op = ema.apply(utils.model_vars())
        ema_getter = functools.partial(utils.getter_ema, ema)
        post_ops.append(ema_op)
        post_ops.extend([tf.assign(v, v * (1 - wd)) for v in utils.model_vars('classify') if 'kernel' in v.name])

        train_op = tf.train.AdamOptimizer(lr).minimize(loss + loss_vat * warmup * vat + entmin_weight * loss_entmin,
                                                       colocate_gradients_with_ops=True)
        with tf.control_dependencies([train_op]):
            train_op = tf.group(*post_ops)

        # Tuning op: only retrain batch norm.
        skip_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        classifier(x_in, training=True)
        train_bn = tf.group(*[v for v in tf.get_collection(tf.GraphKeys.UPDATE_OPS)
                              if v not in skip_ops])

        return EasyDict(
            x=x_in, y=y_in, label=l_in, train_op=train_op, tune_op=train_bn,
            classify_raw=tf.nn.softmax(classifier(x_in, training=False)),  # No EMA, for debugging.
            classify_op=tf.nn.softmax(classifier(x_in, getter=ema_getter, training=False)))


def main(argv):
    del argv  # Unused.
    dataset = data.DATASETS[FLAGS.dataset]()

    log_width = utils.ilog2(dataset.width)