jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def isemhash_bottleneck(x,
                        bottleneck_bits,
                        bottleneck_noise,
                        discretize_warmup_steps,
                        mode,
                        isemhash_noise_dev=0.5,
                        isemhash_mix_prob=0.5):
  """Improved semantic hashing bottleneck."""
  with tf.variable_scope("isemhash_bottleneck"):
    x = tf.layers.dense(x, bottleneck_bits, name="dense")
    y = common_layers.saturating_sigmoid(x)
    if isemhash_noise_dev > 0 and mode == tf.estimator.ModeKeys.TRAIN:
      noise = tf.truncated_normal(
          common_layers.shape_list(x), mean=0.0, stddev=isemhash_noise_dev)
      y = common_layers.saturating_sigmoid(x + noise)
    d = tf.to_float(tf.less(0.5, y)) + y - tf.stop_gradient(y)
    d = 2.0 * d - 1.0  # Move from [0, 1] to [-1, 1].
    if mode == tf.estimator.ModeKeys.TRAIN:  # Flip some bits.
      noise = tf.random_uniform(common_layers.shape_list(x))
      noise = 2.0 * tf.to_float(tf.less(bottleneck_noise, noise)) - 1.0
      d *= noise
      d = common_layers.mix(
          d,
          2.0 * y - 1.0,
          discretize_warmup_steps,
          mode == tf.estimator.ModeKeys.TRAIN,
          max_prob=isemhash_mix_prob)
    return d, 0.0
//  [python]
def isemhash_bottleneck(x,
                        bottleneck_bits,
                        bottleneck_noise,
                        discretize_warmup_steps,
                        mode,
                        isemhash_noise_dev=0.5,
                        isemhash_mix_prob=0.5):
  """Improved semantic hashing bottleneck."""
  with tf.variable_scope("isemhash_bottleneck"):
    x = tf.layers.dense(x, bottleneck_bits, name="dense")
    y = common_layers.saturating_sigmoid(x)
    if isemhash_noise_dev > 0 and mode == tf.estimator.ModeKeys.TRAIN:
      noise = tf.truncated_normal(
          common_layers.shape_list(x), mean=0.0, stddev=isemhash_noise_dev)
      y = common_layers.saturating_sigmoid(x + noise)
    d = tf.cast(tf.less(0.5, y), tf.float32) + y - tf.stop_gradient(y)
    d = 2.0 * d - 1.0  # Move from [0, 1] to [-1, 1].
    if mode == tf.estimator.ModeKeys.TRAIN:  # Flip some bits.
      noise = tf.random_uniform(common_layers.shape_list(x))
      noise = 2.0 * tf.cast(tf.less(bottleneck_noise, noise), tf.float32) - 1.0
      d *= noise
      d = common_layers.mix(
          d,
          2.0 * y - 1.0,
          discretize_warmup_steps,
          mode == tf.estimator.ModeKeys.TRAIN,
          max_prob=isemhash_mix_prob)
    return d, 0.0
//  [python]
def gram_matrix(feature_maps):
  """Computes the Gram matrix for a set of feature maps."""
  batch_size, height, width, channels = tf.unstack(tf.shape(feature_maps))
  denominator = tf.to_float(height * width)
  feature_maps = tf.reshape(
      feature_maps, tf.stack([batch_size, height * width, channels]))
  matrix = tf.matmul(feature_maps, feature_maps, adjoint_a=True)
  return matrix / denominator
//  [python]
def gram_matrix(feature_maps):
  """Computes the Gram matrix for a set of feature maps."""
  batch_size, height, width, channels = tf.unstack(tf.shape(feature_maps))
  denominator = tf.cast(height * width, tf.float32)
  feature_maps = tf.reshape(
      feature_maps, tf.stack([batch_size, height * width, channels]))
  matrix = tf.matmul(feature_maps, feature_maps, adjoint_a=True)
  return matrix / denominator
//  [python]
def make_graph(features, labels, num_hidden=8):
  hidden_weights = tf.Variable(tf.truncated_normal(
      [2, num_hidden],
      stddev=1/math.sqrt(2)
  ))
  tf.summary.image('hidden_weights', tf.expand_dims([hidden_weights], -1))

  # Shape [4, num_hidden]
  hidden_activations = tf.nn.relu(tf.matmul(features, hidden_weights))

  output_weights = tf.Variable(tf.truncated_normal(
      [num_hidden, 1],
      stddev=1/math.sqrt(num_hidden)
  ))

  # Shape [4, 1]
  logits = tf.matmul(hidden_activations, output_weights)

  # Shape [4]
  predictions = tf.sigmoid(tf.squeeze(logits))
  loss = tf.reduce_mean(tf.square(predictions - tf.to_float(labels)))
  tf.summary.scalar('loss', loss)

  accuracy, update_acc = tf.contrib.metrics.streaming_accuracy(
      predictions > 0.5, labels)
  tf.summary.scalar('accuracy', accuracy)

  gs = tf.Variable(0, trainable=False)
  optimizer = tf.train.GradientDescentOptimizer(0.2)

  grads_and_vars = optimizer.compute_gradients(loss)

  gradients = list(zip(grads_and_vars))[0]
  tf.summary.histogram('gradients', gradients)

  train_op = optimizer.apply_gradients(grads_and_vars, global_step=gs)

  return train_op, loss, gs, update_acc
//  [python]
def make_graph(features, labels, num_hidden=8):
  hidden_weights = tf.Variable(tf.truncated_normal(
      [2, num_hidden],
      stddev=1/math.sqrt(2)
  ))
  tf.summary.image('hidden_weights', tf.expand_dims([hidden_weights], -1))

  # Shape [4, num_hidden]
  hidden_activations = tf.nn.relu(tf.matmul(features, hidden_weights))

  output_weights = tf.Variable(tf.truncated_normal(
      [num_hidden, 1],
      stddev=1/math.sqrt(num_hidden)
  ))

  # Shape [4, 1]
  logits = tf.matmul(hidden_activations, output_weights)

  # Shape [4]
  predictions = tf.sigmoid(tf.squeeze(logits))
  loss = tf.reduce_mean(tf.square(predictions - tf.cast(labels, tf.float32)))
  tf.summary.scalar('loss', loss)

  accuracy, update_acc = tf.contrib.metrics.streaming_accuracy(
      predictions > 0.5, labels)
  tf.summary.scalar('accuracy', accuracy)

  gs = tf.Variable(0, trainable=False)
  optimizer = tf.train.GradientDescentOptimizer(0.2)

  grads_and_vars = optimizer.compute_gradients(loss)

  gradients = list(zip(grads_and_vars))[0]
  tf.summary.histogram('gradients', gradients)

  train_op = optimizer.apply_gradients(grads_and_vars, global_step=gs)

  return train_op, loss, gs, update_acc
