jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
if is_tf_available():
    import tensorflow as tf


if is_keras_nlp_available():
    from transformers.models.gpt2 import TFGPT2Tokenizer


TOKENIZER_CHECKPOINTS = ["openai-community/gpt2"]
TINY_MODEL_CHECKPOINT = "openai-community/gpt2"

if is_tf_available():

    class ModelToSave(tf.Module):
        def __init__(self, tokenizer):
            super().__init__()
            self.tokenizer = tokenizer
            config = AutoConfig.from_pretrained(TINY_MODEL_CHECKPOINT)
            self.model = TFGPT2LMHeadModel.from_config(config)

        @tf.function(input_signature=(tf.TensorSpec((None,), tf.string, name="text"),))
        def serving(self, text):
            tokenized = self.tokenizer(text)
            input_ids_dense = tokenized["input_ids"].to_tensor()

            input_mask = tf.to_int32(input_ids_dense > 0)
            # input_mask = tf.reshape(input_mask, [-1, MAX_SEQ_LEN])

            outputs = self.model(input_ids=input_ids_dense, attention_mask=input_mask)["logits"]

            return outputs
//  [python]
if is_tf_available():
    import tensorflow as tf


if is_keras_nlp_available():
    from transformers.models.gpt2 import TFGPT2Tokenizer


TOKENIZER_CHECKPOINTS = ["openai-community/gpt2"]
TINY_MODEL_CHECKPOINT = "openai-community/gpt2"

if is_tf_available():

    class ModelToSave(tf.Module):
        def __init__(self, tokenizer):
            super().__init__()
            self.tokenizer = tokenizer
            config = AutoConfig.from_pretrained(TINY_MODEL_CHECKPOINT)
            self.model = TFGPT2LMHeadModel.from_config(config)

        @tf.function(input_signature=(tf.TensorSpec((None,), tf.string, name="text"),))
        def serving(self, text):
            tokenized = self.tokenizer(text)
            input_ids_dense = tokenized["input_ids"].to_tensor()

            input_mask = tf.cast(input_ids_dense > 0, tf.int32)
            # input_mask = tf.reshape(input_mask, [-1, MAX_SEQ_LEN])

            outputs = self.model(input_ids=input_ids_dense, attention_mask=input_mask)["logits"]

            return outputs
//  [python]
def _nelem(x):
    nelem = t.reduce_sum(t.cast(~t.math.is_nan(x), t.float32))
    return t.cast(t.where(t.equal(nelem, 0.), 1., nelem), x.dtype)
//  [python]
def _nelem(x):
    nelem = tf.reduce_sum(tf.cast(~tf.math.is_nan(x), tf.float32))
    return tf.cast(tf.where(tf.equal(nelem, 0.), 1., nelem), x.dtype)
//  [python]
def create_optimizer(self):
        """
        建立模型公用的优化器
        """
        self.global_step = tf.get_variable("step", dtype=tf.int32, initializer=tf.constant(0))
        self.lr = tf.get_variable("learning_rate", dtype=tf.float32, initializer=tf.constant(self.args.lr))

        global_step = tf.to_float(self.global_step)
        lr_next = tf.reduce_max([self.args.lr * tf.pow(self.args.lr_decay_rate,
            tf.to_float(global_step / self.args.lr_decay_steps)), self.args.min_lr])

        tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, tf.assign(self.lr, lr_next, name="update_lr"))
        optimizer = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=self.args.beta1, beta2=self.args.beta2)
        return optimizer
//  [python]
def create_optimizer(self):
        """
        建立模型公用的优化器
        """
        self.global_step = tf.get_variable("step", dtype=tf.int32, initializer=tf.constant(0))
        self.lr = tf.get_variable("learning_rate", dtype=tf.float32, initializer=tf.constant(self.args.lr))

        global_step = tf.cast(self.global_step, dtype=tf.float32)
        lr_next = tf.reduce_max([self.args.lr * tf.pow(self.args.lr_decay_rate,
            tf.cast(global_step / self.args.lr_decay_steps, dtype=tf.float32)), self.args.min_lr])

        tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, tf.assign(self.lr, lr_next, name="update_lr"))
        optimizer = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=self.args.beta1, beta2=self.args.beta2)
        return optimizer
//  [python]
def poisson_loss(y_true, y_pred):
    y_pred = t.cast(y_pred, t.float32)
    y_true = t.cast(y_true, t.float32)

    # we can use the Possion PMF from TensorFlow as well
    # dist = tf.contrib.distributions
    # return -tf.reduce_mean(dist.Poisson(y_pred).log_pmf(y_true))

    nelem = _nelem(y_true)
    y_true = _nan2zero(y_true)

    # last term can be avoided since it doesn't depend on y_pred
    # however keeping it gives a nice lower bound to zero
    ret = y_pred - y_true*tf.math.log(y_pred+1e-10) + tf.math.lgamma(y_true+1.0)

    return tf.divide(tf.reduce_sum(ret), nelem)
//  [python]
def poisson_loss(y_true, y_pred):
    y_pred = tf.cast(y_pred, tf.float32)
    y_true = tf.cast(y_true, tf.float32)

    # we can use the Possion PMF from TensorFlow as well
    # dist = tf.contrib.distributions
    # return -tf.reduce_mean(dist.Poisson(y_pred).log_pmf(y_true))

    nelem = _nelem(y_true)
    y_true = _nan2zero(y_true)

    # last term can be avoided since it doesn't depend on y_pred
    # however keeping it gives a nice lower bound to zero
    ret = y_pred - y_true*tf.math.log(y_pred+1e-10) + tf.math.lgamma(y_true+1.0)

    return tf.divide(tf.reduce_sum(ret), nelem)
