# 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

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




FLAGS = tf.flags.FLAGS


def model_fn(data, mode):
  """Produces a loss for the rotation task with semi-supervision.

  Args:
    data: Dict of inputs containing, among others, "image" and "label."
    mode: model's mode: training, eval or prediction

  Returns:
    EstimatorSpec
  """
  num_angles = 4

  # In this mode (called once at the end of training), we create the tf.Hub
  # module in order to export the model, and use that to do one last prediction.
  if mode == tf.estimator.ModeKeys.PREDICT:
    # This defines a function called by the hub module to create the model.
    def model_building_fn(img, is_training):
      # This is an example of calling `apply_model_semi` with only one of the
      # inputs provided. The outputs will simply use the given names:

# 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


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




FLAGS = tf.flags.FLAGS


def model_fn(data, mode):
  """Produces a loss for the rotation task with semi-supervision.

  Args:
    data: Dict of inputs containing, among others, "image" and "label."
    mode: model's mode: training, eval or prediction

  Returns:
    EstimatorSpec
  """
  num_angles = 4

  # In this mode (called once at the end of training), we create the tf.Hub
  # module in order to export the model, and use that to do one last prediction.
  if mode == tf.estimator.ModeKeys.PREDICT:
    # This defines a function called by the hub module to create the model.
    def model_building_fn(img, is_training):
      # This is an example of calling `apply_model_semi` with only one of the
      # inputs provided. The outputs will simply use the given names:

      end_points = ss_utils.apply_model_semi(img, None, is_training, outputs={