//  [python]
def compute_forces(overlap, pivots, num_inputs, overlap_ratio, scope=None):
    """Computes accumulated forces based on pairwise hypersphere intersection.

    A (repulsive or attractive) force between two hyperspheres centered in
    corresponding pivots is computed based on the difference between their
    overlap versus desired overlap amount.

    An accumulated force is the average of all forces computed from all the
    other pivots.

    Args:
      overlap: A float32 `Tensor` of shape [num_pivots, num_pivots].
      pivots: A float32 `Tensor` of shape [num_pivots, input_dims].
      num_inputs: Number of training points.
      overlap_ratio: Constant to satisfy independence condition.
      scope: Optional variable scope.

    Returns:
      A float32 `Tensor` of shape [num_pivots, input_dims].
    """
    with tf.variable_scope(scope, 'ComputeForces', [overlap, pivots]):
        num_pivots = tf.shape(pivots)[0]
        indices = tf.reshape(tf.range(num_pivots), [-1, 1])
        indices = tf.reshape(tf.tile(indices, [1, num_pivots]), [-1])
        pivots_i = tf.tile(pivots, [num_pivots, 1])
        pivots_j = tf.gather(pivots, indices)
        overlap_adj = tf.to_float(num_inputs) * overlap_ratio
        forces = 0.5 * tf.multiply(
            tf.reshape(pivots_i - pivots_j, shape=[num_pivots, num_pivots, -1]),
            tf.expand_dims((overlap - overlap_adj) / overlap_adj, 2),
        )
        return tf.reduce_sum(forces / tf.to_float(num_pivots), 0)
//  [python]
def compute_forces(overlap, pivots, num_inputs, overlap_ratio, scope=None):
    """Computes accumulated forces based on pairwise hypersphere intersection.

    A (repulsive or attractive) force between two hyperspheres centered in
    corresponding pivots is computed based on the difference between their
    overlap versus desired overlap amount.

    An accumulated force is the average of all forces computed from all the
    other pivots.

    Args:
      overlap: A float32 `Tensor` of shape [num_pivots, num_pivots].
      pivots: A float32 `Tensor` of shape [num_pivots, input_dims].
      num_inputs: Number of training points.
      overlap_ratio: Constant to satisfy independence condition.
      scope: Optional variable scope.

    Returns:
      A float32 `Tensor` of shape [num_pivots, input_dims].
    """
    with tf.variable_scope(scope, 'ComputeForces', [overlap, pivots]):
        num_pivots = tf.shape(pivots)[0]
        indices = tf.reshape(tf.range(num_pivots), [-1, 1])
        indices = tf.reshape(tf.tile(indices, [1, num_pivots]), [-1])
        pivots_i = tf.tile(pivots, [num_pivots, 1])
        pivots_j = tf.gather(pivots, indices)
        overlap_adj = tf.cast(num_inputs, tf.float32) * overlap_ratio
        forces = 0.5 * tf.multiply(
            tf.reshape(pivots_i - pivots_j, shape=[num_pivots, num_pivots, -1]),
            tf.expand_dims((overlap - overlap_adj) / overlap_adj, 2),
        )
        return tf.reduce_sum(forces / tf.cast(num_pivots, tf.float32), 0)
//  [python]
def compute_statistics(inputs, pivots, overlap_ratio, scope=None):
    """Computes a matrix of hypersphere overlaps and associated statistics.

    Args:
      inputs: A float32 `Tensor` of shape [num_inputs, input_dims].
      pivots: A float32 `Tensor` of shape [num_pivots, input_dims].
      overlap_ratio: Constant to satisfy independence condition.
      scope: Optional variable scope.

    Returns:
      A tuple (`overlap`, `radii`, `avg`, `stdev`).
        overlap: 2-D `Tensor` of hypersphere overlaps.
        radii: 1-D `Tensor` of hypersphere radii.
        avg: Mean of the `overlap` `Tensor`.
        stdev: Standard deviation of the `overlap` `Tensor`.
    """
    with tf.variable_scope(scope, 'ComputeStatistics', [inputs, pivots]):
        dists = pairwise_distances(pivots, inputs)
        num_inputs = tf.shape(inputs)[0]
        median = (num_inputs + 1) // 2
        sorted_dists, _ = tf.nn.top_k(dists, median + 1)
        sorted_dists = tf.transpose(sorted_dists)
        radii = tf.gather(sorted_dists, [median])
        radii = tf.transpose(radii)
        inside_sphere = tf.cast(dists <= radii, tf.float32)
        overlap = tf.matmul(inside_sphere, tf.transpose(inside_sphere))
        num_pivots = tf.shape(pivots)[0]
        bits = tf.to_float(num_pivots * (num_pivots - 1))
        mean = tf.reduce_sum(zero_diag(overlap)) / bits
        stdev = tf.reduce_sum(zero_diag((overlap - mean) ** 2))
        stdev = tf.sqrt(stdev / bits)
        overlap_adj = tf.to_float(num_inputs) * overlap_ratio
        adj = tf.abs(overlap - overlap_adj)
        avg = tf.reduce_sum(zero_diag(adj)) / bits
        return overlap, radii, avg, stdev
//  [python]
def compute_statistics(inputs, pivots, overlap_ratio, scope=None):
    """Computes a matrix of hypersphere overlaps and associated statistics.

    Args:
      inputs: A float32 `Tensor` of shape [num_inputs, input_dims].
      pivots: A float32 `Tensor` of shape [num_pivots, input_dims].
      overlap_ratio: Constant to satisfy independence condition.
      scope: Optional variable scope.

    Returns:
      A tuple (`overlap`, `radii`, `avg`, `stdev`).
        overlap: 2-D `Tensor` of hypersphere overlaps.
        radii: 1-D `Tensor` of hypersphere radii.
        avg: Mean of the `overlap` `Tensor`.
        stdev: Standard deviation of the `overlap` `Tensor`.
    """
    with tf.variable_scope(scope, 'ComputeStatistics', [inputs, pivots]):
        dists = pairwise_distances(pivots, inputs)
        num_inputs = tf.shape(inputs)[0]
        median = (num_inputs + 1) // 2
        sorted_dists, _ = tf.nn.top_k(dists, median + 1)
        sorted_dists = tf.transpose(sorted_dists)
        radii = tf.gather(sorted_dists, [median])
        radii = tf.transpose(radii)
        inside_sphere = tf.cast(dists <= radii, tf.float32)
        overlap = tf.matmul(inside_sphere, tf.transpose(inside_sphere))
        num_pivots = tf.shape(pivots)[0]
        bits = tf.cast(num_pivots * (num_pivots - 1), tf.float32)
        mean = tf.reduce_sum(zero_diag(overlap)) / bits
        stdev = tf.reduce_sum(zero_diag((overlap - mean) ** 2))
        stdev = tf.sqrt(stdev / bits)
        overlap_adj = tf.cast(num_inputs, tf.float32) * overlap_ratio
        adj = tf.abs(overlap - overlap_adj)
        avg = tf.reduce_sum(zero_diag(adj)) / bits
        return overlap, radii, avg, stdev