//  [python]
def loss(self, y_true, y_pred, mean=True):
        scale_factor = self.scale_factor
        eps = self.eps

        with tf.name_scope(self.scope):
            y_true = tf.cast(y_true, tf.float32)
            y_pred = tf.cast(y_pred, tf.float32) * scale_factor

            if self.masking:
                nelem = _nelem(y_true)
                y_true = _nan2zero(y_true)

            # Clip theta
            theta = tf.minimum(self.theta, 1e6)

            t1 = tf.math.lgamma(theta+eps) + tf.math.lgamma(y_true+1.0) - tf.math.lgamma(y_true+theta+eps)
            t2 = (theta+y_true) * tf.math.log(1.0 + (y_pred/(theta+eps))) + (y_true * (tf.math.log(theta+eps) - tf.math.log(y_pred+eps)))

            if self.debug:
                assert_ops = [
                        tf.verify_tensor_all_finite(y_pred, 'y_pred has inf/nans'),
                        tf.verify_tensor_all_finite(t1, 't1 has inf/nans'),
                        tf.verify_tensor_all_finite(t2, 't2 has inf/nans')]

                tf.summary.histogram('t1', t1)
                tf.summary.histogram('t2', t2)

                with tf.control_dependencies(assert_ops):
                    final = t1 + t2

            else:
                final = t1 + t2

            final = _nan2inf(final)

            if mean:
                if self.masking:
                    final = tf.divide(tf.reduce_sum(final), nelem)
                else:
                    final = tf.reduce_mean(final)


        return final
//  [python]
def loss(self, y_true, y_pred, mean=True):
        scale_factor = self.scale_factor
        eps = self.eps

        with tf.name_scope(self.scope):
            # reuse existing NB neg.log.lik.
            # mean is always False here, because everything is calculated
            # element-wise. we take the mean only in the end
            nb_case = super().loss(y_true, y_pred, mean=False) - tf.math.log(1.0-self.pi+eps)

            y_true = tf.cast(y_true, tf.float32)
            y_pred = tf.cast(y_pred, tf.float32) * scale_factor
            theta = tf.minimum(self.theta, 1e6)

            zero_nb = tf.pow(theta/(theta+y_pred+eps), theta)
            zero_case = -tf.math.log(self.pi + ((1.0-self.pi)*zero_nb)+eps)
            result = tf.where(tf.less(y_true, 1e-8), zero_case, nb_case)
            ridge = self.ridge_lambda*tf.square(self.pi)
            result += ridge

            if mean:
                if self.masking:
                    result = _reduce_mean(result)
                else:
                    result = tf.reduce_mean(result)

            result = _nan2inf(result)

            if self.debug:
                tf.summary.histogram('nb_case', nb_case)
                tf.summary.histogram('zero_nb', zero_nb)
                tf.summary.histogram('zero_case', zero_case)
                tf.summary.histogram('ridge', ridge)

        return result
//  [python]
def _build_inputs(self):

        self._queries = tf.placeholder(tf.float32, [self._batch_size, self._input_dim], name="queries") # utterances to be classified

        # Histories of the utterances
        self._histories_own = tf.placeholder(tf.float32, [self._batch_size, self._timesteps, self._input_dim], name="histories_own") 
        self._histories_other = tf.placeholder(tf.float32, [self._batch_size, self._timesteps, self._input_dim], name="histories_other")
        self._histories_own_mask = tf.cast(tf.placeholder(tf.float32, [self._batch_size, self._timesteps], name="histories_own_mask"), dtype=tf.bool)
        self._histories_other_mask = tf.cast(tf.placeholder(tf.float32, [self._batch_size, self._timesteps], name="histories_other_mask"), dtype=tf.bool)
        self._mask = tf.cast(tf.placeholder(tf.float32, [self._batch_size, self._timesteps], name="global_mask"), dtype=tf.bool)

        # True Labels
        self._labels = tf.placeholder(tf.int32, [self._batch_size, self._class_size], name="labels")
        

        # Learning Rate
        self._lr = tf.placeholder(tf.float32, [], name="learning_rate")

        # Dropout Probability
        self._dropout = tf.placeholder(tf.float32, [], name="dropout_keep_rate")

        # 
        self._training = tf.placeholder(tf.bool, [], name="training_testing_mode")
//  [python]
def build_graph(self, real_sample):
        real_sample = tf.expand_dims(real_sample, -1)

        # sample the latent code:
        zc = shapeless_placeholder(sample_prior(BATCH), 0, name='z_code')
        z_noise = shapeless_placeholder(
            tf.random_uniform([BATCH, NOISE_DIM], -1, 1), 0, name='z_noise')
        z = tf.concat([zc, z_noise], 1, name='z')

        with argscope([Conv2D, Conv2DTranspose, FullyConnected],
                      kernel_initializer=tf.truncated_normal_initializer(stddev=0.02)):
            with tf.variable_scope('gen'):
                fake_sample = self.generator(z)
                fake_sample_viz = tf.cast((fake_sample) * 255.0, tf.uint8, name='viz')
                tf.summary.image('gen', fake_sample_viz, max_outputs=30)

            # may need to investigate how bn stats should be updated across two discrim
            with tf.variable_scope('discrim'):
                real_pred, _ = self.discriminator(real_sample)
                fake_pred, dist_param = self.discriminator(fake_sample)

        """
        Mutual information between x (i.e. zc in this case) and some
        information s (the generated samples in this case):

                    I(x;s) = H(x) - H(x|s)
                           = H(x) + E[\log P(x|s)]

        The distribution from which zc is sampled, in this case, is set to a fixed prior already.
        So the first term is a constant.
        For the second term, we can maximize its variational lower bound:
                    E_{x \sim P(x|s)}[\log Q(x|s)]
        where Q(x|s) is a proposal distribution to approximate P(x|s).

        Here, Q(x|s) is assumed to be a distribution which shares the form
        of P, and whose parameters are predicted by the discriminator network.
        """
        with tf.name_scope("mutual_information"):
            with tf.name_scope('prior_entropy'):
                cat, uni = get_distributions(DIST_PRIOR_PARAM[:NUM_CLASS], DIST_PRIOR_PARAM[NUM_CLASS:])
                ents = [cat.entropy(name='cat_entropy'), tf.reduce_sum(uni.entropy(), name='uni_entropy')]
                entropy = tf.add_n(ents, name='total_entropy')
                # Note that the entropy of prior is a constant. The paper mentioned it but didn't use it.

            with tf.name_scope('conditional_entropy'):
                cond_ents = entropy_from_samples(zc, dist_param)
                cond_entropy = tf.add_n(cond_ents, name="total_entropy")

            MI = tf.subtract(entropy, cond_entropy, name='mutual_information')
            summary.add_moving_summary(entropy, cond_entropy, MI, *cond_ents)

        # default GAN objective
        self.build_losses(real_pred, fake_pred)

        # subtract mutual information for latent factors (we want to maximize them)
        self.g_loss = tf.subtract(self.g_loss, MI, name='total_g_loss')
        self.d_loss = tf.subtract(self.d_loss, MI, name='total_d_loss')

        summary.add_moving_summary(self.g_loss, self.d_loss)

        # distinguish between variables of generator and discriminator updates
        self.collect_variables()
