jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _log_prob(self, data, num_samples=1):
    """Compute a lower bound on the log likelihood."""
    mean_centered_data = data - self.data_mean
    # Tile by num_samples on the batch dimension.
    tiled_mean_centered_data = tf.tile(mean_centered_data,
                                       [num_samples] + [1] * len(self.data_dim))
    tiled_data = tf.tile(data,
                         [num_samples] + [1] * len(self.data_dim))

    # Construct approximate posterior and sample z.
    q_z = self.q(tiled_mean_centered_data)
    z = q_z.sample()  # [num_samples * batch_size, data_dim]
    log_q_z = q_z.log_prob(z)  # [num_samples * batch_size]

    # compute the proposal prob of z, #[num_samples * batch_size]
    try:
      log_p_z = self.proposal.log_prob(z, log_q_data=log_q_z)
    except TypeError:
      log_p_z = self.proposal.log_prob(z)

    # Compute the model logprob of the data
    p_x_given_z = self.decoder(z)
    # [num_samples * batch_size]
    log_p_x_given_z = p_x_given_z.log_prob(tiled_data)

    elbo = log_p_x_given_z + self.kl_weight * (log_p_z - log_q_z)
    iwae = (tf.reduce_logsumexp(tf.reshape(elbo, [num_samples, -1]), axis=0)
            - tf.log(tf.to_float(num_samples)))
    return iwae
//  [python]
def _log_prob(self, data, num_samples=1):
    """Compute a lower bound on the log likelihood."""
    mean_centered_data = data - self.data_mean
    # Tile by num_samples on the batch dimension.
    tiled_mean_centered_data = tf.tile(mean_centered_data,
                                       [num_samples] + [1] * len(self.data_dim))
    tiled_data = tf.tile(data,
                         [num_samples] + [1] * len(self.data_dim))

    # Construct approximate posterior and sample z.
    q_z = self.q(tiled_mean_centered_data)
    z = q_z.sample()  # [num_samples * batch_size, data_dim]
    log_q_z = q_z.log_prob(z)  # [num_samples * batch_size]

    # compute the proposal prob of z, #[num_samples * batch_size]
    try:
      log_p_z = self.proposal.log_prob(z, log_q_data=log_q_z)
    except TypeError:
      log_p_z = self.proposal.log_prob(z)

    # Compute the model logprob of the data
    p_x_given_z = self.decoder(z)
    # [num_samples * batch_size]
    log_p_x_given_z = p_x_given_z.log_prob(tiled_data)

    elbo = log_p_x_given_z + self.kl_weight * (log_p_z - log_q_z)
    iwae = (tf.reduce_logsumexp(tf.reshape(elbo, [num_samples, -1]), axis=0)
            - tf.log(tf.cast(num_samples, tf.float32)))
    return iwae
//  [python]
def create_input_layer(mode, features, params, embeddings):
    """Create input layer."""

    opts = params['opts']
    embedding_dim = opts.embedding_dim
    use_batch_normalization = opts.use_batch_normalization
    training = (mode == tf.estimator.ModeKeys.TRAIN)

    with tf.name_scope("input_layer"):
        if training:
            records_column = params['records_column']
        else:
            records_column = params['predict_records_column']

        records_column = tf.feature_column.input_layer(
            features, records_column)

        nonzero_mask = tf.to_float(
            tf.not_equal(records_column, model_keys.PADDING_ID))
        tf.summary.histogram('nonzero_cnt', nonzero_mask)
        nonzero_mask = nonzero_mask + tf.to_float(
            tf.equal(records_column, model_keys.PADDING_ID)) * (-1000000.0)
        tf.summary.histogram('nonzero_mask', nonzero_mask)

        # [batch, 1]
        nonzeros = tf.count_nonzero(records_column, 1, keepdims=True)
        nonzeros = tf.maximum(nonzeros, 1)  # avoid divide zero
        tf.summary.scalar('nonzeros_sum', tf.reduce_sum(nonzeros))

        if not training and opts.normalize_embeddings:
            embeds = tf.nn.embedding_lookup(
                embeddings, tf.cast(records_column, tf.int32))
        else:
            embeds = mask_padding_embedding_lookup(
                embeddings, embedding_dim, records_column,
                model_keys.PADDING_ID)

        embeds_concat = []
        if opts.add_average_pooling:
            embeds_average = average_pooling(embeds, nonzeros)
            embeds_concat.append(embeds_average)
        if opts.add_max_pooling:
            embeds_max = max_pooling(embeds, nonzeros)
            embeds_concat.append(embeds_max)
        if opts.add_min_pooling:
            embeds_min = min_pooling(embeds, nonzeros)
            embeds_concat.append(embeds_min)
        if opts.add_hierarchical_pooling:
            embeds_hier = hierarchical_pooling(embeds, nonzeros, opts)
            embeds_concat.append(embeds_hier)
        if opts.add_attention_layer:
            embeds_att = attention_layer(embeds, nonzero_mask, opts)
            embeds_concat.append(embeds_att)
        embeds_pool = tf.concat(embeds_concat, 1)

        if use_batch_normalization:
            embeds_pool = batch_normalization(
                embeds_pool, training, "bn_input")

        if opts.use_user_features:
            tf.logging.info("Use user features")
            user_features = tf.feature_column.input_layer(
                features, params['user_features_columns'])
            concat_features = [embeds_pool, user_features]
            input_layer = tf.concat(concat_features, axis=1,
                                    name='concat_user_features')
        else:
            input_layer = embeds_pool

    return input_layer