//  [python]
def total_variation_loss(stylized_inputs, total_variation_weight):
  """Total variation regularization loss.

  This loss improves the smoothness of the image by expressing high frequency
  variations as a loss.
  http://link.springer.com/article/10.1023/B:JMIV.0000011325.36760.1e

  Args:
    stylized_inputs: The batched set of images.
    total_variation_weight: Weight of total variation loss.

  Returns:
    Tensor for the total variation loss, dict mapping loss names to losses.
  """
  shape = tf.shape(stylized_inputs)
  batch_size = shape[0]
  height = shape[1]
  width = shape[2]
  channels = shape[3]
  y_size = tf.to_float((height - 1) * width * channels)
  x_size = tf.to_float(height * (width - 1) * channels)
  y_loss = tf.nn.l2_loss(
      stylized_inputs[:, 1:, :, :] - stylized_inputs[:, :-1, :, :]) / y_size
  x_loss = tf.nn.l2_loss(
      stylized_inputs[:, :, 1:, :] - stylized_inputs[:, :, :-1, :]) / x_size
  loss = (y_loss + x_loss) / tf.to_float(batch_size)
  weighted_loss = loss * total_variation_weight
  return weighted_loss, {
      'total_variation_loss': loss,
      'weighted_total_variation_loss': weighted_loss
  }
//  [python]
def total_variation_loss(stylized_inputs, total_variation_weight):
  """Total variation regularization loss.

  This loss improves the smoothness of the image by expressing high frequency
  variations as a loss.
  http://link.springer.com/article/10.1023/B:JMIV.0000011325.36760.1e

  Args:
    stylized_inputs: The batched set of images.
    total_variation_weight: Weight of total variation loss.

  Returns:
    Tensor for the total variation loss, dict mapping loss names to losses.
  """
  shape = tf.shape(stylized_inputs)
  batch_size = shape[0]
  height = shape[1]
  width = shape[2]
  channels = shape[3]
  y_size = tf.cast((height - 1) * width * channels, tf.float32)
  x_size = tf.cast(height * (width - 1) * channels, tf.float32)
  y_loss = tf.nn.l2_loss(
      stylized_inputs[:, 1:, :, :] - stylized_inputs[:, :-1, :, :]) / y_size
  x_loss = tf.nn.l2_loss(
      stylized_inputs[:, :, 1:, :] - stylized_inputs[:, :, :-1, :]) / x_size
  loss = (y_loss + x_loss) / tf.cast(batch_size, tf.float32)
  weighted_loss = loss * total_variation_weight
  return weighted_loss, {
      'total_variation_loss': loss,
      'weighted_total_variation_loss': weighted_loss
  }