//  [python]
def compute_loss(self, logits, target_output, target_sequence_length, unk_helper=True):
        """Compute optimization loss."""
        max_time = target_output.shape[1].value or tf.shape(target_output)[1]

        crossent = self._softmax_cross_entropy_loss(
            logits, target_output)

        target_weights = tf.sequence_mask(target_sequence_length, max_time, dtype=tf.float32)
        if unk_helper:
            is_unk = tf.equal(target_output, 0)
            unk_val = tf.to_float(is_unk)
            unk_val = unk_val / tf.reduce_sum(unk_val, keep_dims=True)
            unk_weights = unk_val * target_weights
            target_weights = tf.where(is_unk, unk_weights, target_weights)
        loss = crossent * target_weights
        return loss
//  [python]
def compute_loss(self, logits, target_output, target_sequence_length, unk_helper=True):
        """Compute optimization loss."""
        max_time = target_output.shape[1].value or tf.shape(target_output)[1]

        crossent = self._softmax_cross_entropy_loss(
            logits, target_output)

        target_weights = tf.sequence_mask(target_sequence_length, max_time, dtype=tf.float32)
        if unk_helper:
            is_unk = tf.equal(target_output, 0)
            unk_val = tf.cast(is_unk, tf.float32)
            unk_val = unk_val / tf.reduce_sum(unk_val, keep_dims=True)
            unk_weights = unk_val * target_weights
            target_weights = tf.where(is_unk, unk_weights, target_weights)
        loss = crossent * target_weights
        return loss
//  [python]
def call(self, inputs):
        inputs = t.cast(inputs, tf.int32)
        outputs = tf.gather(self.embeddings, inputs)
        return outputs
//  [python]
def call(self, inputs):
        inputs = tf.cast(inputs, tf.int32)
        outputs = tf.gather(self.embeddings, inputs)
        return outputs
//  [python]
def build_mock_model(path):
   """Build and save the mock model with the given signature."""
   module = tf.Module()
   # We have to set this useless variable in order for the TF C API to correctly
   # intake it
   module.var = tf.Variable(0, dtype=tf.float32)

   def action(*inputs):
     s1 = tf.reduce_sum([
         tf.to_float(inputs[0][key]) for key in PER_LIVEINTERVAL_FEATURE_LIST
     ],
         axis=0)
     s2 = tf.reduce_sum(
         [tf.to_float(inputs[0][key]) for key in CONTEXT_FEATURE_LIST])
     # Add a large number so s won't be 0.
     s = s1 + s2
     result = s + module.var
     return {POLICY_DECISION_LABEL: result}
   module.action = tf.function()(action)
   action = {
       'action': module.action.get_concrete_function(get_input_signature())
   }

   tf.saved_model.save(module, path, signatures=action)
   output_spec_path = get_output_spec_path(path)
   with open(output_spec_path, 'w') as f:
     print(f'Writing output spec to {output_spec_path}.')
     f.write(POLICY_OUTPUT_SPEC)
//  [python]
def build_mock_model(path):
   """Build and save the mock model with the given signature."""
   module = tf.Module()
   # We have to set this useless variable in order for the TF C API to correctly
   # intake it
   module.var = tf.Variable(0, dtype=tf.float32)

   def action(*inputs):
     s1 = tf.reduce_sum([
         tf.cast(inputs[0][key], tf.float32) for key in PER_LIVEINTERVAL_FEATURE_LIST
     ],
         axis=0)
     s2 = tf.reduce_sum(
         [tf.cast(inputs[0][key], tf.float32) for key in CONTEXT_FEATURE_LIST])
     # Add a large number so s won't be 0.
     s = s1 + s2
     result = s + module.var
     return {POLICY_DECISION_LABEL: result}
   module.action = tf.function()(action)
   action = {
       'action': module.action.get_concrete_function(get_input_signature())
   }

   tf.saved_model.save(module, path, signatures=action)
   output_spec_path = get_output_spec_path(path)
   with open(output_spec_path, 'w') as f:
     print(f'Writing output spec to {output_spec_path}.')
     f.write(POLICY_OUTPUT_SPEC)
//  [python]
def create_multiple_gaussian_map(coords_wh, output_size, sigma, valid_vec=None, extra=False):
        """ Creates a map of size (output_shape[0], output_shape[1]) at (center[0], center[1])
            with variance sigma for multiple coordinates."""
        with tf.name_scope('create_multiple_gaussian_map'):
            sigma = tf.to_float(sigma)
            assert len(output_size) == 2
            s = coords_wh.get_shape().as_list()
            coords_wh = tf.to_int32(coords_wh)
            if valid_vec is not None:
                valid_vec = tf.to_float(valid_vec)
                valid_vec = tf.squeeze(valid_vec)
                cond_val = tf.greater(valid_vec, 0.5)
            else:
                cond_val = tf.ones_like(coords_wh[:, 0], dtype=tf.float32)
                cond_val = tf.greater(cond_val, 0.5)

            cond_1_in = tf.logical_and(tf.less(coords_wh[:, 0], output_size[0] - 1), tf.greater(coords_wh[:, 0], 0))
            cond_2_in = tf.logical_and(tf.less(coords_wh[:, 1], output_size[1] - 1), tf.greater(coords_wh[:, 1], 0))
            cond_in = tf.logical_and(cond_1_in, cond_2_in)
            cond = tf.logical_and(cond_val, cond_in)

            coords_wh = tf.to_float(coords_wh)

            # create meshgrid
            x_range = tf.expand_dims(tf.range(output_size[0]), 1)
            y_range = tf.expand_dims(tf.range(output_size[1]), 0)

            X = tf.to_float(tf.tile(x_range, [1, output_size[1]]))
            Y = tf.to_float(tf.tile(y_range, [output_size[0], 1]))

            X.set_shape((output_size[0], output_size[1]))
            Y.set_shape((output_size[0], output_size[1]))

            X = tf.expand_dims(X, -1)
            Y = tf.expand_dims(Y, -1)

            X_b = tf.tile(X, [1, 1, s[0]])
            Y_b = tf.tile(Y, [1, 1, s[0]])

            X_b -= coords_wh[:, 0]
            Y_b -= coords_wh[:, 1]

            dist = tf.square(X_b) + tf.square(Y_b)

            scoremap = tf.exp(-dist / (2 * tf.square(sigma))) * tf.to_float(cond)

            if extra:
                negative = 1 - tf.reduce_sum(scoremap, axis=2, keep_dims=True)
                negative = tf.minimum(tf.maximum(negative, 0.0), 1.0)
                scoremap = tf.concat([scoremap, negative], axis=2)

            return scoremap
