# Path: semi_supervised/utils.py
# def apply_model_semi(img_unsup, img_sup, is_training, outputs, **kw):
# 
# Path: semi_supervised/vat_utils.py
# def virtual_adversarial_perturbation_direction(
#     x, logits, net, xi=1e-6, num_power_method_iters=1):
# def kl_divergence_from_logits(logits1, logits2):
# def normalize(v):


import functools
import tensorflow as tf
import datasets
import trainer
import utils
import tpu_ops
from semi_supervised import utils as ss_utils
from semi_supervised import vat_utils
  # Forward them both through the model. The scope is needed for tf.Hub export.
  with tf.variable_scope('module'):
    # Here, we pass both inputs to `apply_model_semi`, and so we now get
    # outputs corresponding to each in `end_points` as "rotations_unsup" and
    # similar, which we will use below.
    end_points = ss_utils.apply_model_semi(
        images_unsup, images_sup,
        is_training=(mode == tf.estimator.ModeKeys.TRAIN),
        outputs={
            'rotations': num_angles,
            'classes': datasets.get_auxiliary_num_classes(),
        }, normalization_fn=tpu_ops.cross_replica_batch_norm)

  # Compute virtual adversarial perturbation
  # =====

  def classification_net_fn(x):  # pylint: disable=missing-docstring
    with tf.variable_scope('module', reuse=True):
      end_points_x = ss_utils.apply_model_semi(
          x, None,
          is_training=(mode == tf.estimator.ModeKeys.TRAIN),
          outputs={'classes': datasets.get_auxiliary_num_classes()},
          # Don't update batch norm stats as we're running this on perturbed
          # (corrupted) inputs. Setting decay=1 is what does the trick.
          normalization_fn=functools.partial(
              tpu_ops.cross_replica_batch_norm, decay=1.0))
      return end_points_x['classes']

  vat_eps = FLAGS.get_flag_value('vat_eps', 1.0)
  vat_num_power_method_iters = FLAGS.get_flag_value('vat_num_power_method_iters', 1)

# Path: semi_supervised/utils.py
# def apply_model_semi(img_unsup, img_sup, is_training, outputs, **kw):
# 
# Path: semi_supervised/vat_utils.py
# def virtual_adversarial_perturbation_direction(
#     x, logits, net, xi=1e-6, num_power_method_iters=1):
# def kl_divergence_from_logits(logits1, logits2):
# def normalize(v):


# Path: semi_supervised/rotation_vat.py
import functools
import tensorflow as tf
import datasets
import trainer
import utils
import tpu_ops
from semi_supervised import utils as ss_utils
from semi_supervised import vat_utils

  # Forward them both through the model. The scope is needed for tf.Hub export.
  with tf.variable_scope('module'):
    # Here, we pass both inputs to `apply_model_semi`, and so we now get
    # outputs corresponding to each in `end_points` as "rotations_unsup" and
    # similar, which we will use below.
    end_points = ss_utils.apply_model_semi(
        images_unsup, images_sup,
        is_training=(mode == tf.estimator.ModeKeys.TRAIN),
        outputs={
            'rotations': num_angles,
            'classes': datasets.get_auxiliary_num_classes(),
        }, normalization_fn=tpu_ops.cross_replica_batch_norm)

  # Compute virtual adversarial perturbation
  # =====

  def classification_net_fn(x):  # pylint: disable=missing-docstring
    with tf.variable_scope('module', reuse=True):
      end_points_x = ss_utils.apply_model_semi(
          x, None,
          is_training=(mode == tf.estimator.ModeKeys.TRAIN),
          outputs={'classes': datasets.get_auxiliary_num_classes()},
          # Don't update batch norm stats as we're running this on perturbed
          # (corrupted) inputs. Setting decay=1 is what does the trick.
          normalization_fn=functools.partial(
              tpu_ops.cross_replica_batch_norm, decay=1.0))
      return end_points_x['classes']

  vat_eps = FLAGS.get_flag_value('vat_eps', 1.0)
  vat_num_power_method_iters = FLAGS.get_flag_value('vat_num_power_method_iters', 1)

  vat_perturbation = vat_utils.virtual_adversarial_perturbation_direction(