# 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 combine_dicts(*args):
# 
# Path: libml/data.py
# DATA_DIR = None
# _DATA_CACHE = None
# SAMPLES_PER_CLASS = [1, 2, 3, 4, 5, 10, 25, 100, 400]
# FLAGS = flags.FLAGS
#     DATA_DIR = FLAGS.data_dir or os.environ['ML_DATA']
# DATASETS = functools.partial(create_datasets, augment_module.augment_function)
# PAIR_DATASETS = functools.partial(create_datasets, augment_module.pair_augment_function)
# MANY_DATASETS = functools.partial(create_datasets, augment_module.many_augment_function)
# def _data_setup():
# def record_parse_mnist(serialized_example, image_shape=None):
# def record_parse(serialized_example, image_shape=None):
# def compute_mean_std(data: tf.data.Dataset):
#         def iterator():
#     def __init__(self, data: tf.data.Dataset, augment_fn: AugmentPair, parse_fn=record_parse, image_shape=None):
#     def from_files(cls, filenames: list, augment_fn: AugmentPair, parse_fn=record_parse, image_shape=None):
#             def fetch_dataset(filename):
#     def empty_data(cls, image_shape, augment_fn: AugmentPair = None):
#         def _get_null_input(_):
#     def __getattr__(self, item):
#         def call_and_update(*args, **kwargs):
#     def parse(self):
#     def numpy_augment(self, *args, **kwargs):
#     def augment(self):
#     def memoize(self):
#         def tf_get(index, image_shape):
#             def get(index):
#     def __init__(self, name, train_labeled: DataSet, train_unlabeled: DataSet, test: DataSet, valid: DataSet,
#                  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=record_parse, do_memoize=False,
#                 nclass=10, colors=3, height=32, width=32):
#         def create():
# def create_datasets(augment_fn):
# class DataSet:
# class DataSets:
# 
# Path: libml/models.py
# class CNN13(ClassifySemi):
# class ResNet(ClassifySemi):
# class ShakeNet(ClassifySemi):
# class MultiModel(CNN13, ResNet, ShakeNet):
#     def classifier(self, x, scales, filters, training, getter=None, **kwargs):
#     def classifier(self, x, scales, filters, repeat, training, getter=None, dropout=0, **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, dropout=0, **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):
#     MODELS = ('cnn13', 'resnet', 'shake')
#     MODEL_CNN13, MODEL_RESNET, MODEL_SHAKE = MODELS
# 
# Path: libml/utils.py
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
#         super(EasyDict, self).__init__(*args, **kwargs)
#         self.__dict__ = self


import functools
import os
import tensorflow as tf
from absl import app
from absl import flags
from libml import utils, data, models
from libml.utils import EasyDict
        wd *= lr
        warmup = tf.clip_by_value(tf.to_float(self.step) / (warmup_pos * (FLAGS.train_kimg << 10)), 0, 1)

        classifier = lambda x, **kw: self.classifier(x, **kw, **kwargs).logits
        logits_x = classifier(xt_in, training=True)
        post_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)  # Take only first call to update batch norm.
        logits_y = classifier(y_in, training=True)
        # Get the pseudo-label loss
        loss_pl = tf.nn.sparse_softmax_cross_entropy_with_logits(
            labels=tf.argmax(logits_y, axis=-1), logits=logits_y
        )
        # Masks denoting which data points have high-confidence predictions
        greater_than_thresh = tf.reduce_any(
            tf.greater(tf.nn.softmax(logits_y), threshold),
            axis=-1,
            keepdims=True,
        )
        greater_than_thresh = tf.cast(greater_than_thresh, loss_pl.dtype)
        # Only enforce the loss when the model is confident
        loss_pl *= greater_than_thresh
        # Note that we also average over examples without confident outputs;
        # this is consistent with the realistic evaluation codebase
        loss_pl = tf.reduce_mean(loss_pl)

        loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=l, logits=logits_x)
        loss = tf.reduce_mean(loss)
        tf.summary.scalar('losses/xe', loss)
        tf.summary.scalar('losses/pl', loss_pl)

        ema = tf.train.ExponentialMovingAverage(decay=ema)