//  [python]
def create_multiple_gaussian_map(coords_wh, output_size, sigma, valid_vec=None, extra=False):
        """ Creates a map of size (output_shape[0], output_shape[1]) at (center[0], center[1])
            with variance sigma for multiple coordinates."""
        with tf.name_scope('create_multiple_gaussian_map'):
            sigma = tf.cast(sigma, tf.float32)
            assert len(output_size) == 2
            s = coords_wh.get_shape().as_list()
            coords_wh = tf.cast(coords_wh, tf.int32)
            if valid_vec is not None:
                valid_vec = tf.cast(valid_vec, tf.float32)
                valid_vec = tf.squeeze(valid_vec)
                cond_val = tf.greater(valid_vec, 0.5)
            else:
                cond_val = tf.ones_like(coords_wh[:, 0], dtype=tf.float32)
                cond_val = tf.greater(cond_val, 0.5)

            cond_1_in = tf.logical_and(tf.less(coords_wh[:, 0], output_size[0] - 1), tf.greater(coords_wh[:, 0], 0))
            cond_2_in = tf.logical_and(tf.less(coords_wh[:, 1], output_size[1] - 1), tf.greater(coords_wh[:, 1], 0))
            cond_in = tf.logical_and(cond_1_in, cond_2_in)
            cond = tf.logical_and(cond_val, cond_in)

            coords_wh = tf.cast(coords_wh, tf.float32)

            # create meshgrid
            x_range = tf.expand_dims(tf.range(output_size[0]), 1)
            y_range = tf.expand_dims(tf.range(output_size[1]), 0)

            X = tf.cast(tf.tile(x_range, [1, output_size[1]]), tf.float32)
            Y = tf.cast(tf.tile(y_range, [output_size[0], 1]), tf.float32)

            X.set_shape((output_size[0], output_size[1]))
            Y.set_shape((output_size[0], output_size[1]))

            X = tf.expand_dims(X, -1)
            Y = tf.expand_dims(Y, -1)

            X_b = tf.tile(X, [1, 1, s[0]])
            Y_b = tf.tile(Y, [1, 1, s[0]])

            X_b -= coords_wh[:, 0]
            Y_b -= coords_wh[:, 1]

            dist = tf.square(X_b) + tf.square(Y_b)

            scoremap = tf.exp(-dist / (2 * tf.square(sigma))) * tf.cast(cond, tf.float32)

            if extra:
                negative = 1 - tf.reduce_sum(scoremap, axis=2, keep_dims=True)
                negative = tf.minimum(tf.maximum(negative, 0.0), 1.0)
                scoremap = tf.concat([scoremap, negative], axis=2)

            return scoremap
//  [python]
def create_multiple_gaussian_map_3d(keypoint_3d, output_size, sigma3d, valid_vec=None, extra=False):
        """ Creates a 3D heatmap for the hand skeleton
        """
        with tf.name_scope('create_multiple_gaussian_map_3d'):
            if valid_vec is not None:
                valid_vec = tf.to_float(valid_vec)
                valid_vec = tf.squeeze(valid_vec)
                cond_val = tf.greater(valid_vec, 0.5)
            else:
                cond_val = tf.ones_like(keypoint_3d[:, 0], dtype=tf.float32)
                cond_val = tf.greater(cond_val, 0.5)

            sigma3d = tf.to_float(sigma3d)
            # reverse the order of axis: tensorflow uses NDHWC
            reverse = keypoint_3d[:, ::-1]

            z_range = tf.expand_dims(tf.expand_dims(tf.range(output_size, dtype=tf.float32), 1), 2)
            y_range = tf.expand_dims(tf.expand_dims(tf.range(output_size, dtype=tf.float32), 0), 2)
            x_range = tf.expand_dims(tf.expand_dims(tf.range(output_size, dtype=tf.float32), 0), 1)

            Z = tf.tile(z_range, [1, output_size, output_size])
            Y = tf.tile(y_range, [output_size, 1, output_size])
            X = tf.tile(x_range, [output_size, output_size, 1])

            Z = tf.expand_dims(Z, -1)
            Y = tf.expand_dims(Y, -1)
            X = tf.expand_dims(X, -1)

            s = reverse.get_shape().as_list()

            Z_b = tf.tile(Z, [1, 1, 1, s[0]])
            Y_b = tf.tile(Y, [1, 1, 1, s[0]])
            X_b = tf.tile(X, [1, 1, 1, s[0]])

            Z_b -= reverse[:, 0]
            Y_b -= reverse[:, 1]
            X_b -= reverse[:, 2]

            dist = tf.square(X_b) + tf.square(Y_b) + tf.square(Z_b)

            scoremap_3d = tf.exp(-dist / (2 * tf.square(sigma3d))) * tf.to_float(cond_val)

            if extra:
                negative = 1 - tf.reduce_sum(scoremap_3d, axis=3, keep_dims=True)
                negative = tf.minimum(tf.maximum(negative, 0.0), 1.0)
                scoremap_3d = tf.concat([scoremap_3d, negative], axis=3)

            return scoremap_3d