//  [python]
def create_input_layer(mode, features, params, embeddings):
    """Create input layer."""

    opts = params['opts']
    embedding_dim = opts.embedding_dim
    use_batch_normalization = opts.use_batch_normalization
    training = (mode == tf.estimator.ModeKeys.TRAIN)

    with tf.name_scope("input_layer"):
        if training:
            records_column = params['records_column']
        else:
            records_column = params['predict_records_column']

        records_column = tf.feature_column.input_layer(
            features, records_column)

        nonzero_mask = tf.cast(
            tf.not_equal(records_column, model_keys.PADDING_ID), tf.float32)
        tf.summary.histogram('nonzero_cnt', nonzero_mask)
        nonzero_mask = nonzero_mask + tf.cast(
            tf.equal(records_column, model_keys.PADDING_ID), tf.float32) * (-1000000.0)
        tf.summary.histogram('nonzero_mask', nonzero_mask)

        # [batch, 1]
        nonzeros = tf.count_nonzero(records_column, 1, keepdims=True)
        nonzeros = tf.maximum(nonzeros, 1)  # avoid divide zero
        tf.summary.scalar('nonzeros_sum', tf.reduce_sum(nonzeros))

        if not training and opts.normalize_embeddings:
            embeds = tf.nn.embedding_lookup(
                embeddings, tf.cast(records_column, tf.int32))
        else:
            embeds = mask_padding_embedding_lookup(
                embeddings, embedding_dim, records_column,
                model_keys.PADDING_ID)

        embeds_concat = []
        if opts.add_average_pooling:
            embeds_average = average_pooling(embeds, nonzeros)
            embeds_concat.append(embeds_average)
        if opts.add_max_pooling:
            embeds_max = max_pooling(embeds, nonzeros)
            embeds_concat.append(embeds_max)
        if opts.add_min_pooling:
            embeds_min = min_pooling(embeds, nonzeros)
            embeds_concat.append(embeds_min)
        if opts.add_hierarchical_pooling:
            embeds_hier = hierarchical_pooling(embeds, nonzeros, opts)
            embeds_concat.append(embeds_hier)
        if opts.add_attention_layer:
            embeds_att = attention_layer(embeds, nonzero_mask, opts)
            embeds_concat.append(embeds_att)
        embeds_pool = tf.concat(embeds_concat, 1)

        if use_batch_normalization:
            embeds_pool = batch_normalization(
                embeds_pool, training, "bn_input")

        if opts.use_user_features:
            tf.logging.info("Use user features")
            user_features = tf.feature_column.input_layer(
                features, params['user_features_columns'])
            concat_features = [embeds_pool, user_features]
            input_layer = tf.concat(concat_features, axis=1,
                                    name='concat_user_features')
        else:
            input_layer = embeds_pool

    return input_layer
//  [python]
def log_prob(self, data, num_samples=1):
    """Computes log probability lower bound."""
    tiled_data = tf.tile(data[None],
                         [num_samples, 1] + [1] * len(self.data_dim))
    tiled_data_flat = tf.reshape(tiled_data, [-1] + self.data_dim)

    # Construct approximate posterior and sample z.
    q_z2_given_x = self.q_z2_given_x(data)
    z2 = q_z2_given_x.sample(sample_shape=[num_samples])  # [S, B, ...]
    z2_flat = tf.reshape(z2, [-1, self.latent_dim])  # [S*B, ...]
    q_z1_given_x_z2 = self.q_z1_given_x_z2(tiled_data_flat, z2_flat)
    z1 = q_z1_given_x_z2.sample()

    log_q = self.kl_weight * (
        tf.reshape(q_z2_given_x.log_prob(z2), [-1]) +
        q_z1_given_x_z2.log_prob(z1))
    log_p = (
        self.kl_weight *
        (self.proposal.log_prob(z2_flat) + self.p_z1_z2(z2_flat).log_prob(z1)) +
        self.p_x_given_z1_z2(z1, z2_flat).log_prob(tiled_data_flat))

    elbo = tf.reduce_logsumexp(tf.reshape(log_p - log_q, [num_samples, -1]),
                               axis=0) - tf.log(tf.to_float(num_samples))
    return elbo
//  [python]
def log_prob(self, data, num_samples=1):
    """Computes log probability lower bound."""
    tiled_data = tf.tile(data[None],
                         [num_samples, 1] + [1] * len(self.data_dim))
    tiled_data_flat = tf.reshape(tiled_data, [-1] + self.data_dim)

    # Construct approximate posterior and sample z.
    q_z2_given_x = self.q_z2_given_x(data)
    z2 = q_z2_given_x.sample(sample_shape=[num_samples])  # [S, B, ...]
    z2_flat = tf.reshape(z2, [-1, self.latent_dim])  # [S*B, ...]
    q_z1_given_x_z2 = self.q_z1_given_x_z2(tiled_data_flat, z2_flat)
    z1 = q_z1_given_x_z2.sample()

    log_q = self.kl_weight * (
        tf.reshape(q_z2_given_x.log_prob(z2), [-1]) +
        q_z1_given_x_z2.log_prob(z1))
    log_p = (
        self.kl_weight *
        (self.proposal.log_prob(z2_flat) + self.p_z1_z2(z2_flat).log_prob(z1)) +
        self.p_x_given_z1_z2(z1, z2_flat).log_prob(tiled_data_flat))

    elbo = tf.reduce_logsumexp(tf.reshape(log_p - log_q, [num_samples, -1]),
                               axis=0) - tf.log(tf.cast(num_samples, tf.float32))
    return elbo
//  [python]
def _torso(self, input_):
    """Processing of all the visual and language inputs to the LSTM core."""

    # Extract the inputs
    last_action, env_output = input_
    last_reward, _, _, observation = env_output
    if type(observation) == list:
      frame = observation[self._idx_frame]
    else:
      frame = observation

    # Convert to image to floats and normalise.
    frame = tf.to_float(frame)
    frame /= 255

    # Feed image through convnet.
    with tf.variable_scope('convnet'):
      conv_out = frame
      for i, (num_ch, num_blocks) in enumerate([(16, 2), (32, 2), (32, 2)]):
        # Downscale.
        conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
        conv_out = tf.nn.pool(
            conv_out,
            window_shape=[3, 3],
            pooling_type='MAX',
            padding='SAME',
            strides=[2, 2])
        # Residual block(s).
        for j in range(num_blocks):
          with tf.variable_scope('residual_%d_%d' % (i, j)):
            block_input = conv_out
            conv_out = tf.nn.relu(conv_out)
            conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
            conv_out = tf.nn.relu(conv_out)
            conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
            conv_out += block_input
    # Fully connected layer.
    conv_out = tf.nn.relu(conv_out)
    conv_out = snt.BatchFlatten()(conv_out)
    conv_out = snt.Linear(256)(conv_out)
    conv_out = tf.nn.relu(conv_out)

    # Concatenate outputs of the visual and instruction pathways.
    if self._feed_action_and_reward:
      # Append clipped last reward and one hot last action.
      tf.logging.info('Append last reward clipped to: %f', self._max_reward)
      clipped_last_reward = tf.expand_dims(
          tf.clip_by_value(last_reward, -self._max_reward, self._max_reward),
          -1)
      tf.logging.info('Append last action (one-hot of %d)', self._num_actions)
      one_hot_last_action = tf.one_hot(last_action, self._num_actions)
      core_input = tf.concat(
          [conv_out, clipped_last_reward, one_hot_last_action],
          axis=1)
    else:
      core_input = conv_out
    return core_input