//  [python]
def spherical_hashing(
    inputs, num_bits, pack_bits=True, reuse=False, scope='SphericalHashing'
):
    """Spherical hashing from https://dl.acm.org/citation.cfm?id=2882197

    Heo, Jae-Pil, et al. "Spherical hashing." Computer Vision and Pattern
    Recognition (CVPR), 2012 IEEE Conference on. IEEE, 2012.

    Args:
      inputs: A float32 `Tensor` of input data.
      num_bits: Number of pivots (equivalently, number of hash bits).
      pack_bits: Whether to pack bits into int64s.
      reuse: Whether or not variables should be reused. To be able to reuse `scope`
        must be given.
      scope: Optional variable scope.

    Returns:
      A tuple of output bits and a dictionary of model variables.
    """
    with tf.variable_scope(scope, 'SphericalHashing', [inputs], reuse=reuse):
        inputs_dims = inputs.shape[1]
        pivots = tf.Variable(
            name='Pivots',
            shape=[num_bits, inputs_dims],
            initial_value=tf.zeros([num_bits, inputs_dims]),
            collections=[tf.GraphKeys.MODEL_VARIABLES],
        )
        radii = tf.Variable(
            name='Radii',
            shape=[num_bits, 1],
            initial_value=tf.zeros([num_bits, 1]),
            collections=[tf.GraphKeys.MODEL_VARIABLES],
        )
        distances = pairwise_distances(inputs, pivots)
        thresholds = distances - tf.transpose(radii)
        bits = tf.cast(thresholds <= 0, tf.uint8, name='Bits')
        if pack_bits:
            pow_2 = tf.pow(tf.to_int64(2), tf.range(0, 64, dtype=tf.int64))
            bits = tf.reshape(bits, [-1, 64])
            bits = tf.reduce_sum(tf.to_int64(bits) * pow_2, axis=1)
            bits = tf.reshape(bits, [-1, num_bits // 64])
        end_points = {'Pivots': pivots, 'Radii': radii}
        return bits, end_points
//  [python]
def spherical_hashing(
    inputs, num_bits, pack_bits=True, reuse=False, scope='SphericalHashing'
):
    """Spherical hashing from https://dl.acm.org/citation.cfm?id=2882197

    Heo, Jae-Pil, et al. "Spherical hashing." Computer Vision and Pattern
    Recognition (CVPR), 2012 IEEE Conference on. IEEE, 2012.

    Args:
      inputs: A float32 `Tensor` of input data.
      num_bits: Number of pivots (equivalently, number of hash bits).
      pack_bits: Whether to pack bits into int64s.
      reuse: Whether or not variables should be reused. To be able to reuse `scope`
        must be given.
      scope: Optional variable scope.

    Returns:
      A tuple of output bits and a dictionary of model variables.
    """
    with tf.variable_scope(scope, 'SphericalHashing', [inputs], reuse=reuse):
        inputs_dims = inputs.shape[1]
        pivots = tf.Variable(
            name='Pivots',
            shape=[num_bits, inputs_dims],
            initial_value=tf.zeros([num_bits, inputs_dims]),
            collections=[tf.GraphKeys.MODEL_VARIABLES],
        )
        radii = tf.Variable(
            name='Radii',
            shape=[num_bits, 1],
            initial_value=tf.zeros([num_bits, 1]),
            collections=[tf.GraphKeys.MODEL_VARIABLES],
        )
        distances = pairwise_distances(inputs, pivots)
        thresholds = distances - tf.transpose(radii)
        bits = tf.cast(thresholds <= 0, tf.uint8, name='Bits')
        if pack_bits:
            pow_2 = tf.pow(tf.cast(tf.range(0, 64, dtype=tf.int64), tf.int64), tf.range(0, 64, dtype=tf.int64))
            bits = tf.reshape(bits, [-1, 64])
            bits = tf.reduce_sum(tf.cast(bits, tf.int64) * pow_2, axis=1)
            bits = tf.reshape(bits, [-1, num_bits // 64])
        end_points = {'Pivots': pivots, 'Radii': radii}
        return bits, end_points
//  [python]
def get_resize_small(smaller_size):
  """Resizes the smaller side to `smaller_size` keeping aspect ratio."""
  def _resize_small_pp(data):
    image = data["image"]
    # A single image: HWC
    # A batch of images: BHWC
    h, w = tf.shape(image)[-3], tf.shape(image)[-2]

    # Figure out the necessary h/w.
    ratio = tf.to_float(smaller_size) / tf.to_float(tf.minimum(h, w))
    h = tf.to_int32(tf.round(tf.to_float(h) * ratio))
    w = tf.to_int32(tf.round(tf.to_float(w) * ratio))

    # NOTE: use align_corners=False for AREA resize, but True for Bilinear.
    # See also https://github.com/tensorflow/tensorflow/issues/6720
    static_rank = len(image.get_shape().as_list())
    if static_rank == 3:  # A single image: HWC
      data["image"] = tf.image.resize_area(image[None], [h, w])[0]
    elif static_rank == 4:  # A batch of images: BHWC
      data["image"] = tf.image.resize_area(image, [h, w])
    return data
  return _resize_small_pp
//  [python]
def get_resize_small(smaller_size):
  """Resizes the smaller side to `smaller_size` keeping aspect ratio."""
  def _resize_small_pp(data):
    image = data["image"]
    # A single image: HWC
    # A batch of images: BHWC
    h, w = tf.shape(image)[-3], tf.shape(image)[-2]

    # Figure out the necessary h/w.
    ratio = tf.to_float(smaller_size) / tf.to_float(tf.minimum(h, w))
    h = tf.cast(tf.round(tf.to_float(h) * ratio), tf.int32)
    w = tf.cast(tf.round(tf.to_float(w) * ratio), tf.int32)

    # NOTE: use align_corners=False for AREA resize, but True for Bilinear.
    # See also https://github.com/tensorflow/tensorflow/issues/6720
    static_rank = len(image.get_shape().as_list())
    if static_rank == 3:  # A single image: HWC
      data["image"] = tf.image.resize_area(image[None], [h, w])[0]
    elif static_rank == 4:  # A batch of images: BHWC
      data["image"] = tf.image.resize_area(image, [h, w])
    return data
  return _resize_small_pp
//  [python]
def train(self, input_sequence, output_sequence, sequence_length, control_sequence=None):
        """
        Train on the given sequences, returning an optimizer.
        :param input_sequence: The sequence to be fed to the encoder.
        :param output_sequence: The sequence expected from the decoder.
        :param sequence_length: The length of the given sequences (which must be
            identical).
        :param control_sequence: (Optional) sequence on which to condition. This will be
            concatenated depthwise to the model inputs for both encoding and decoding.
        :return:
            optimizer: A tf.train.Optimizer.
        """

        scalars_to_summarize = self._compute_model_loss(
            input_sequence, output_sequence, sequence_length, control_sequence
        )

        hparams = self.hparams
        lr = ((hparams.learning_rate - hparams.min_learning_rate) *
              tf.pow(hparams.decay_rate, tf.to_float(self.global_step)) +
              hparams.min_learning_rate)

        tf.summary.scalar('learning_rate', lr)
        for n, t in scalars_to_summarize.items():
            tf.summary.scalar(n, tf.reduce_mean(t))

        optimizer = tf.train.AdamOptimizer(lr)

        return optimizer
//  [python]
def train(self, input_sequence, output_sequence, sequence_length, control_sequence=None):
        """
        Train on the given sequences, returning an optimizer.
        :param input_sequence: The sequence to be fed to the encoder.
        :param output_sequence: The sequence expected from the decoder.
        :param sequence_length: The length of the given sequences (which must be
            identical).
        :param control_sequence: (Optional) sequence on which to condition. This will be
            concatenated depthwise to the model inputs for both encoding and decoding.
        :return:
            optimizer: A tf.train.Optimizer.
        """

        scalars_to_summarize = self._compute_model_loss(
            input_sequence, output_sequence, sequence_length, control_sequence
        )

        hparams = self.hparams
        lr = ((hparams.learning_rate - hparams.min_learning_rate) *
              tf.pow(hparams.decay_rate, tf.cast(self.global_step, tf.float32)) +
              hparams.min_learning_rate)

        tf.summary.scalar('learning_rate', lr)
        for n, t in scalars_to_summarize.items():
            tf.summary.scalar(n, tf.reduce_mean(t))

        optimizer = tf.train.AdamOptimizer(lr)

        return optimizer
//  [python]
def call(self, inputs):
        y_pred = inputs[0]

        cell_x = tf.to_float(
            tf.reshape(tf.tile(tf.range(config.GRID_W), [config.GRID_H]), (1, config.GRID_H, config.GRID_W, 1, 1)))
        cell_y = tf.transpose(cell_x, (0, 2, 1, 3, 4))

        cell_grid = tf.tile(tf.concat([cell_x, cell_y], -1), [config.BATCH_SIZE, 1, 1, config.N_BOX, 1])

        # compute x and y based on YOLOv2 paper formula:
        # x, y = sigma(tx) + cell_grid, sigma(ty) + cell_grid
        pred_box_xy = tf.sigmoid(y_pred[..., :2]) + cell_grid
        pred_box_xy = pred_box_xy / tf.cast(config.GRID_W, tf.float32)    # normalize to 0~1

        # compute w and h based on YOLOv2 paper formula:
        # w, h = prior_width (anchor widths) * exp(tw), prior_height (anchor heights) * exp(th)
        pred_box_wh = tf.exp(y_pred[..., 2:4]) * np.reshape(config.ANCHORS, [1, 1, 1, config.N_BOX, 2])
        pred_box_wh = pred_box_wh / tf.cast(config.GRID_W, tf.float32)    # normalize

        # get diagonal coordinates, xmin, ymin, xmax, ymax
        pred_wh_half = pred_box_wh / 2.
        pred_mins = pred_box_xy - pred_wh_half
        pred_maxes = pred_box_xy + pred_wh_half

        # concatenate and reshape
        output_boxes = tf.concat([pred_mins, pred_maxes], axis=-1)

        output_boxes = tf.reshape(output_boxes, [-1,
                                                 output_boxes.shape[1] * output_boxes.shape[2] * output_boxes.shape[3],
                                                 output_boxes.shape[-1]])

        return output_boxes
//  [python]
def call(self, inputs):
        y_pred = inputs[0]

        cell_x = tf.cast(
            tf.reshape(tf.tile(tf.range(config.GRID_W), [config.GRID_H]), (1, config.GRID_H, config.GRID_W, 1, 1)), dtype=tf.float32)
        cell_y = tf.transpose(cell_x, (0, 2, 1, 3, 4))

        cell_grid = tf.tile(tf.concat([cell_x, cell_y], -1), [config.BATCH_SIZE, 1, 1, config.N_BOX, 1])

        # compute x and y based on YOLOv2 paper formula:
        # x, y = sigma(tx) + cell_grid, sigma(ty) + cell_grid
        pred_box_xy = tf.sigmoid(y_pred[..., :2]) + cell_grid
        pred_box_xy = pred_box_xy / tf.cast(config.GRID_W, tf.float32)    # normalize to 0~1

        # compute w and h based on YOLOv2 paper formula:
        # w, h = prior_width (anchor widths) * exp(tw), prior_height (anchor heights) * exp(th)
        pred_box_wh = tf.exp(y_pred[..., 2:4]) * np.reshape(config.ANCHORS, [1, 1, 1, config.N_BOX, 2])
        pred_box_wh = pred_box_wh / tf.cast(config.GRID_W, tf.float32)    # normalize

        # get diagonal coordinates, xmin, ymin, xmax, ymax
        pred_wh_half = pred_box_wh / 2.
        pred_mins = pred_box_xy - pred_wh_half
        pred_maxes = pred_box_xy + pred_wh_half

        # concatenate and reshape
        output_boxes = tf.concat([pred_mins, pred_maxes], axis=-1)

        output_boxes = tf.reshape(output_boxes, [-1,
                                                 output_boxes.shape[1] * output_boxes.shape[2] * output_boxes.shape[3],
                                                 output_boxes.shape[-1]])

        return output_boxes
//  [python]
def train(self, input_sequence, output_sequence, sequence_length, control_sequence=None):
        """
        Train on the given sequences, returning an optimizer.
        :param input_sequence: The sequence to be fed to the encoder.
        :param output_sequence: The sequence expected from the decoder.
        :param sequence_length: The length of the given sequences (which must be
            identical).
        :param control_sequence: (Optional) sequence on which to condition. This will be
            concatenated depthwise to the model inputs for both encoding and decoding.
        :return:
            optimizer: A tf.train.Optimizer.
        """

        _, scalars_to_summarize = self._compute_model_loss(
            input_sequence, output_sequence, sequence_length, control_sequence
        )

        hparams = self.hparams
        lr = ((hparams.learning_rate - hparams.min_learning_rate) *
              tf.pow(hparams.decay_rate, tf.to_float(self.global_step)) +
              hparams.min_learning_rate)

        tf.summary.scalar('learning_rate', lr)
        for n, t in scalars_to_summarize.items():
            tf.summary.scalar(n, tf.reduce_mean(t))

        optimizer = tf.train.AdamOptimizer(lr)

        return optimizer
//  [python]
def train(self, input_sequence, output_sequence, sequence_length, control_sequence=None):
        """
        Train on the given sequences, returning an optimizer.
        :param input_sequence: The sequence to be fed to the encoder.
        :param output_sequence: The sequence expected from the decoder.
        :param sequence_length: The length of the given sequences (which must be
            identical).
        :param control_sequence: (Optional) sequence on which to condition. This will be
            concatenated depthwise to the model inputs for both encoding and decoding.
        :return:
            optimizer: A tf.train.Optimizer.
        """

        _, scalars_to_summarize = self._compute_model_loss(
            input_sequence, output_sequence, sequence_length, control_sequence
        )

        hparams = self.hparams
        lr = ((hparams.learning_rate - hparams.min_learning_rate) *
              tf.pow(hparams.decay_rate, tf.cast(self.global_step, tf.float32)) +
              hparams.min_learning_rate)

        tf.summary.scalar('learning_rate', lr)
        for n, t in scalars_to_summarize.items():
            tf.summary.scalar(n, tf.reduce_mean(t))

        optimizer = tf.train.AdamOptimizer(lr)

        return optimizer
//  [python]
def encode(self, sequence, sequence_length, control_sequence=None):
        """
        Encodes input sequences into a MultivariateNormalDiag distribution.
        :param sequence: A Tensor with shape `[num_sequences, max_length, input_depth]`
            containing the sequences to encode.
        :param sequence_length: The length of each sequence in the `sequence` Tensor.
        :param control_sequence: (Optional) A Tensor with shape
            `[num_sequences, max_length, control_length]` containing control sequences
            on which to condition. These will be concatenated depthwise to the input sequences.
        :return:
            A tfp.distribution.MultivariateNormalDiag representing the posterior distribution
            for each sequence.
        """
        hparams = self.hparams
        z_size = hparams.z_size

        sequence = tf.to_float(sequence)
        if control_sequence is not None:
            control_sequence = tf.to_float(control_sequence)
            sequence = tf.concat([sequence, control_sequence], axis=-1)
        encoder_output = self.encoder.encode(sequence, sequence_length)

        mu = tf.layers.Dense(
            z_size,
            name='encoder/mu',
            kernel_initializer=tf.random_normal_initializer(stddev=0.001),
        )(encoder_output)
        sigma = tf.layers.Dense(
            z_size,
            activation=tf.nn.softplus,
            name='encoder/sigma',
            kernel_initializer=tf.random_normal_initializer(stddev=0.001),
        )(encoder_output)

        return ds.MultivariateNormalDiag(loc=mu, scale_diag=sigma)
//  [python]
def encode(self, sequence, sequence_length, control_sequence=None):
        """
        Encodes input sequences into a MultivariateNormalDiag distribution.
        :param sequence: A Tensor with shape `[num_sequences, max_length, input_depth]`
            containing the sequences to encode.
        :param sequence_length: The length of each sequence in the `sequence` Tensor.
        :param control_sequence: (Optional) A Tensor with shape
            `[num_sequences, max_length, control_length]` containing control sequences
            on which to condition. These will be concatenated depthwise to the input sequences.
        :return:
            A tfp.distribution.MultivariateNormalDiag representing the posterior distribution
            for each sequence.
        """
        hparams = self.hparams
        z_size = hparams.z_size

        sequence = tf.cast(sequence, tf.float32)
        if control_sequence is not None:
            control_sequence = tf.cast(control_sequence, tf.float32)
            sequence = tf.concat([sequence, control_sequence], axis=-1)
        encoder_output = self.encoder.encode(sequence, sequence_length)

        mu = tf.layers.Dense(
            z_size,
            name='encoder/mu',
            kernel_initializer=tf.random_normal_initializer(stddev=0.001),
        )(encoder_output)
        sigma = tf.layers.Dense(
            z_size,
            activation=tf.nn.softplus,
            name='encoder/sigma',
            kernel_initializer=tf.random_normal_initializer(stddev=0.001),
        )(encoder_output)

        return ds.MultivariateNormalDiag(loc=mu, scale_diag=sigma)
//  [python]
def _compute_model_loss(
            self, input_sequence, output_sequence, sequence_length, control_sequence):
        """Builds a model with loss for train/eval."""
        hparams = self.hparams
        batch_size = hparams.batch_size

        input_sequence = tf.to_float(input_sequence)
        output_sequence = tf.to_float(output_sequence)

        max_seq_len = tf.minimum(tf.shape(output_sequence)[1], hparams.max_seq_len)

        input_sequence = input_sequence[:, :max_seq_len]

        if control_sequence is not None:
            control_depth = control_sequence.shape[-1]
            control_sequence = tf.to_float(control_sequence)
            control_sequence = control_sequence[:, :max_seq_len]
            # Shouldn't be necessary, but the slice loses shape information when
            # control depth is zero.
            control_sequence.set_shape([batch_size, None, control_depth])

        # Inputs to be fed to decoder, including zero padding for the initial input.
        x_length = tf.minimum(sequence_length, max_seq_len)
        q_z = self.encode(input_sequence, x_length, control_sequence)
        z = q_z.sample()
        q_encode_z = self.encode_latent(z)
        latent_z = q_encode_z.sample()

        # KL Divergence (nats)
        kl_cost = self._kl_latent_loss(q_encode_z)

        # Reconstruction loss
        r_loss = self._latent_reconstruction_loss(z, latent_z)

        beta = ((1.0 - tf.pow(hparams.beta_rate, tf.to_float(self.global_step))) * hparams.max_beta)
        self.loss = tf.reduce_mean(r_loss) + beta * tf.reduce_mean(kl_cost)

        scalars_to_summarize = {
            'loss': self.loss,
            'losses/r_loss': r_loss,
            'losses/kl_loss': kl_cost,
            'losses/kl_beta': beta,
        }
        return scalars_to_summarize
//  [python]
def _compute_model_loss(
            self, input_sequence, output_sequence, sequence_length, control_sequence):
        """Builds a model with loss for train/eval."""
        hparams = self.hparams
        batch_size = hparams.batch_size

        input_sequence = tf.cast(input_sequence, tf.float32)
        output_sequence = tf.cast(output_sequence, tf.float32)

        max_seq_len = tf.minimum(tf.shape(output_sequence)[1], hparams.max_seq_len)

        input_sequence = input_sequence[:, :max_seq_len]

        if control_sequence is not None:
            control_depth = control_sequence.shape[-1]
            control_sequence = tf.cast(control_sequence, tf.float32)
            control_sequence = control_sequence[:, :max_seq_len]
            # Shouldn't be necessary, but the slice loses shape information when
            # control depth is zero.
            control_sequence.set_shape([batch_size, None, control_depth])

        # Inputs to be fed to decoder, including zero padding for the initial input.
        x_length = tf.minimum(sequence_length, max_seq_len)
        q_z = self.encode(input_sequence, x_length, control_sequence)
        z = q_z.sample()
        q_encode_z = self.encode_latent(z)
        latent_z = q_encode_z.sample()

        # KL Divergence (nats)
        kl_cost = self._kl_latent_loss(q_encode_z)

        # Reconstruction loss
        r_loss = self._latent_reconstruction_loss(z, latent_z)

        beta = ((1.0 - tf.pow(hparams.beta_rate, tf.cast(self.global_step, tf.float32))) * hparams.max_beta)
        self.loss = tf.reduce_mean(r_loss) + beta * tf.reduce_mean(kl_cost)

        scalars_to_summarize = {
            'loss': self.loss,
            'losses/r_loss': r_loss,
            'losses/kl_loss': kl_cost,
            'losses/kl_beta': beta,
        }
        return scalars_to_summarize
//  [python]
def call(self, inputs):
        y_pred = inputs[0]
        # mask_shape = tf.shape(y_pred)[:4]
        # grid_w, grid_h = mask_shape[1], mask_shape[2]

        cell_x = tf.to_float(
            tf.reshape(tf.tile(tf.range(config.GRID_W), [config.GRID_H]), (1, config.GRID_H, config.GRID_W, 1, 1)))
        cell_y = tf.transpose(cell_x, (0, 2, 1, 3, 4))

        cell_grid = tf.tile(tf.concat([cell_x, cell_y], -1), [config.BATCH_SIZE, 1, 1, config.N_BOX, 1])

        """ Adjust prediction """
        # adjust x and y
        pred_box_xy = tf.sigmoid(y_pred[..., :2]) + cell_grid
        pred_box_xy = pred_box_xy / tf.cast(config.GRID_W, tf.float32)
        # adjust w and h
        pred_box_wh = tf.exp(y_pred[..., 2:4]) * np.reshape(config.ANCHORS, [1, 1, 1, config.N_BOX, 2])
        pred_box_wh = pred_box_wh / tf.cast(config.GRID_W, tf.float32)    # normalize

        """ get x, y coordinates """
        # pred_xy = tf.expand_dims(pred_box_xy, 4)
        # pred_wh = tf.expand_dims(pred_box_wh, 4)

        pred_wh_half = pred_box_wh / 2.
        pred_mins = pred_box_xy - pred_wh_half
        pred_maxes = pred_box_xy + pred_wh_half

        # xmin, ymin, xmax, ymax
        # output_boxes = tf.concat([pred_mins, pred_maxes], axis=-1)
        #
        # output_boxes = tf.reshape(output_boxes, [-1,
        #                                          output_boxes.shape[1] * output_boxes.shape[2] * output_boxes.shape[3],
        #                                          output_boxes.shape[-1]])

        pred_box_conf = tf.sigmoid(y_pred[..., 4])
        pred_box_conf = tf.expand_dims(pred_box_conf, axis=-1)

        pred_box_class = tf.argmax(y_pred[..., 5:], axis=-1)
        pred_box_class = tf.expand_dims(tf.to_float(pred_box_class), axis=-1)

        detections = tf.concat([pred_mins, pred_maxes, pred_box_conf, pred_box_class], axis=-1)
        detections = tf.reshape(detections, [-1,
                                             detections.shape[1] * detections.shape[2] * detections.shape[3],
                                             detections.shape[-1]])

        return detections
//  [python]
def call(self, inputs):
        y_pred = inputs[0]
        # mask_shape = tf.shape(y_pred)[:4]
        # grid_w, grid_h = mask_shape[1], mask_shape[2]

        cell_x = tf.cast(
            tf.reshape(tf.tile(tf.range(config.GRID_W), [config.GRID_H]), (1, config.GRID_H, config.GRID_W, 1, 1)), tf.float32)
        cell_y = tf.transpose(cell_x, (0, 2, 1, 3, 4))

        cell_grid = tf.tile(tf.concat([cell_x, cell_y], -1), [config.BATCH_SIZE, 1, 1, config.N_BOX, 1])

        """ Adjust prediction """
        # adjust x and y
        pred_box_xy = tf.sigmoid(y_pred[..., :2]) + cell_grid
        pred_box_xy = pred_box_xy / tf.cast(config.GRID_W, tf.float32)
        # adjust w and h
        pred_box_wh = tf.exp(y_pred[..., 2:4]) * np.reshape(config.ANCHORS, [1, 1, 1, config.N_BOX, 2])
        pred_box_wh = pred_box_wh / tf.cast(config.GRID_W, tf.float32)    # normalize

        """ get x, y coordinates """
        # pred_xy = tf.expand_dims(pred_box_xy, 4)
        # pred_wh = tf.expand_dims(pred_box_wh, 4)

        pred_wh_half = pred_box_wh / 2.
        pred_mins = pred_box_xy - pred_wh_half
        pred_maxes = pred_box_xy + pred_wh_half

        # xmin, ymin, xmax, ymax
        # output_boxes = tf.concat([pred_mins, pred_maxes], axis=-1)
        #
        # output_boxes = tf.reshape(output_boxes, [-1,
        #                                          output_boxes.shape[1] * output_boxes.shape[2] * output_boxes.shape[3],
        #                                          output_boxes.shape[-1]])

        pred_box_conf = tf.sigmoid(y_pred[..., 4])
        pred_box_conf = tf.expand_dims(pred_box_conf, axis=-1)

        pred_box_class = tf.argmax(y_pred[..., 5:], axis=-1)
        pred_box_class = tf.expand_dims(tf.cast(pred_box_class, tf.float32), axis=-1)

        detections = tf.concat([pred_mins, pred_maxes, pred_box_conf, pred_box_class], axis=-1)
        detections = tf.reshape(detections, [-1,
                                             detections.shape[1] * detections.shape[2] * detections.shape[3],
                                             detections.shape[-1]])

        return detections
//  [python]
def get_random_bounding_box(
    hw,
    aspect_ratio=1,
    min_crop_fraction=0.5,
    new_height=None,
    dtype=tf.int32):
  """Random bounding box with aspect_ratio and within (0,0,hw[:,0],hw[:,1])."""
  hw = tf.to_float(hw)
  h = hw[Ellipsis, 0]
  w = hw[Ellipsis, 1]
  if new_height is None:
    min_height = tf.minimum(h * min_crop_fraction,
                            w * min_crop_fraction / aspect_ratio)
    max_height = tf.minimum(h, w / aspect_ratio)
    new_h = tf.random.uniform(
        shape=h.shape.as_list(),
        minval=min_height,
        maxval=max_height,
        dtype=tf.float32)
  else:
    new_h = tf.to_float(new_height)
  new_hw = tf.stack([new_h, new_h / aspect_ratio], axis=-1)
  if dtype in [tf.int32, tf.int64]:
    new_hw = tf.round(new_hw)

  minval = tf.zeros_like(hw)
  maxval = hw - new_hw
  tl = tf.random.uniform(shape=hw.shape, minval=minval,
                         maxval=maxval, dtype=tf.float32)
  if dtype in [tf.int32, tf.int64]:
    tl = tf.round(tl)

  br = tl + new_hw
  boxes = tf.concat((tl, br), axis=-1)
  if dtype in [tf.int32, tf.int64]:
    boxes = tf.round(boxes)
  return tf.cast(boxes, dtype=dtype)
//  [python]
def get_random_bounding_box(
    hw,
    aspect_ratio=1,
    min_crop_fraction=0.5,
    new_height=None,
    dtype=tf.int32):
  """Random bounding box with aspect_ratio and within (0,0,hw[:,0],hw[:,1])."""
  hw = tf.cast(hw, tf.float32)
  h = hw[Ellipsis, 0]
  w = hw[Ellipsis, 1]
  if new_height is None:
    min_height = tf.minimum(h * min_crop_fraction,
                            w * min_crop_fraction / aspect_ratio)
    max_height = tf.minimum(h, w / aspect_ratio)
    new_h = tf.random.uniform(
        shape=h.shape.as_list(),
        minval=min_height,
        maxval=max_height,
        dtype=tf.float32)
  else:
    new_h = tf.cast(new_height, tf.float32)
  new_hw = tf.stack([new_h, new_h / aspect_ratio], axis=-1)
  if dtype in [tf.int32, tf.int64]:
    new_hw = tf.round(new_hw)

  minval = tf.zeros_like(hw)
  maxval = hw - new_hw
  tl = tf.random.uniform(shape=hw.shape, minval=minval,
                         maxval=maxval, dtype=tf.float32)
  if dtype in [tf.int32, tf.int64]:
    tl = tf.round(tl)

  br = tl + new_hw
  boxes = tf.concat((tl, br), axis=-1)
  if dtype in [tf.int32, tf.int64]:
    boxes = tf.round(boxes)
  return tf.cast(boxes, dtype=dtype)
//  [python]
def bleu_score(predictions, labels, **unused_kwargs):
  """BLEU score computation between labels and predictions.

  An approximate BLEU scoring method since we do not glue word pieces or
  decode the ids and tokenize the output. By default, we use ngram order of 4
  and use brevity penalty. Also, this does not have beam search.

  Args:
    predictions: tensor, model predictions
    labels: tensor, gold output.

  Returns:
    bleu: int, approx bleu score
  """
  outputs = tf.to_int32(tf.argmax(predictions, axis=-1))
  # Convert the outputs and labels to a [batch_size, input_length] tensor.
  outputs = tf.squeeze(outputs, axis=[-1, -2])
  labels = tf.squeeze(labels, axis=[-1, -2])

  bleu = tf.py_func(compute_bleu, (labels, outputs), tf.float32)
  return bleu, tf.constant(1.0)
//  [python]
def bleu_score(predictions, labels, **unused_kwargs):
  """BLEU score computation between labels and predictions.

  An approximate BLEU scoring method since we do not glue word pieces or
  decode the ids and tokenize the output. By default, we use ngram order of 4
  and use brevity penalty. Also, this does not have beam search.

  Args:
    predictions: tensor, model predictions
    labels: tensor, gold output.

  Returns:
    bleu: int, approx bleu score
  """
  outputs = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)
  # Convert the outputs and labels to a [batch_size, input_length] tensor.
  outputs = tf.squeeze(outputs, axis=[-1, -2])
  labels = tf.squeeze(labels, axis=[-1, -2])

  bleu = tf.py_func(compute_bleu, (labels, outputs), tf.float32)
  return bleu, tf.constant(1.0)
//  [python]
def __call__(self, features, is_training):
        seqs_id = features['seqs_i']
        seqs_ts = features['seqs_t'] / self.time_scale
        seqs_spans = seqs_ts[:, 1:] - seqs_ts[:, :-1]

        # positional encoding
        seqs_marks = tf.nn.embedding_lookup(self.mark_lookup_table, seqs_id)
        seqs_units = self.item_embs(seqs_id)
        seqs_units = self.pcoding(seqs_units)

        # Dropout
        seqs_units = tf.layers.dropout(seqs_units, rate=self.hidden_dropout_rate,
                                       training=tf.convert_to_tensor(is_training))
        seqs_masks = tf.expand_dims(tf.to_float(tf.not_equal(seqs_id, 0)), 1)
        seqs_masks = tf.tile(seqs_masks, [self.num_heads, tf.shape(seqs_id)[1], 1])

        # multi-head self-modulating attention
        seqs_outs = seqs_units
        for i, (attention, dense) in enumerate(zip(self.list_attention, self.list_dense)):
            with tf.variable_scope("num_blocks_%d" % i):
                # sequential-temporal representations
                with tf.variable_scope("attention"):
                    seqs_outs, seqs_intny = attention(layernorm(seqs_outs), seqs_outs, seqs_masks,
                                                      seqs_spans, seqs_marks, is_training, causality=True)

                # feed-forward
                with tf.variable_scope("feed-forward"):
                    seqs_outs = dense(layernorm(seqs_outs), is_training)

                # likelihood for point process
                if is_training:
                    tf.add_to_collection("LLE_PP", seqs_intny)

        with tf.variable_scope("outln"):
            seqs_outs = layernorm(seqs_outs)

        if is_training:
            seqs_outs = tf.reshape(seqs_outs, [tf.shape(seqs_id)[0] * self.seqslen, self.num_units])
        else:
            # only using the latest representations for making predictions
            seqs_outs = tf.reshape(seqs_outs[:, -1], [tf.shape(seqs_id)[0], self.num_units])

        # compute logits
        logits = tf.matmul(seqs_outs, self.item_embs.lookup_table, transpose_b=True)
        logits = tf.nn.bias_add(logits, self.output_bias)
        return logits
//  [python]
def __call__(self, features, is_training):
        seqs_id = features['seqs_i']
        seqs_ts = features['seqs_t'] / self.time_scale
        seqs_spans = seqs_ts[:, 1:] - seqs_ts[:, :-1]

        # positional encoding
        seqs_marks = tf.nn.embedding_lookup(self.mark_lookup_table, seqs_id)
        seqs_units = self.item_embs(seqs_id)
        seqs_units = self.pcoding(seqs_units)

        # Dropout
        seqs_units = tf.layers.dropout(seqs_units, rate=self.hidden_dropout_rate,
                                       training=tf.convert_to_tensor(is_training))
        seqs_masks = tf.expand_dims(tf.cast(tf.not_equal(seqs_id, 0), tf.float32), 1)
        seqs_masks = tf.tile(seqs_masks, [self.num_heads, tf.shape(seqs_id)[1], 1])

        # multi-head self-modulating attention
        seqs_outs = seqs_units
        for i, (attention, dense) in enumerate(zip(self.list_attention, self.list_dense)):
            with tf.variable_scope("num_blocks_%d" % i):
                # sequential-temporal representations
                with tf.variable_scope("attention"):
                    seqs_outs, seqs_intny = attention(layernorm(seqs_outs), seqs_outs, seqs_masks,
                                                      seqs_spans, seqs_marks, is_training, causality=True)

                # feed-forward
                with tf.variable_scope("feed-forward"):
                    seqs_outs = dense(layernorm(seqs_outs), is_training)

                # likelihood for point process
                if is_training:
                    tf.add_to_collection("LLE_PP", seqs_intny)

        with tf.variable_scope("outln"):
            seqs_outs = layernorm(seqs_outs)

        if is_training:
            seqs_outs = tf.reshape(seqs_outs, [tf.shape(seqs_id)[0] * self.seqslen, self.num_units])
        else:
            # only using the latest representations for making predictions
            seqs_outs = tf.reshape(seqs_outs[:, -1], [tf.shape(seqs_id)[0], self.num_units])

        # compute logits
        logits = tf.matmul(seqs_outs, self.item_embs.lookup_table, transpose_b=True)
        logits = tf.nn.bias_add(logits, self.output_bias)
        return logits