//  [python]
def create_multiple_gaussian_map_3d(keypoint_3d, output_size, sigma3d, valid_vec=None, extra=False):
        """ Creates a 3D heatmap for the hand skeleton
        """
        with tf.name_scope('create_multiple_gaussian_map_3d'):
            if valid_vec is not None:
                valid_vec = tf.cast(valid_vec, tf.float32)
                valid_vec = tf.squeeze(valid_vec)
                cond_val = tf.greater(valid_vec, 0.5)
            else:
                cond_val = tf.ones_like(keypoint_3d[:, 0], dtype=tf.float32)
                cond_val = tf.greater(cond_val, 0.5)

            sigma3d = tf.cast(sigma3d, tf.float32)
            # reverse the order of axis: tensorflow uses NDHWC
            reverse = keypoint_3d[:, ::-1]

            z_range = tf.expand_dims(tf.expand_dims(tf.range(output_size, dtype=tf.float32), 1), 2)
            y_range = tf.expand_dims(tf.expand_dims(tf.range(output_size, dtype=tf.float32), 0), 2)
            x_range = tf.expand_dims(tf.expand_dims(tf.range(output_size, dtype=tf.float32), 0), 1)

            Z = tf.tile(z_range, [1, output_size, output_size])
            Y = tf.tile(y_range, [output_size, 1, output_size])
            X = tf.tile(x_range, [output_size, output_size, 1])

            Z = tf.expand_dims(Z, -1)
            Y = tf.expand_dims(Y, -1)
            X = tf.expand_dims(X, -1)

            s = reverse.get_shape().as_list()

            Z_b = tf.tile(Z, [1, 1, 1, s[0]])
            Y_b = tf.tile(Y, [1, 1, 1, s[0]])
            X_b = tf.tile(X, [1, 1, 1, s[0]])

            Z_b -= reverse[:, 0]
            Y_b -= reverse[:, 1]
            X_b -= reverse[:, 2]

            dist = tf.square(X_b) + tf.square(Y_b) + tf.square(Z_b)

            scoremap_3d = tf.exp(-dist / (2 * tf.square(sigma3d))) * tf.cast(cond_val, tf.float32)

            if extra:
                negative = 1 - tf.reduce_sum(scoremap_3d, axis=3, keep_dims=True)
                negative = tf.minimum(tf.maximum(negative, 0.0), 1.0)
                scoremap_3d = tf.concat([scoremap_3d, negative], axis=3)

            return scoremap_3d
//  [python]
def call(self, inputs, training=False):
        input_tensor, attention_mask = inputs

        self_outputs = self.self_attention(
            [input_tensor, attention_mask], training=training
        )
        attention_output = self.dense_output(
            [self_outputs[0], input_tensor], training=training
        )
        masked_attention_output = attention_output * tf.cast(
            tf.expand_dims(attention_mask, 2), dtype=attention_output.dtype
        )
        outputs = (masked_attention_output,) + self_outputs[
            1:
        ]  # add attentions if we output them
        return outputs