//  [python]
def _torso(self, input_):
    """Processing of all the visual and language inputs to the LSTM core."""

    # Extract the inputs
    last_action, env_output = input_
    last_reward, _, _, observation = env_output
    if type(observation) == list:
      frame = observation[self._idx_frame]
    else:
      frame = observation

    # Convert to image to floats and normalise.
    frame = tf.cast(frame, tf.float32)
    frame /= 255

    # Feed image through convnet.
    with tf.variable_scope('convnet'):
      conv_out = frame
      for i, (num_ch, num_blocks) in enumerate([(16, 2), (32, 2), (32, 2)]):
        # Downscale.
        conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
        conv_out = tf.nn.pool(
            conv_out,
            window_shape=[3, 3],
            pooling_type='MAX',
            padding='SAME',
            strides=[2, 2])
        # Residual block(s).
        for j in range(num_blocks):
          with tf.variable_scope('residual_%d_%d' % (i, j)):
            block_input = conv_out
            conv_out = tf.nn.relu(conv_out)
            conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
            conv_out = tf.nn.relu(conv_out)
            conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
            conv_out += block_input
    # Fully connected layer.
    conv_out = tf.nn.relu(conv_out)
    conv_out = snt.BatchFlatten()(conv_out)
    conv_out = snt.Linear(256)(conv_out)
    conv_out = tf.nn.relu(conv_out)

    # Concatenate outputs of the visual and instruction pathways.
    if self._feed_action_and_reward:
      # Append clipped last reward and one hot last action.
      tf.logging.info('Append last reward clipped to: %f', self._max_reward)
      clipped_last_reward = tf.expand_dims(
          tf.clip_by_value(last_reward, -self._max_reward, self._max_reward),
          -1)
      tf.logging.info('Append last action (one-hot of %d)', self._num_actions)
      one_hot_last_action = tf.one_hot(last_action, self._num_actions)
      core_input = tf.concat(
          [conv_out, clipped_last_reward, one_hot_last_action],
          axis=1)
    else:
      core_input = conv_out
    return core_input
//  [python]
def body(self, features):
    """Body of the model.

    Args:
      features: a dictionary with the tensors.

    Returns:
      A pair (predictions, losses) where predictions is the generated image
      and losses is a dictionary of losses (that get added for the final loss).
    """
    features["targets"] = features["inputs"]
    is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN

    # Input images.
    inputs = tf.to_float(features["targets_raw"])

    # Noise vector.
    z = tf.random_uniform([self.hparams.batch_size,
                           self.hparams.bottleneck_bits],
                          minval=-1, maxval=1, name="z")

    # Generator output: fake images.
    out_shape = common_layers.shape_list(inputs)[1:4]
    g = self.generator(z, is_training, out_shape)

    losses = self.losses(inputs, g)  # pylint: disable=not-callable

    summary_g_image = tf.reshape(
        g[0, :], [1] + common_layers.shape_list(inputs)[1:])
    tf.summary.image("generated", summary_g_image, max_outputs=1)

    if is_training:  # Returns an dummy output and the losses dictionary.
      return tf.zeros_like(inputs), losses
    return tf.reshape(g, tf.shape(inputs)), losses
//  [python]
def body(self, features):
    """Body of the model.

    Args:
      features: a dictionary with the tensors.

    Returns:
      A pair (predictions, losses) where predictions is the generated image
      and losses is a dictionary of losses (that get added for the final loss).
    """
    features["targets"] = features["inputs"]
    is_training = self.hparams.mode == tf.estimator.ModeKeys.TRAIN

    # Input images.
    inputs = tf.cast(features["targets_raw"], tf.float32)

    # Noise vector.
    z = tf.random.uniform([self.hparams.batch_size,
                           self.hparams.bottleneck_bits],
                          minval=-1, maxval=1, name="z")

    # Generator output: fake images.
    out_shape = common_layers.shape_list(inputs)[1:4]
    g = self.generator(z, is_training, out_shape)

    losses = self.losses(inputs, g)  # pylint: disable=not-callable

    summary_g_image = tf.reshape(
        g[0, :], [1] + common_layers.shape_list(inputs)[1:])
    tf.summary.image("generated", summary_g_image, max_outputs=1)

    if is_training:  # Returns an dummy output and the losses dictionary.
      return tf.zeros_like(inputs), losses
    return tf.reshape(g, tf.shape(inputs)), losses