# 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 combine_dicts(*args):
# 
# Path: libml/data.py
# DATA_DIR = None
# _DATA_CACHE = None
# SAMPLES_PER_CLASS = [1, 2, 3, 4, 5, 10, 25, 100, 400]
# FLAGS = flags.FLAGS
#     DATA_DIR = FLAGS.data_dir or os.environ['ML_DATA']
# DATASETS = functools.partial(create_datasets, augment_module.augment_function)
# PAIR_DATASETS = functools.partial(create_datasets, augment_module.pair_augment_function)
# MANY_DATASETS = functools.partial(create_datasets, augment_module.many_augment_function)
# def _data_setup():
# def record_parse_mnist(serialized_example, image_shape=None):
# def record_parse(serialized_example, image_shape=None):
# def compute_mean_std(data: tf.data.Dataset):
#         def iterator():
#     def __init__(self, data: tf.data.Dataset, augment_fn: AugmentPair, parse_fn=record_parse, image_shape=None):
#     def from_files(cls, filenames: list, augment_fn: AugmentPair, parse_fn=record_parse, image_shape=None):
#             def fetch_dataset(filename):
#     def empty_data(cls, image_shape, augment_fn: AugmentPair = None):
#         def _get_null_input(_):
#     def __getattr__(self, item):
#         def call_and_update(*args, **kwargs):
#     def parse(self):
#     def numpy_augment(self, *args, **kwargs):
#     def augment(self):
#     def memoize(self):
#         def tf_get(index, image_shape):
#             def get(index):
#     def __init__(self, name, train_labeled: DataSet, train_unlabeled: DataSet, test: DataSet, valid: DataSet,
#                  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=record_parse, do_memoize=False,
#                 nclass=10, colors=3, height=32, width=32):
#         def create():
# def create_datasets(augment_fn):
# class DataSet:
# class DataSets:
# 
# Path: libml/models.py
# class CNN13(ClassifySemi):
# class ResNet(ClassifySemi):
# class ShakeNet(ClassifySemi):
# class MultiModel(CNN13, ResNet, ShakeNet):
#     def classifier(self, x, scales, filters, training, getter=None, **kwargs):
#     def classifier(self, x, scales, filters, repeat, training, getter=None, dropout=0, **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, dropout=0, **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):
#     MODELS = ('cnn13', 'resnet', 'shake')
#     MODEL_CNN13, MODEL_RESNET, MODEL_SHAKE = MODELS
# 
# Path: libml/utils.py
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
#         super(EasyDict, self).__init__(*args, **kwargs)
#         self.__dict__ = self


# Path: pseudo_label.py
import functools
import os
import tensorflow as tf
from absl import app
from absl import flags
from libml import utils, data, models
from libml.utils import EasyDict

        wd *= lr
        warmup = tf.clip_by_value(tf.to_float(self.step) / (warmup_pos * (FLAGS.train_kimg << 10)), 0, 1)

        classifier = lambda x, **kw: self.classifier(x, **kw, **kwargs).logits
        logits_x = classifier(xt_in, training=True)
        post_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)  # Take only first call to update batch norm.
        logits_y = classifier(y_in, training=True)
        # Get the pseudo-label loss
        loss_pl = tf.nn.sparse_softmax_cross_entropy_with_logits(
            labels=tf.argmax(logits_y, axis=-1), logits=logits_y
        )
        # Masks denoting which data points have high-confidence predictions
        greater_than_thresh = tf.reduce_any(
            tf.greater(tf.nn.softmax(logits_y), threshold),
            axis=-1,
            keepdims=True,
        )
        greater_than_thresh = tf.cast(greater_than_thresh, loss_pl.dtype)
        # Only enforce the loss when the model is confident
        loss_pl *= greater_than_thresh
        # Note that we also average over examples without confident outputs;
        # this is consistent with the realistic evaluation codebase
        loss_pl = tf.reduce_mean(loss_pl)

        loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=l, logits=logits_x)
        loss = tf.reduce_mean(loss)
        tf.summary.scalar('losses/xe', loss)
        tf.summary.scalar('losses/pl', loss_pl)

        ema = tf.train.ExponentialMovingAverage(decay=ema)

        ema_op = ema.apply(utils.model_vars())