//  [python]
def call(self, inputs):
        """Call logic."""
        hidden_states, attention_mask = inputs

        hidden_states = self.conv1d_1(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        hidden_states = self.conv1d_2(hidden_states)

        masked_hidden_states = hidden_states * tf.cast(
            tf.expand_dims(attention_mask, 2), dtype=hidden_states.dtype
        )
        return masked_hidden_states
//  [python]
def _build_losses(self, json_data):
        super()._build_losses(json_data)

        disc_weight_decay = 0 if (self.DISC_WEIGHT_DECAY_KEY not in json_data) else json_data[self.DISC_WEIGHT_DECAY_KEY]
        disc_logit_reg_weight = 0 if (self.DISC_LOGIT_REG_WEIGHT_KEY not in json_data) else json_data[self.DISC_LOGIT_REG_WEIGHT_KEY]
        disc_grad_penalty = 0.0 if (self.DISC_GRAD_PENALTY_KEY not in json_data) else json_data[self.DISC_GRAD_PENALTY_KEY]
        
        disc_loss_expert_tf = self.build_disc_loss_pos(self._disc_logits_expert_tf)
        disc_loss_agent_tf = self.build_disc_loss_neg(self._disc_logits_agent_tf)
        disc_loss_expert_tf = tf.reduce_mean(disc_loss_expert_tf)
        disc_loss_agent_tf = tf.reduce_mean(disc_loss_agent_tf)

        self._disc_loss_tf = 0.5 * (disc_loss_agent_tf + disc_loss_expert_tf)

        self._acc_expert_tf = tf.reduce_mean(tf.to_float(tf.greater(self._disc_logits_expert_tf, 0)))
        self._acc_agent_tf = tf.reduce_mean(tf.to_float(tf.less(self._disc_logits_agent_tf, 0)))
        
        if (disc_weight_decay != 0):
            self._disc_loss_tf += disc_weight_decay * self._weight_decay_loss(self.MAIN_SCOPE + "/" + self.DISC_SCOPE)
                    
        if (disc_logit_reg_weight != 0):
            self._disc_loss_tf += disc_logit_reg_weight * self._disc_logit_reg_loss()
        
        if (disc_grad_penalty != 0):
            self._grad_penalty_loss_tf = self._disc_grad_penalty_loss(in_tfs=self._disc_expert_inputs, out_tf=self._disc_logits_expert_tf)
            self._disc_loss_tf += disc_grad_penalty * self._grad_penalty_loss_tf
        else:
            self._grad_penalty_loss_tf = tf.constant(0.0, dtype=tf.float32)

        return
//  [python]
def _build_losses(self, json_data):
        super()._build_losses(json_data)

        disc_weight_decay = 0 if (self.DISC_WEIGHT_DECAY_KEY not in json_data) else json_data[self.DISC_WEIGHT_DECAY_KEY]
        disc_logit_reg_weight = 0 if (self.DISC_LOGIT_REG_WEIGHT_KEY not in json_data) else json_data[self.DISC_LOGIT_REG_WEIGHT_KEY]
        disc_grad_penalty = 0.0 if (self.DISC_GRAD_PENALTY_KEY not in json_data) else json_data[self.DISC_GRAD_PENALTY_KEY]
        
        disc_loss_expert_tf = self.build_disc_loss_pos(self._disc_logits_expert_tf)
        disc_loss_agent_tf = self.build_disc_loss_neg(self._disc_logits_agent_tf)
        disc_loss_expert_tf = tf.reduce_mean(disc_loss_expert_tf)
        disc_loss_agent_tf = tf.reduce_mean(disc_loss_agent_tf)

        self._disc_loss_tf = 0.5 * (disc_loss_agent_tf + disc_loss_expert_tf)

        self._acc_expert_tf = tf.reduce_mean(tf.cast(tf.greater(self._disc_logits_expert_tf, 0), tf.float32))
        self._acc_agent_tf = tf.reduce_mean(tf.cast(tf.less(self._disc_logits_agent_tf, 0), tf.float32))
        
        if (disc_weight_decay != 0):
            self._disc_loss_tf += disc_weight_decay * self._weight_decay_loss(self.MAIN_SCOPE + "/" + self.DISC_SCOPE)
                    
        if (disc_logit_reg_weight != 0):
            self._disc_loss_tf += disc_logit_reg_weight * self._disc_logit_reg_loss()
        
        if (disc_grad_penalty != 0):
            self._grad_penalty_loss_tf = self._disc_grad_penalty_loss(in_tfs=self._disc_expert_inputs, out_tf=self._disc_logits_expert_tf)
            self._disc_loss_tf += disc_grad_penalty * self._grad_penalty_loss_tf
        else:
            self._grad_penalty_loss_tf = tf.constant(0.0, dtype=tf.float32)

        return
//  [python]
def call(self, inputs, training=False):
        """Call logic."""
        hidden_states, attention_mask = inputs

        attention_outputs = self.attention(
            [hidden_states, attention_mask], training=training
        )
        attention_output = attention_outputs[0]
        intermediate_output = self.intermediate(
            [attention_output, attention_mask], training=training
        )
        layer_output = self.bert_output(
            [intermediate_output, attention_output], training=training
        )
        masked_layer_output = layer_output * tf.cast(
            tf.expand_dims(attention_mask, 2), dtype=layer_output.dtype
        )
        outputs = (masked_layer_output,) + attention_outputs[
            1:
        ]  # add attentions if we output them
        return outputs
//  [python]
def call(self, inputs, training=False):
        """Call logic."""
        outputs, mask = inputs
        extended_mask = tf.to_float(tf.expand_dims(mask, axis=2))
        for i, (conv, bn) in enumerate(self.conv_batch_norm):
            outputs = conv(outputs)
            outputs = bn(outputs)
            outputs = self.activation[i](outputs)
            outputs = self.dropout(outputs, training=training)
        return outputs * extended_mask
//  [python]
def call(self, inputs, training=False):
        """Call logic."""
        outputs, mask = inputs
        extended_mask = tf.cast(tf.expand_dims(mask, axis=2), outputs.dtype)
        for i, (conv, bn) in enumerate(self.conv_batch_norm):
            outputs = conv(outputs)
            outputs = bn(outputs)
            outputs = self.activation[i](outputs)
            outputs = self.dropout(outputs, training=training)
        return outputs * extended_mask
//  [python]
def call(self, inputs, training=False):
        """Call logic."""
        hidden_states, attention_mask = inputs

        batch_size = tf.shape(hidden_states)[0]
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.to_float(
            tf.shape(key_layer)[-1]
        )  # scale attention_scores
        attention_scores = attention_scores / tf.math.sqrt(dk)

        if attention_mask is not None:
            # extended_attention_masks for self attention encoder.
            extended_attention_mask = attention_mask[:, tf.newaxis, tf.newaxis, :]
            extended_attention_mask = tf.to_float(
                extended_attention_mask
            )
            extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
            attention_scores = attention_scores + extended_attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(attention_scores, axis=-1)
        attention_probs = self.dropout(attention_probs, training=training)

        context_layer = tf.matmul(attention_probs, value_layer)
        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        context_layer = tf.reshape(context_layer, (batch_size, -1, self.all_head_size))

        outputs = (
            (context_layer, attention_probs)
            if self.output_attentions
            else (context_layer,)
        )
        return outputs
//  [python]
def call(self, inputs, training=False):
        """Call logic."""
        hidden_states, attention_mask = inputs

        batch_size = tf.shape(hidden_states)[0]
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.cast(
            tf.shape(key_layer)[-1], attention_scores.dtype
        )  # scale attention_scores
        attention_scores = attention_scores / tf.math.sqrt(dk)

        if attention_mask is not None:
            # extended_attention_masks for self attention encoder.
            extended_attention_mask = attention_mask[:, tf.newaxis, tf.newaxis, :]
            extended_attention_mask = tf.cast(
                extended_attention_mask, attention_scores.dtype
            )
            extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
            attention_scores = attention_scores + extended_attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(attention_scores, axis=-1)
        attention_probs = self.dropout(attention_probs, training=training)

        context_layer = tf.matmul(attention_probs, value_layer)
        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        context_layer = tf.reshape(context_layer, (batch_size, -1, self.all_head_size))

        outputs = (
            (context_layer, attention_probs)
            if self.output_attentions
            else (context_layer,)
        )
        return outputs
//  [python]
def call(self, inputs, training=False):
        """Call logic."""
        encoder_hidden_states, attention_mask = inputs
        attention_mask = t.cast(
            t.expand_dims(attention_mask, 2), encoder_hidden_states.dtype
        )

        # mask encoder hidden states
        masked_encoder_hidden_states = encoder_hidden_states * attention_mask

        # pass though first layer
        outputs = self.conv_layers_sequence(masked_encoder_hidden_states)
        outputs = self.output_layer(outputs)
        masked_outputs = outputs * attention_mask
        return t.squeeze(t.nn.relu6(masked_outputs), -1)
//  [python]
def call(self, inputs, training=False):
        """Call logic."""
        encoder_hidden_states, attention_mask = inputs
        attention_mask = tf.cast(
            tf.expand_dims(attention_mask, 2), encoder_hidden_states.dtype
        )

        # mask encoder hidden states
        masked_encoder_hidden_states = encoder_hidden_states * attention_mask

        # pass though first layer
        outputs = self.conv_layers_sequence(masked_encoder_hidden_states)
        outputs = self.output_layer(outputs)
        masked_outputs = outputs * attention_mask
        return tf.squeeze(tf.nn.relu6(masked_outputs), -1)
//  [python]
def call(self, inputs, training=False):
        hidden_states, speaker_ids, encoder_mask, decoder_pos = inputs

        if self.is_compatible_encoder is False:
            hidden_states = self.project_compatible_decoder(hidden_states)

        # calculate new hidden states.
        hidden_states += t.cast(
            self.decoder_positional_embeddings(decoder_pos), hidden_states.dtype
        )

        if self.config.n_speakers > 1:
            speaker_embeddings = self.decoder_speaker_embeddings(speaker_ids)
            speaker_features = tf.math.softplus(self.speaker_fc(speaker_embeddings))
            # extended speaker embeddings
            extended_speaker_features = speaker_features[:, tf.newaxis, :]
            hidden_states += extended_speaker_features

        return super().call([hidden_states, encoder_mask], training=training)
//  [python]
def call(self, inputs, training=False):
        hidden_states, speaker_ids, encoder_mask, decoder_pos = inputs

        if self.is_compatible_encoder is False:
            hidden_states = self.project_compatible_decoder(hidden_states)

        # calculate new hidden states.
        hidden_states += tf.cast(
            self.decoder_positional_embeddings(decoder_pos), hidden_states.dtype
        )

        if self.config.n_speakers > 1:
            speaker_embeddings = self.decoder_speaker_embeddings(speaker_ids)
            speaker_features = tf.math.softplus(self.speaker_fc(speaker_embeddings))
            # extended speaker embeddings
            extended_speaker_features = speaker_features[:, tf.newaxis, :]
            hidden_states += extended_speaker_features

        return super().call([hidden_states, encoder_mask], training=training)
//  [python]
def eval_metrics_fn():
    return {
        "accuracy": lambda labels, predictions: tf.equal(
            tf.argmax(predictions, 1, output_type=tf.int32),
            tf.to_int32(tf.reshape(labels, [-1])),
        )
    }
//  [python]
def eval_metrics_fn():
    return {
        "accuracy": lambda labels, predictions: tf.equal(
            tf.argmax(predictions, 1, output_type=tf.int32),
            tf.cast(tf.reshape(labels, [-1]), tf.int32),
        )
    }
//  [python]
def _parse_data(record):
    feature_description = {
        "image": tf.io.FixedLenFeature([28, 28], tf.float32),
        "label": tf.io.FixedLenFeature([1], tf.int64),
    }
    r = tf.io.parse_single_example(record, feature_description)
    features = {
        "image": tf.math.divide(tf.cast(r["image"], tf.float32), 255.0)
    }
    return features["image"], tf.to_int32(r["label"])
//  [python]
def _parse_data(record):
    feature_description = {
        "image": tf.io.FixedLenFeature([28, 28], tf.float32),
        "label": tf.io.FixedLenFeature([1], tf.int64),
    }
    r = tf.io.parse_single_example(record, feature_description)
    features = {
        "image": tf.math.divide(tf.cast(r["image"], tf.float32), 255.0)
    }
    return features["image"], tf.cast(r["label"], tf.int32)
//  [python]
def _augment_hsv_darknet(image, rh, rs, rv, seed=None):
  """Randomize the hue, saturation, and brightness via the darknet method."""
  if rh > 0.0:
    deltah = random_uniform_strong(-rh, rh, seed=seed)
    image = tf.image.adjust_hue(image, deltah)
  if rs > 0.0:
    deltas = random_scale(rs, seed=seed)
    image = tf.image.adjust_saturation(image, deltas)
  if rv > 0.0:
    deltav = random_scale(rv, seed=seed)
    image *= tf.to_float(deltav)

  # clip the values of the image between 0.0 and 1.0
  image = tf.clip_by_value(image, 0.0, 1.0)
  return image
//  [python]
def _augment_hsv_darknet(image, rh, rs, rv, seed=None):
  """Randomize the hue, saturation, and brightness via the darknet method."""
  if rh > 0.0:
    deltah = random_uniform_strong(-rh, rh, seed=seed)
    image = tf.image.adjust_hue(image, deltah)
  if rs > 0.0:
    deltas = random_scale(rs, seed=seed)
    image = tf.image.adjust_saturation(image, deltas)
  if rv > 0.0:
    deltav = random_scale(rv, seed=seed)
    image *= tf.cast(deltav, image.dtype)

  # clip the values of the image between 0.0 and 1.0
  image = tf.clip_by_value(image, 0.0, 1.0)
  return image
//  [python]
def _inference(self, input_ids, speaker_ids, speed_ratios, **kwargs):
        """Call logic."""
        attention_mask = tf.math.not_equal(input_ids, 0)
        embedding_output = self.embeddings([input_ids, speaker_ids], training=False)
        encoder_output = self.encoder(
            [embedding_output, attention_mask], training=False
        )
        last_encoder_hidden_states = encoder_output[0]

        # duration predictor, here use last_encoder_hidden_states, u can use more hidden_states layers
        # rather than just use last_hidden_states of encoder for duration_predictor.
        duration_outputs = self.duration_predictor(
            [last_encoder_hidden_states, attention_mask]
        )  # [batch_size, length]
        duration_outputs = tf.math.exp(duration_outputs) - 1.0

        if speed_ratios is None:
            speed_ratios = tf.convert_to_tensor(np.array([1.0]), dtype=tf.float32)

        speed_ratios = tf.expand_dims(speed_ratios, 1)

        duration_outputs = tf.to_int32(
            tf.math.round(duration_outputs * speed_ratios)
        )

        length_regulator_outputs, encoder_masks = self.length_regulator(
            [last_encoder_hidden_states, duration_outputs], training=False
        )

        # create decoder positional embedding
        decoder_pos = tf.range(
            1, tf.shape(length_regulator_outputs)[1] + 1, dtype=tf.int32
        )
        masked_decoder_pos = tf.expand_dims(decoder_pos, 0) * encoder_masks

        decoder_output = self.decoder(
            [length_regulator_outputs, speaker_ids, encoder_masks, masked_decoder_pos],
            training=False,
        )
        last_decoder_hidden_states = decoder_output[0]

        # here u can use sum or concat more than 1 hidden states layers from decoder.
        mel_before = self.mel_dense(last_decoder_hidden_states)
        mel_after = (
            self.postnet([mel_before, encoder_masks], training=False) + mel_before
        )

        outputs = (mel_before, mel_after, duration_outputs)
        return outputs
//  [python]
def _inference(self, input_ids, speaker_ids, speed_ratios, **kwargs):
        """Call logic."""
        attention_mask = tf.math.not_equal(input_ids, 0)
        embedding_output = self.embeddings([input_ids, speaker_ids], training=False)
        encoder_output = self.encoder(
            [embedding_output, attention_mask], training=False
        )
        last_encoder_hidden_states = encoder_output[0]

        # duration predictor, here use last_encoder_hidden_states, u can use more hidden_states layers
        # rather than just use last_hidden_states of encoder for duration_predictor.
        duration_outputs = self.duration_predictor(
            [last_encoder_hidden_states, attention_mask]
        )  # [batch_size, length]
        duration_outputs = tf.math.exp(duration_outputs) - 1.0

        if speed_ratios is None:
            speed_ratios = tf.convert_to_tensor(np.array([1.0]), dtype=tf.float32)

        speed_ratios = tf.expand_dims(speed_ratios, 1)

        duration_outputs = tf.cast(
            tf.math.round(duration_outputs * speed_ratios), tf.int32
        )

        length_regulator_outputs, encoder_masks = self.length_regulator(
            [last_encoder_hidden_states, duration_outputs], training=False
        )

        # create decoder positional embedding
        decoder_pos = tf.range(
            1, tf.shape(length_regulator_outputs)[1] + 1, dtype=tf.int32
        )
        masked_decoder_pos = tf.expand_dims(decoder_pos, 0) * encoder_masks

        decoder_output = self.decoder(
            [length_regulator_outputs, speaker_ids, encoder_masks, masked_decoder_pos],
            training=False,
        )
        last_decoder_hidden_states = decoder_output[0]

        # here u can use sum or concat more than 1 hidden states layers from decoder.
        mel_before = self.mel_dense(last_decoder_hidden_states)
        mel_after = (
            self.postnet([mel_before, encoder_masks], training=False) + mel_before
        )

        outputs = (mel_before, mel_after, duration_outputs)
        return outputs
//  [python]
def conv_model(feature, target, mode):
    """2-layer convolution model."""
    # Convert the target to a one-hot tensor of shape (batch_size, 10) and
    # with a on-value of 1 for each one-hot vector of length 10.
    target = tf.one_hot(tf.to_int32(target), 10, 1, 0)

    # Reshape feature to 4d tensor with 2nd and 3rd dimensions being
    # image width and height final dimension being the number of color
    # channels.
    feature = tf.reshape(feature, [-1, 28, 28, 1])

    # First conv layer will compute 32 features for each 5x5 patch
    with tf.variable_scope("conv_layer1"):
        h_conv1 = layers.conv2d(
            feature,
            32,
            kernel_size=[5, 5],
            activation=tf.nn.relu,
            padding="SAME",
        )
        h_pool1 = tf.nn.max_pool(
            h_conv1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME"
        )

    # Second conv layer will compute 64 features for each 5x5 patch.
    with tf.variable_scope("conv_layer2"):
        h_conv2 = layers.conv2d(
            h_pool1,
            64,
            kernel_size=[5, 5],
            activation=tf.nn.relu,
            padding="SAME",
        )
        h_pool2 = tf.nn.max_pool(
            h_conv2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME"
        )
        # reshape tensor into a batch of vectors
        h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])

    # Densely connected layer with 1024 neurons.
    h_fc1 = layers.dropout(
        layers.dense(h_pool2_flat, 1024, activation=tf.nn.relu),
        rate=0.5,
        training=mode == tf.estimator.ModeKeys.TRAIN,
    )

    # Compute logits (1 per class) and compute loss.
    logits = layers.dense(h_fc1, 10, activation=None)
    loss = tf.losses.softmax_cross_entropy(target, logits)

    return tf.argmax(logits, 1), loss