//  [python]
def noise_from_step_num():
    """Quantization noise equal to (phi * (step_num + 1)) mod 1.0.

    Not using random_uniform here due to a problem on TPU in that random seeds
    are not respected, which may cause the parameters on different replicas
    to go out-of-sync.

    Returns:
      a float32 scalar
    """
    step = tf.to_int32(tf.train.get_or_create_global_step()) + 1
    phi = ((5 ** 0.5) - 1) / 2
    # Naive computation tf.mod(phi * step, 1.0) in float32 would be disastrous
    # due to loss of precision when the step number gets large.
    # Computation in doubles does not work on TPU, so we use this complicated
    # alternative computation which does not suffer from these roundoff errors.
    ret = 0.0
    for i in range(30):
        ret += (((phi * (2 ** i)) % 1.0)  # double-precision computation in python
                * tf.to_float(tf.mod(step // (2 ** i), 2)))
    return tf.mod(ret, 1.0)
//  [python]
def noise_from_step_num():
    """Quantization noise equal to (phi * (step_num + 1)) mod 1.0.

    Not using random_uniform here due to a problem on TPU in that random seeds
    are not respected, which may cause the parameters on different replicas
    to go out-of-sync.

    Returns:
      a float32 scalar
    """
    step = tf.cast(tf.train.get_or_create_global_step(), tf.int32) + 1
    phi = ((5 ** 0.5) - 1) / 2
    # Naive computation tf.mod(phi * step, 1.0) in float32 would be disastrous
    # due to loss of precision when the step number gets large.
    # Computation in doubles does not work on TPU, so we use this complicated
    # alternative computation which does not suffer from these roundoff errors.
    ret = 0.0
    for i in range(30):
        ret += (((phi * (2 ** i)) % 1.0)  # double-precision computation in python
                * tf.to_float(tf.mod(step // (2 ** i), 2)))
    return tf.mod(ret, 1.0)
//  [python]
def _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2):
    """Round-off x to cand1 or to cand2 in an unbiased way.

    Cand1 and cand2 are the same shape as x.
    For every element of x, the corresponding elements of cand1 and cand2 should
    be the two closest bfloat16 values to x.  Order does not matter.
    cand1 and cand2 must differ from each other.

    Args:
      x: A float32 Tensor.
      noise: A Tensor broadcastable to the shape of x containing
      random uniform values in [0.0, 1.0].
      cand1: A bfloat16 Tensor the same shape as x.
      cand2: A bfloat16 Tensor the same shape as x.

    Returns:
      A bfloat16 Tensor.
    """
    cand1_f = tf.to_float(cand1)
    cand2_f = tf.to_float(cand2)
    step_size = cand2_f - cand1_f
    fpart = (x - cand1_f) / step_size
    ret = tf.where(tf.greater(fpart, noise), cand2, cand1)
    return ret
//  [python]
def _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2):
    """Round-off x to cand1 or to cand2 in an unbiased way.

    Cand1 and cand2 are the same shape as x.
    For every element of x, the corresponding elements of cand1 and cand2 should
    be the two closest bfloat16 values to x.  Order does not matter.
    cand1 and cand2 must differ from each other.

    Args:
      x: A float32 Tensor.
      noise: A Tensor broadcastable to the shape of x containing
      random uniform values in [0.0, 1.0].
      cand1: A bfloat16 Tensor the same shape as x.
      cand2: A bfloat16 Tensor the same shape as x.

    Returns:
      A bfloat16 Tensor.
    """
    cand1_f = tf.cast(cand1, tf.float32)
    cand2_f = tf.cast(cand2, tf.float32)
    step_size = cand2_f - cand1_f
    fpart = (x - cand1_f) / step_size
    ret = tf.where(tf.greater(fpart, noise), cand2, cand1)
    return ret
//  [python]
def encode(self, x, noise):
        x = tf.to_float(x)
        # we can't use tf.pow(..., 8.0) because of a high-error approximation
        # on TPU.  Instead we square three times.
        x = tf.sign(x) * tf.square(tf.square(tf.square(tf.abs(x) * 128.0)))
        x = _to_bfloat16_unbiased(x, noise)
        return x
//  [python]
def encode(self, x, noise):
        x = tf.cast(x, tf.float32)
        # we can't use tf.pow(..., 8.0) because of a high-error approximation
        # on TPU.  Instead we square three times.
        x = tf.sign(x) * tf.square(tf.square(tf.square(tf.abs(x) * 128.0)))
        x = _to_bfloat16_unbiased(x, noise)
        return x
//  [python]
def decode(self, x):
        x = tf.to_float(x)
        # we can't use tf.pow(..., 0.125) because of a high-error approximation
        # on TPU.  Instead we sqrt three times.
        return tf.sign(x) * (tf.sqrt(tf.sqrt(tf.sqrt(tf.abs(x)))) / 128.0)
//  [python]
def decode(self, x):
        x = tf.cast(x, tf.float32)
        # we can't use tf.pow(..., 0.125) because of a high-error approximation
        # on TPU.  Instead we sqrt three times.
        return tf.sign(x) * (tf.sqrt(tf.sqrt(tf.sqrt(tf.abs(x)))) / 128.0)
//  [python]
def _build_train_op(self):
    """Builds a training op.

    Returns:
      train_op: An op performing one step of training.
    """
    target_distribution = tf.stop_gradient(self._build_target_distribution())

    # size of indices: batch_size x 1.
    indices = tf.range(tf.shape(self._replay_net_outputs.logits)[0])[:, None]
    # size of reshaped_actions: batch_size x 2.
    reshaped_actions = tf.concat([indices, self._replay.actions[:, None]], 1)
    # For each element of the batch, fetch the logits for its selected action.
    chosen_action_logits = tf.gather_nd(self._replay_net_outputs.logits,
                                        reshaped_actions)

    bellman_errors = (target_distribution[:, None, :] -
                      chosen_action_logits[:, :, None])  # Input `u' of Eq. 9.
    huber_loss = (  # Eq. 9 of paper.
        tf.to_float(tf.abs(bellman_errors) <= self.kappa) *
        0.5 * bellman_errors ** 2 +
        tf.to_float(tf.abs(bellman_errors) > self.kappa) *
        self.kappa * (tf.abs(bellman_errors) - 0.5 * self.kappa))

    tau_hat = ((tf.range(self._num_atoms, dtype=tf.float32) + 0.5) /
               self._num_atoms)  # Quantile midpoints.  See Lemma 2 of paper.

    quantile_huber_loss = (  # Eq. 10 of paper.
        tf.abs(tau_hat[None, :, None] - tf.to_float(bellman_errors < 0)) *
        huber_loss)

    # Sum over tau dimension, average over target value dimension.
    loss = tf.reduce_sum(tf.reduce_mean(quantile_huber_loss, 2), 1)

    if self._replay_scheme == 'prioritized':
      target_priorities = self._replay.tf_get_priority(self._replay.indices)
      # The original prioritized experience replay uses a linear exponent
      # schedule 0.4 -> 1.0. Comparing the schedule to a fixed exponent of 0.5
      # on 5 games (Asterix, Pong, Q*Bert, Seaquest, Space Invaders) suggested
      # a fixed exponent actually performs better, except on Pong.
      loss_weights = 1.0 / tf.sqrt(target_priorities + 1e-10)
      loss_weights /= tf.reduce_max(loss_weights)

      # Rainbow and prioritized replay are parametrized by an exponent alpha,
      # but in both cases it is set to 0.5 - for simplicity's sake we leave it
      # as is here, using the more direct tf.sqrt(). Taking the square root
      # "makes sense", as we are dealing with a squared loss.
      # Add a small nonzero value to the loss to avoid 0 priority items. While
      # technically this may be okay, setting all items to 0 priority will cause
      # troubles, and also result in 1.0 / 0.0 = NaN correction terms.
      update_priorities_op = self._replay.tf_set_priority(
          self._replay.indices, tf.sqrt(loss + 1e-10))

      # Weight loss by inverse priorities.
      loss = loss_weights * loss
    else:
      update_priorities_op = tf.no_op()

    with tf.control_dependencies([update_priorities_op]):
      if self.summary_writer is not None:
        with tf.variable_scope('Losses'):
          tf.summary.scalar('QuantileLoss', tf.reduce_mean(loss))
      return self.optimizer.minimize(tf.reduce_mean(loss)), loss
//  [python]
def _build_train_op(self):
    """Builds a training op.

    Returns:
      train_op: An op performing one step of training.
    """
    target_distribution = tf.stop_gradient(self._build_target_distribution())

    # size of indices: batch_size x 1.
    indices = tf.range(tf.shape(self._replay_net_outputs.logits)[0])[:, None]
    # size of reshaped_actions: batch_size x 2.
    reshaped_actions = tf.concat([indices, self._replay.actions[:, None]], 1)
    # For each element of the batch, fetch the logits for its selected action.
    chosen_action_logits = tf.gather_nd(self._replay_net_outputs.logits,
                                        reshaped_actions)

    bellman_errors = (target_distribution[:, None, :] -
                      chosen_action_logits[:, :, None])  # Input `u' of Eq. 9.
    huber_loss = (  # Eq. 9 of paper.
        tf.cast(tf.abs(bellman_errors) <= self.kappa, tf.float32) *
        0.5 * bellman_errors ** 2 +
        tf.cast(tf.abs(bellman_errors) > self.kappa, tf.float32) *
        self.kappa * (tf.abs(bellman_errors) - 0.5 * self.kappa))

    tau_hat = ((tf.range(self._num_atoms, dtype=tf.float32) + 0.5) /
               self._num_atoms)  # Quantile midpoints.  See Lemma 2 of paper.

    quantile_huber_loss = (  # Eq. 10 of paper.
        tf.abs(tau_hat[None, :, None] - tf.cast(bellman_errors < 0, tf.float32)) *
        huber_loss)

    # Sum over tau dimension, average over target value dimension.
    loss = tf.reduce_sum(tf.reduce_mean(quantile_huber_loss, 2), 1)

    if self._replay_scheme == 'prioritized':
      target_priorities = self._replay.tf_get_priority(self._replay.indices)
      # The original prioritized experience replay uses a linear exponent
      # schedule 0.4 -> 1.0. Comparing the schedule to a fixed exponent of 0.5
      # on 5 games (Asterix, Pong, Q*Bert, Seaquest, Space Invaders) suggested
      # a fixed exponent actually performs better, except on Pong.
      loss_weights = 1.0 / tf.sqrt(target_priorities + 1e-10)
      loss_weights /= tf.reduce_max(loss_weights)

      # Rainbow and prioritized replay are parametrized by an exponent alpha,
      # but in both cases it is set to 0.5 - for simplicity's sake we leave it
      # as is here, using the more direct tf.sqrt(). Taking the square root
      # "makes sense", as we are dealing with a squared loss.
      # Add a small nonzero value to the loss to avoid 0 priority items. While
      # technically this may be okay, setting all items to 0 priority will cause
      # troubles, and also result in 1.0 / 0.0 = NaN correction terms.
      update_priorities_op = self._replay.tf_set_priority(
          self._replay.indices, tf.sqrt(loss + 1e-10))

      # Weight loss by inverse priorities.
      loss = loss_weights * loss
    else:
      update_priorities_op = tf.no_op()

    with tf.control_dependencies([update_priorities_op]):
      if self.summary_writer is not None:
        with tf.variable_scope('Losses'):
          tf.summary.scalar('QuantileLoss', tf.reduce_mean(loss))
      return self.optimizer.minimize(tf.reduce_mean(loss)), loss
//  [python]
def _to_bfloat16_unbiased(x, noise):
    """Convert a float32 to a bfloat16 using randomized roundoff.

    Args:
      x: A float32 Tensor.
      noise: a float32 Tensor with values in [0, 1), broadcastable to tf.shape(x)
    Returns:
      A float32 Tensor.
    """
    x_sign = tf.sign(x)
    # Make sure x is positive.  If it is zero, the two candidates are
    # identical.
    x = x * x_sign + 1e-30
    cand1 = tf.to_bfloat16(x)
    cand1_f = tf.to_float(cand1)
    # This relies on the fact that for a positive bfloat16 b,
    # b * 1.005 gives you the next higher bfloat16 and b*0.995 gives you the
    # next lower one. Both 1.005 and 0.995 are ballpark estimation.
    cand2 = tf.to_bfloat16(
        tf.where(tf.greater(x, cand1_f), cand1_f * 1.005, cand1_f * 0.995))
    ret = _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2)
    return ret * tf.to_bfloat16(x_sign)
//  [python]
def _to_bfloat16_unbiased(x, noise):
    """Convert a float32 to a bfloat16 using randomized roundoff.

    Args:
      x: A float32 Tensor.
      noise: a float32 Tensor with values in [0, 1), broadcastable to tf.shape(x)
    Returns:
      A float32 Tensor.
    """
    x_sign = tf.sign(x)
    # Make sure x is positive.  If it is zero, the two candidates are
    # identical.
    x = x * x_sign + 1e-30
    cand1 = tf.cast(x, tf.bfloat16)
    cand1_f = tf.cast(cand1, tf.float32)
    # This relies on the fact that for a positive bfloat16 b,
    # b * 1.005 gives you the next higher bfloat16 and b*0.995 gives you the
    # next lower one. Both 1.005 and 0.995 are ballpark estimation.
    cand2 = tf.cast(
        tf.where(tf.greater(x, cand1_f), cand1_f * 1.005, cand1_f * 0.995), tf.bfloat16)
    ret = _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2)
    return ret * tf.cast(x_sign, tf.bfloat16)
//  [python]
def _shard_features(self, features):  # pylint: disable=missing-docstring
    sharded_features = {}
    for k, v in sorted(six.iteritems(features)):
      v = tf.convert_to_tensor(v)
      v_shape = common_layers.shape_list(v)
      if not v_shape:
        v = tf.expand_dims(v, axis=-1)
        v_shape = [1]
      if v_shape == [1]:
        v = tf.tile(v, tf.to_int32([self._num_datashards]))
      sharded_features[k] = self._data_parallelism(
          tf.identity, tf.split(v, self._num_datashards, 0))
    return sharded_features
//  [python]
def _shard_features(self, features):  # pylint: disable=missing-docstring
    sharded_features = {}
    for k, v in sorted(six.iteritems(features)):
      v = tf.convert_to_tensor(v)
      v_shape = common_layers.shape_list(v)
      if not v_shape:
        v = tf.expand_dims(v, axis=-1)
        v_shape = [1]
      if v_shape == [1]:
        v = tf.tile(v, tf.cast([self._num_datashards], tf.int32))
      sharded_features[k] = self._data_parallelism(
          tf.identity, tf.split(v, self._num_datashards, 0))
    return sharded_features
//  [python]
def train(self, features, labels):
        logits = self.__call__(features, is_training=True)  # batch_size, mask_len
        labels = tf.reshape(labels, [-1])  # batch_size, mask_len
        label_weights = tf.to_float(tf.not_equal(labels, 0))

        per_example_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=logits)
        numerator = tf.reduce_sum(label_weights * per_example_loss)
        denominator = tf.reduce_sum(label_weights) + 1e-5
        loss = numerator / denominator + tf.losses.get_regularization_loss()
        train_op = self.trainOp(loss)

        with tf.variable_scope("Sequential/TRAIN"):
            _, loss_op = tf.metrics.mean(loss, name='loss')

        loss_vars = tf.get_collection(tf.GraphKeys.LOCAL_VARIABLES, scope="main/Sequential/TRAIN")
        assert len(loss_vars) > 0, "(train)metric local variables should not be None."
        loss_init_op = tf.variables_initializer(loss_vars)

        return train_op, loss_op, loss_init_op
//  [python]
def train(self, features, labels):
        logits = self.__call__(features, is_training=True)  # batch_size, mask_len
        labels = tf.reshape(labels, [-1])  # batch_size, mask_len
        label_weights = tf.cast(tf.not_equal(labels, 0), tf.float32)

        per_example_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=logits)
        numerator = tf.reduce_sum(label_weights * per_example_loss)
        denominator = tf.reduce_sum(label_weights) + 1e-5
        loss = numerator / denominator + tf.losses.get_regularization_loss()
        train_op = self.trainOp(loss)

        with tf.variable_scope("Sequential/TRAIN"):
            _, loss_op = tf.metrics.mean(loss, name='loss')

        loss_vars = tf.get_collection(tf.GraphKeys.LOCAL_VARIABLES, scope="main/Sequential/TRAIN")
        assert len(loss_vars) > 0, "(train)metric local variables should not be None."
        loss_init_op = tf.variables_initializer(loss_vars)

        return train_op, loss_op, loss_init_op
