# Path: awesome_gans/magan/adamax.py
# class AdamaxOptimizer(optimizer.Optimizer):
#     """Optimizer that implements the Adamax algorithm.
#     See [Kingma et. al., 2014](http://arxiv.org/abs/1412.6980)
#     ([pdf](http://arxiv.org/pdf/1412.6980.pdf)).
#     @@__init__
#     """
# 
#     def __init__(self, learning_rate=1e-3, beta1=0.9, beta2=0.999, use_locking=False, name="AdaMaxOptimizer"):
#         super(AdamaxOptimizer, self).__init__(use_locking, name)
#         self._lr = learning_rate
#         self._beta1 = beta1
#         self._beta2 = beta2
# 
#         # Tensor versions of the constructor arguments, created in _prepare().
#         self._lr_t = None
#         self._beta1_t = None
#         self._beta2_t = None
# 
#     def _prepare(self):
#         self._lr_t = ops.convert_to_tensor(self._lr, name="learning_rate")
#         self._beta1_t = ops.convert_to_tensor(self._beta1, name="beta1")
#         self._beta2_t = ops.convert_to_tensor(self._beta2, name="beta2")
# 
#     def _create_slots(self, var_list):
#         # Create slots for the first and second moments.
#         for v in var_list:
#             self._zeros_slot(v, "m", self._name)
#             self._zeros_slot(v, "v", self._name)
# 
#     def _apply_dense(self, grad, var):
#         lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype)
#         beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype)
#         beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype)
#         if var.dtype.base_dtype == tf.float16:
#             eps = 1e-7  # Can't use 1e-8 due to underflow -- not sure if it makes a big difference.
#         else:
#             eps = 1e-8
# 
#         v = self.get_slot(var, "v")
#         v_t = v.assign(beta1_t * v + (1.0 - beta1_t) * grad)
#         m = self.get_slot(var, "m")
#         m_t = m.assign(tf.maximum(beta2_t * m + eps, tf.abs(grad)))
#         g_t = v_t / m_t
# 
#         var_update = state_ops.assign_sub(var, lr_t * g_t)
#         return control_flow_ops.group(*[var_update, m_t, v_t])
# 
#     def _apply_sparse(self, grad, var):
#         raise NotImplementedError("Sparse gradient updates are not supported.")


import sys
import tensorflow as tf
import awesome_gans.modules as t
from awesome_gans.magan.adamax import AdamaxOptimizer

            x = t.deconv2d(x, self.channel, 4, 2, name='gen-deconv2d-5')
            x = tf.nn.tanh(x)
            return x

    def build_magan(self):
        # Generator
        self.g = self.generator(self.z)

        # Discriminator
        _, d_real = self.discriminator(self.x)
        _, d_fake = self.discriminator(self.g, reuse=True)

        self.d_real_loss = t.mse_loss(self.x, d_real, self.batch_size)
        self.d_fake_loss = t.mse_loss(self.g, d_fake, self.batch_size)
        self.d_loss = self.d_real_loss + tf.maximum(0.0, self.m - self.d_fake_loss)
        self.g_loss = self.d_fake_loss

        # Summary
        tf.summary.scalar("loss/d_loss", self.d_loss)
        tf.summary.scalar("loss/d_real_loss", self.d_real_loss)
        tf.summary.scalar("loss/d_fake_loss", self.d_fake_loss)
        tf.summary.scalar("loss/g_loss", self.g_loss)

        # Optimizer
        t_vars = tf.trainable_variables()
        d_params = [v for v in t_vars if v.name.startswith('d')]
        g_params = [v for v in t_vars if v.name.startswith('g')]

        self.d_op = AdamaxOptimizer(learning_rate=self.lr, beta1=self.beta1).minimize(self.d_loss, var_list=d_params)

# Path: awesome_gans/magan/adamax.py
# class AdamaxOptimizer(optimizer.Optimizer):
#     """Optimizer that implements the Adamax algorithm.
#     See [Kingma et. al., 2014](http://arxiv.org/abs/1412.6980)
#     ([pdf](http://arxiv.org/pdf/1412.6980.pdf)).
#     @@__init__
#     """
# 
#     def __init__(self, learning_rate=1e-3, beta1=0.9, beta2=0.999, use_locking=False, name="AdaMaxOptimizer"):
#         super(AdamaxOptimizer, self).__init__(use_locking, name)
#         self._lr = learning_rate
#         self._beta1 = beta1
#         self._beta2 = beta2
# 
#         # Tensor versions of the constructor arguments, created in _prepare().
#         self._lr_t = None
#         self._beta1_t = None
#         self._beta2_t = None
# 
#     def _prepare(self):
#         self._lr_t = ops.convert_to_tensor(self._lr, name="learning_rate")
#         self._beta1_t = ops.convert_to_tensor(self._beta1, name="beta1")
#         self._beta2_t = ops.convert_to_tensor(self._beta2, name="beta2")
# 
#     def _create_slots(self, var_list):
#         # Create slots for the first and second moments.
#         for v in var_list:
#             self._zeros_slot(v, "m", self._name)
#             self._zeros_slot(v, "v", self._name)
# 
#     def _apply_dense(self, grad, var):
#         lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype)
#         beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype)
#         beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype)
#         if var.dtype.base_dtype == tf.float16:
#             eps = 1e-7  # Can't use 1e-8 due to underflow -- not sure if it makes a big difference.
#         else:
#             eps = 1e-8
# 
#         v = self.get_slot(var, "v")
#         v_t = v.assign(beta1_t * v + (1.0 - beta1_t) * grad)
#         m = self.get_slot(var, "m")
#         m_t = m.assign(tf.maximum(beta2_t * m + eps, tf.abs(grad)))
#         g_t = v_t / m_t
# 
#         var_update = state_ops.assign_sub(var, lr_t * g_t)
#         return control_flow_ops.group(*[var_update, m_t, v_t])
# 
#     def _apply_sparse(self, grad, var):
#         raise NotImplementedError("Sparse gradient updates are not supported.")


# Path: awesome_gans/magan/magan_model.py
import sys
import tensorflow as tf
import awesome_gans.modules as t
from awesome_gans.magan.adamax import AdamaxOptimizer


            x = t.deconv2d(x, self.channel, 4, 2, name='gen-deconv2d-5')
            x = tf.nn.tanh(x)
            return x

    def build_magan(self):
        # Generator
        self.g = self.generator(self.z)

        # Discriminator
        _, d_real = self.discriminator(self.x)
        _, d_fake = self.discriminator(self.g, reuse=True)

        self.d_real_loss = t.mse_loss(self.x, d_real, self.batch_size)
        self.d_fake_loss = t.mse_loss(self.g, d_fake, self.batch_size)
        self.d_loss = self.d_real_loss + tf.maximum(0.0, self.m - self.d_fake_loss)
        self.g_loss = self.d_fake_loss

        # Summary
        tf.summary.scalar("loss/d_loss", self.d_loss)
        tf.summary.scalar("loss/d_real_loss", self.d_real_loss)
        tf.summary.scalar("loss/d_fake_loss", self.d_fake_loss)
        tf.summary.scalar("loss/g_loss", self.g_loss)

        # Optimizer
        t_vars = tf.trainable_variables()
        d_params = [v for v in t_vars if v.name.startswith('d')]
        g_params = [v for v in t_vars if v.name.startswith('g')]

        self.d_op = AdamaxOptimizer(learning_rate=self.lr, beta1=self.beta1).minimize(self.d_loss, var_list=d_params)

        self.g_op = AdamaxOptimizer(learning_rate=self.lr, beta1=self.beta1).minimize(self.g_loss, var_list=g_params)