//  [python]
def conv_model(feature, target, mode):
    """2-layer convolution model."""
    # Convert the target to a one-hot tensor of shape (batch_size, 10) and
    # with a on-value of 1 for each one-hot vector of length 10.
    target = tf.one_hot(tf.cast(target, tf.int32), 10, 1, 0)

    # Reshape feature to 4d tensor with 2nd and 3rd dimensions being
    # image width and height final dimension being the number of color
    # channels.
    feature = tf.reshape(feature, [-1, 28, 28, 1])

    # First conv layer will compute 32 features for each 5x5 patch
    with tf.variable_scope("conv_layer1"):
        h_conv1 = layers.conv2d(
            feature,
            32,
            kernel_size=[5, 5],
            activation=tf.nn.relu,
            padding="SAME",
        )
        h_pool1 = tf.nn.max_pool(
            h_conv1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME"
        )

    # Second conv layer will compute 64 features for each 5x5 patch.
    with tf.variable_scope("conv_layer2"):
        h_conv2 = layers.conv2d(
            h_pool1,
            64,
            kernel_size=[5, 5],
            activation=tf.nn.relu,
            padding="SAME",
        )
        h_pool2 = tf.nn.max_pool(
            h_conv2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME"
        )
        # reshape tensor into a batch of vectors
        h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])

    # Densely connected layer with 1024 neurons.
    h_fc1 = layers.dropout(
        layers.dense(h_pool2_flat, 1024, activation=tf.nn.relu),
        rate=0.5,
        training=mode == tf.estimator.ModeKeys.TRAIN,
    )

    # Compute logits (1 per class) and compute loss.
    logits = layers.dense(h_fc1, 10, activation=None)
    loss = tf.losses.softmax_cross_entropy(target, logits)

    return tf.argmax(logits, 1), loss