//  [python]
def summarize_features(features, num_shards=1):
  """Generate summaries for features."""
  if not common_layers.should_generate_summaries():
    return

  with tf.name_scope("input_stats"):
    for (k, v) in sorted(six.iteritems(features)):
      if (isinstance(v, tf.Tensor) and (v.get_shape().ndims > 1) and
          (v.dtype != tf.string)):
        tf.summary.scalar("%s_batch" % k, tf.shape(v)[0] // num_shards)
        tf.summary.scalar("%s_length" % k, tf.shape(v)[1])
        nonpadding = tf.to_float(tf.not_equal(v, 0))
        nonpadding_tokens = tf.reduce_sum(nonpadding)
        tf.summary.scalar("%s_nonpadding_tokens" % k, nonpadding_tokens)
        tf.summary.scalar("%s_nonpadding_fraction" % k,
                          tf.reduce_mean(nonpadding))
//  [python]
def summarize_features(features, num_shards=1):
  """Generate summaries for features."""
  if not common_layers.should_generate_summaries():
    return

  with tf.name_scope("input_stats"):
    for (k, v) in sorted(six.iteritems(features)):
      if (isinstance(v, tf.Tensor) and (v.get_shape().ndims > 1) and
          (v.dtype != tf.string)):
        tf.summary.scalar("%s_batch" % k, tf.shape(v)[0] // num_shards)
        tf.summary.scalar("%s_length" % k, tf.shape(v)[1])
        nonpadding = tf.cast(tf.not_equal(v, 0), tf.float32)
        nonpadding_tokens = tf.reduce_sum(nonpadding)
        tf.summary.scalar("%s_nonpadding_tokens" % k, nonpadding_tokens)
        tf.summary.scalar("%s_nonpadding_fraction" % k,
                          tf.reduce_mean(nonpadding))
//  [python]
def _global_step(hparams):
  """Adjust global step if a multi-step optimizer is used."""
  step = tf.to_float(tf.train.get_or_create_global_step())
  multiplier = hparams.optimizer_multistep_accumulate_steps
  if not multiplier:
    return step

  tf.logging.info("Dividing global step by %d for multi-step optimizer."
                  % multiplier)
  return step / tf.to_float(multiplier)
//  [python]
def _global_step(hparams):
  """Adjust global step if a multi-step optimizer is used."""
  step = tf.cast(tf.train.get_or_create_global_step(), tf.float32)
  multiplier = hparams.optimizer_multistep_accumulate_steps
  if not multiplier:
    return step

  tf.logging.info("Dividing global step by %d for multi-step optimizer."
                  % multiplier)
  return step / tf.cast(multiplier, tf.float32)
//  [python]
def build_subnetwork(self,
                       features,
                       logits_dimension,
                       training,
                       iteration_step,
                       summary,
                       previous_ensemble=None):
    """See `adanet.subnetwork.Builder`."""

    # Prepare the input.
    assert len(self._feature_columns) == 1, "Got feature columns: {}".format(
        self._feature_columns)
    images = tf.to_float(features[self._feature_columns[0].name])
    self._name_scope = tf.get_default_graph().get_name_scope()

    seed = self._seed
    if seed is not None and previous_ensemble:
      # Deterministically change the seed for different iterations so that
      # subnetworks are not correlated.
      seed += len(previous_ensemble.weighted_subnetworks)

    arg_scope = nasnet.nasnet_cifar_arg_scope(weight_decay=self._weight_decay)

    with tf.contrib.slim.arg_scope(arg_scope):
      build_fn = nasnet.build_nasnet_cifar
      logits, end_points = build_fn(
          images,
          num_classes=logits_dimension,
          is_training=training,
          config=self._hparams)
    last_layer = end_points["global_pool"]

    subnetwork_shared_data = {
        _PREVIOUS_NUM_CELLS: tf.constant(self._num_cells),
        _PREVIOUS_CONV_FILTERS: tf.constant(self._num_conv_filters)
    }

    return adanet.Subnetwork(
        last_layer=last_layer,
        logits=logits,
        complexity=1,
        shared=subnetwork_shared_data)
//  [python]
def build_subnetwork(self,
                       features,
                       logits_dimension,
                       training,
                       iteration_step,
                       summary,
                       previous_ensemble=None):
    """See `adanet.subnetwork.Builder`."""

    # Prepare the input.
    assert len(self._feature_columns) == 1, "Got feature columns: {}".format(
        self._feature_columns)
    images = tf.cast(features[self._feature_columns[0].name], tf.float32)
    self._name_scope = tf.get_default_graph().get_name_scope()

    seed = self._seed
    if seed is not None and previous_ensemble:
      # Deterministically change the seed for different iterations so that
      # subnetworks are not correlated.
      seed += len(previous_ensemble.weighted_subnetworks)

    arg_scope = nasnet.nasnet_cifar_arg_scope(weight_decay=self._weight_decay)

    with tf.contrib.slim.arg_scope(arg_scope):
      build_fn = nasnet.build_nasnet_cifar
      logits, end_points = build_fn(
          images,
          num_classes=logits_dimension,
          is_training=training,
          config=self._hparams)
    last_layer = end_points["global_pool"]

    subnetwork_shared_data = {
        _PREVIOUS_NUM_CELLS: tf.constant(self._num_cells),
        _PREVIOUS_CONV_FILTERS: tf.constant(self._num_conv_filters)
    }

    return adanet.Subnetwork(
        last_layer=last_layer,
        logits=logits,
        complexity=1,
        shared=subnetwork_shared_data)
//  [python]
def legacy_learning_rate_schedule(hparams):
  """Backwards-compatible learning-rate schedule."""
  step_num = _global_step(hparams)
  warmup_steps = tf.to_float(hparams.learning_rate_warmup_steps)
  if hparams.learning_rate_decay_scheme == "noam":
    ret = 5000.0 * hparams.hidden_size**-0.5 * tf.minimum(
        (step_num + 1) * warmup_steps**-1.5, (step_num + 1)**-0.5)
  else:
    warmup_steps = hparams.learning_rate_warmup_steps
    warmup = _learning_rate_warmup(warmup_steps, hparams=hparams)
    decay = _learning_rate_decay(hparams, warmup_steps)
    ret = tf.where(step_num < warmup_steps, warmup, decay)
  optimizer_correction = 0.002 if "adam" in hparams.optimizer else 1.0
  tf.logging.info("Base learning rate: %f", hparams.learning_rate)
  return ret * optimizer_correction * hparams.learning_rate
//  [python]
def legacy_learning_rate_schedule(hparams):
  """Backwards-compatible learning-rate schedule."""
  step_num = _global_step(hparams)
  warmup_steps = tf.cast(hparams.learning_rate_warmup_steps, tf.float32)
  if hparams.learning_rate_decay_scheme == "noam":
    ret = 5000.0 * hparams.hidden_size**-0.5 * tf.minimum(
        (step_num + 1) * warmup_steps**-1.5, (step_num + 1)**-0.5)
  else:
    warmup_steps = hparams.learning_rate_warmup_steps
    warmup = _learning_rate_warmup(warmup_steps, hparams=hparams)
    decay = _learning_rate_decay(hparams, warmup_steps)
    ret = tf.where(step_num < warmup_steps, warmup, decay)
  optimizer_correction = 0.002 if "adam" in hparams.optimizer else 1.0
  tf.logging.info("Base learning rate: %f", hparams.learning_rate)
  return ret * optimizer_correction * hparams.learning_rate
//  [python]
def _learning_rate_warmup(warmup_steps, warmup_schedule="exp", hparams=None):
  """Learning rate warmup multiplier."""
  if not warmup_steps:
    return tf.constant(1.)

  tf.logging.info("Applying %s learning rate warmup for %d steps",
                  warmup_schedule, warmup_steps)

  warmup_steps = tf.to_float(warmup_steps)
  global_step = _global_step(hparams)

  if warmup_schedule == "exp":
    return tf.exp(tf.log(0.01) / warmup_steps)**(warmup_steps - global_step)
  else:
    assert warmup_schedule == "linear"
    start = tf.constant(0.35)
    return ((tf.constant(1.) - start) / warmup_steps) * global_step + start
//  [python]
def _learning_rate_warmup(warmup_steps, warmup_schedule="exp", hparams=None):
  """Learning rate warmup multiplier."""
  if not warmup_steps:
    return tf.constant(1.)

  tf.logging.info("Applying %s learning rate warmup for %d steps",
                  warmup_schedule, warmup_steps)

  warmup_steps = tf.cast(warmup_steps, tf.float32)
  global_step = _global_step(hparams)

  if warmup_schedule == "exp":
    return tf.exp(tf.log(0.01) / warmup_steps)**(warmup_steps - global_step)
  else:
    assert warmup_schedule == "linear"
    start = tf.constant(0.35)
    return ((tf.constant(1.) - start) / warmup_steps) * global_step + start
//  [python]
def create_host_call(model_dir):
  """Construct a host_call writing scalar summaries.

  Args:
    model_dir: String containing path to train

  Returns:
    (fn, args) Pair to be called by TPUEstimator as the host_call.
  """
  graph = tf.get_default_graph()
  summaries = graph.get_collection(tf.GraphKeys.SUMMARIES)
  gs_t = tf.reshape(tf.to_int32(tf.train.get_global_step()), [1])
  summary_kwargs = collections.OrderedDict()
  for t in summaries:
    # TODO(aidangomez): enable ImageSummary support when we have a faster method
    # see @shibow's comment in cl/202344570
    if t.op.type not in ["ScalarSummary"]:
      tf.logging.warn("Ignoring unsupported tf.Summary type %s" % t.op.type)
      continue

    name = t.op.name
    tensor = t.op.inputs[1]
    if t.op.type == "ScalarSummary":
      assert tensor.shape.is_compatible_with([])
      if tensor.dtype == tf.int64:
        tensor = tf.to_int32(tensor)
      summary_kwargs["ScalarSummary" + name] = tf.reshape(tensor, [1])
    elif t.op.type == "ImageSummary":
      # TODO(aidangomez): as we move to support more types, update
      # common_layers.tpu_safe_image_summary
      if tensor.dtype != tf.float32:
        tf.logging.warn(
            "Currently T2T on TPU only supports ImageSummary of "
            "tf.float32-type Tensors. Skipping Tensor "
            "%s with dtype %s..." % (tensor.name, tensor.dtype))
        continue
      # tensor = tf.to_float(tensor)
      summary_kwargs["ImageSummary" + name] = tensor
  # When no supported summaries are found, don't create host_call. Otherwise,
  # TPU outfeed queue would enqueue global_step while host_call doesn't dequeue
  # it, eventually causing hang.
  if not summary_kwargs:
    return None
  summary_kwargs["global_step"] = gs_t
  log_info("summary_kwargs %s" % str(summary_kwargs))

  def host_call_fn(**kwargs):
    """Training host call. Creates summaries for training metrics.

    Args:
      **kwargs: Dict of {str: Tensor} , with `Tensor` of shape `[batch]`. Must
        contain key "global_step" with value of current global_step Tensor.

    Returns:
      List of summary ops to run on the CPU host.
    """
    gs = tf.to_int64(kwargs.pop("global_step")[0])
    with contrib.summary().create_file_writer(model_dir).as_default():
      with contrib.summary().always_record_summaries():
        # We need to use tf.contrib.summary in order to feed the `step`.
        for name, value in sorted(six.iteritems(kwargs)):
          if name.startswith("ScalarSummary"):
            name = name[len("ScalarSummary"):]
            contrib.summary().scalar(
                name, tf.reduce_mean(tf.to_float(value)), step=gs)
          elif name.startswith("ImageSummary"):
            name = name[len("ImageSummary"):]
            contrib.summary().image(name, value, step=gs)

        return contrib.summary().all_summary_ops()

  return (host_call_fn, summary_kwargs)
//  [python]
def create_host_call(model_dir):
  """Construct a host_call writing scalar summaries.

  Args:
    model_dir: String containing path to train

  Returns:
    (fn, args) Pair to be called by TPUEstimator as the host_call.
  """
  graph = tf.get_default_graph()
  summaries = graph.get_collection(tf.GraphKeys.SUMMARIES)
  gs_t = tf.reshape(tf.cast(tf.train.get_global_step(), tf.int32), [1])
  summary_kwargs = collections.OrderedDict()
  for t in summaries:
    # TODO(aidangomez): enable ImageSummary support when we have a faster method
    # see @shibow's comment in cl/202344570
    if t.op.type not in ["ScalarSummary"]:
      tf.logging.warn("Ignoring unsupported tf.Summary type %s" % t.op.type)
      continue

    name = t.op.name
    tensor = t.op.inputs[1]
    if t.op.type == "ScalarSummary":
      assert tensor.shape.is_compatible_with([])
      if tensor.dtype == tf.int64:
        tensor = tf.cast(tensor, tf.int32)
      summary_kwargs["ScalarSummary" + name] = tf.reshape(tensor, [1])
    elif t.op.type == "ImageSummary":
      # TODO(aidangomez): as we move to support more types, update
      # common_layers.tpu_safe_image_summary
      if tensor.dtype != tf.float32:
        tf.logging.warn(
            "Currently T2T on TPU only supports ImageSummary of "
            "tf.float32-type Tensors. Skipping Tensor "
            "%s with dtype %s..." % (tensor.name, tensor.dtype))
        continue
      # tensor = tf.to_float(tensor)
      summary_kwargs["ImageSummary" + name] = tensor
  # When no supported summaries are found, don't create host_call. Otherwise,
  # TPU outfeed queue would enqueue global_step while host_call doesn't dequeue
  # it, eventually causing hang.
  if not summary_kwargs:
    return None
  summary_kwargs["global_step"] = gs_t
  log_info("summary_kwargs %s" % str(summary_kwargs))

  def host_call_fn(**kwargs):
    """Training host call. Creates summaries for training metrics.

    Args:
      **kwargs: Dict of {str: Tensor} , with `Tensor` of shape `[batch]`. Must
        contain key "global_step" with value of current global_step Tensor.

    Returns:
      List of summary ops to run on the CPU host.
    """
    gs = tf.to_int64(kwargs.pop("global_step")[0])
    with contrib.summary().create_file_writer(model_dir).as_default():
      with contrib.summary().always_record_summaries():
        # We need to use tf.contrib.summary in order to feed the `step`.
        for name, value in sorted(six.iteritems(kwargs)):
          if name.startswith("ScalarSummary"):
            name = name[len("ScalarSummary"):]
            contrib.summary().scalar(
                name, tf.reduce_mean(tf.to_float(value)), step=gs)
          elif name.startswith("ImageSummary"):
            name = name[len("ImageSummary"):]
            contrib.summary().image(name, value, step=gs)

        return contrib.summary().all_summary_ops()

  return (host_call_fn, summary_kwargs)