jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, num_outputs, reuse=False, trainable=True):
    self.num_outputs = num_outputs

    # Placeholders for our input
    # Our input are 4 RGB frames of shape 160, 160 each
    self.states = tf.placeholder(shape=[None, 84, 84, 4], dtype=tf.uint8, name="X")
    # The TD target value
    self.targets = tf.placeholder(shape=[None], dtype=tf.float32, name="y")
    # Integer id of which action was selected
    self.actions = tf.placeholder(shape=[None], dtype=tf.int32, name="actions")

    # Normalize
    X = tf.to_float(self.states) / 255.0
    batch_size = tf.shape(self.states)[0]

    # Graph shared with Value Net
    with tf.variable_scope("shared", reuse=reuse):
      fc1 = build_shared_network(X, add_summaries=(not reuse))


    with tf.variable_scope("policy_net"):
      self.logits = tf.contrib.layers.fully_connected(fc1, num_outputs, activation_fn=None)
      self.probs = tf.nn.softmax(self.logits) + 1e-8

      self.predictions = {
        "logits": self.logits,
        "probs": self.probs
      }

      # We add entropy to the loss to encourage exploration
      self.entropy = -tf.reduce_sum(self.probs * tf.log(self.probs), 1, name="entropy")
      self.entropy_mean = tf.reduce_mean(self.entropy, name="entropy_mean")

      # Get the predictions for the chosen actions only
      gather_indices = tf.range(batch_size) * tf.shape(self.probs)[1] + self.actions
      self.picked_action_probs = tf.gather(tf.reshape(self.probs, [-1]), gather_indices)

      self.losses = - (tf.log(self.picked_action_probs) * self.targets + 0.01 * self.entropy)
      self.loss = tf.reduce_sum(self.losses, name="loss")

      tf.summary.scalar(self.loss.op.name, self.loss)
      tf.summary.scalar(self.entropy_mean.op.name, self.entropy_mean)
      tf.summary.histogram(self.entropy.op.name, self.entropy)

      if trainable:
        # self.optimizer = tf.train.AdamOptimizer(1e-4)
        self.optimizer = tf.train.RMSPropOptimizer(0.00025, 0.99, 0.0, 1e-6)
        self.grads_and_vars = self.optimizer.compute_gradients(self.loss)
        self.grads_and_vars = [[grad, var] for grad, var in self.grads_and_vars if grad is not None]
        self.train_op = self.optimizer.apply_gradients(self.grads_and_vars,
          global_step=tf.contrib.framework.get_global_step())

    # Merge summaries from this network and the shared network (but not the value net)
    var_scope_name = tf.get_variable_scope().name
    summary_ops = tf.get_collection(tf.GraphKeys.SUMMARIES)
    sumaries = [s for s in summary_ops if "policy_net" in s.name or "shared" in s.name]
    sumaries = [s for s in summary_ops if var_scope_name in s.name]
    self.summaries = tf.summary.merge(sumaries)
//  [python]
def __init__(self, num_outputs, reuse=False, trainable=True):
    self.num_outputs = num_outputs

    # Placeholders for our input
    # Our input are 4 RGB frames of shape 160, 160 each
    self.states = tf.placeholder(shape=[None, 84, 84, 4], dtype=tf.uint8, name="X")
    # The TD target value
    self.targets = tf.placeholder(shape=[None], dtype=tf.float32, name="y")
    # Integer id of which action was selected
    self.actions = tf.placeholder(shape=[None], dtype=tf.int32, name="actions")

    # Normalize
    X = tf.cast(self.states, tf.float32) / 255.0
    batch_size = tf.shape(self.states)[0]

    # Graph shared with Value Net
    with tf.variable_scope("shared", reuse=reuse):
      fc1 = build_shared_network(X, add_summaries=(not reuse))


    with tf.variable_scope("policy_net"):
      self.logits = tf.contrib.layers.fully_connected(fc1, num_outputs, activation_fn=None)
      self.probs = tf.nn.softmax(self.logits) + 1e-8

      self.predictions = {
        "logits": self.logits,
        "probs": self.probs
      }

      # We add entropy to the loss to encourage exploration
      self.entropy = -tf.reduce_sum(self.probs * tf.log(self.probs), 1, name="entropy")
      self.entropy_mean = tf.reduce_mean(self.entropy, name="entropy_mean")

      # Get the predictions for the chosen actions only
      gather_indices = tf.range(batch_size) * tf.shape(self.probs)[1] + self.actions
      self.picked_action_probs = tf.gather(tf.reshape(self.probs, [-1]), gather_indices)

      self.losses = - (tf.log(self.picked_action_probs) * self.targets + 0.01 * self.entropy)
      self.loss = tf.reduce_sum(self.losses, name="loss")

      tf.summary.scalar(self.loss.op.name, self.loss)
      tf.summary.scalar(self.entropy_mean.op.name, self.entropy_mean)
      tf.summary.histogram(self.entropy.op.name, self.entropy)

      if trainable:
        # self.optimizer = tf.train.AdamOptimizer(1e-4)
        self.optimizer = tf.train.RMSPropOptimizer(0.00025, 0.99, 0.0, 1e-6)
        self.grads_and_vars = self.optimizer.compute_gradients(self.loss)
        self.grads_and_vars = [[grad, var] for grad, var in self.grads_and_vars if grad is not None]
        self.train_op = self.optimizer.apply_gradients(self.grads_and_vars,
          global_step=tf.contrib.framework.get_global_step())

    # Merge summaries from this network and the shared network (but not the value net)
    var_scope_name = tf.get_variable_scope().name
    summary_ops = tf.get_collection(tf.GraphKeys.SUMMARIES)
    sumaries = [s for s in summary_ops if "policy_net" in s.name or "shared" in s.name]
    sumaries = [s for s in summary_ops if var_scope_name in s.name]
    self.summaries = tf.summary.merge(sumaries)
//  [python]
def smooth_l1_loss_rcnn_r(bbox_pred, bbox_targets, label, num_classes, sigma=1.0):
    '''

    :param bbox_pred: [-1, (cfgs.CLS_NUM +1) * 5]
    :param bbox_targets:[-1, (cfgs.CLS_NUM +1) * 5]
    :param label:[-1]
    :param num_classes:
    :param sigma:
    :return:
    '''

    outside_mask = tf.stop_gradient(tf.to_float(tf.greater(label, 0)))

    bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 5])
    bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 5])

    value = _smooth_l1_loss_base(bbox_pred,
                                 bbox_targets,
                                 sigma=sigma)
    value = tf.reduce_sum(value, 2)
    value = tf.reshape(value, [-1, num_classes])

    inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]),
                             depth=num_classes, axis=1)

    inside_mask = tf.stop_gradient(
        tf.to_float(tf.reshape(inside_mask, [-1, num_classes])))

    normalizer = tf.to_float(tf.shape(bbox_pred)[0])
    bbox_loss = tf.reduce_sum(
        tf.reduce_sum(value * inside_mask, 1)*outside_mask) / normalizer

    return bbox_loss
//  [python]
def smooth_l1_loss_rcnn_r(bbox_pred, bbox_targets, label, num_classes, sigma=1.0):
    '''

    :param bbox_pred: [-1, (cfgs.CLS_NUM +1) * 5]
    :param bbox_targets:[-1, (cfgs.CLS_NUM +1) * 5]
    :param label:[-1]
    :param num_classes:
    :param sigma:
    :return:
    '''

    outside_mask = tf.stop_gradient(tf.cast(tf.greater(label, 0), tf.float32))

    bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 5])
    bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 5])

    value = _smooth_l1_loss_base(bbox_pred,
                                 bbox_targets,
                                 sigma=sigma)
    value = tf.reduce_sum(value, 2)
    value = tf.reshape(value, [-1, num_classes])

    inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]),
                             depth=num_classes, axis=1)

    inside_mask = tf.stop_gradient(
        tf.cast(tf.reshape(inside_mask, [-1, num_classes]), tf.float32))

    normalizer = tf.cast(tf.shape(bbox_pred)[0], tf.float32)
    bbox_loss = tf.reduce_sum(
        tf.reduce_sum(value * inside_mask, 1)*outside_mask) / normalizer

    return bbox_loss
//  [python]
def dataset(directory, images_file, labels_file):
    """Download and parse MNIST dataset."""

    # images_file = download(directory, images_file)
    # labels_file = download(directory, labels_file)
    #
    # check_image_file_header(images_file)
    # check_labels_file_header(labels_file)

    def decode_image(image):
        # Normalize from [0, 255] to [0.0, 1.0]
        image = tf.decode_raw(image, tf.uint8)
        image = tf.cast(image, tf.float32)
        image = tf.reshape(image, [50176])
        return image / 255.0

    def decode_label(label):
        label = tf.decode_raw(label, tf.uint8)  # tf.string -> [tf.uint8]
        label = tf.reshape(label, [])  # label is a scalar
        return tf.to_int32(label)

    images = tf.data.FixedLengthRecordDataset(images_file, 224 * 224, header_bytes=16)
    images = images.map(decode_image)
    labels = tf.data.FixedLengthRecordDataset(labels_file, 1, header_bytes=8).map(decode_label)

    return tf.data.Dataset.zip((images, labels))
//  [python]
def dataset(directory, images_file, labels_file):
    """Download and parse MNIST dataset."""

    # images_file = download(directory, images_file)
    # labels_file = download(directory, labels_file)
    #
    # check_image_file_header(images_file)
    # check_labels_file_header(labels_file)

    def decode_image(image):
        # Normalize from [0, 255] to [0.0, 1.0]
        image = tf.decode_raw(image, tf.uint8)
        image = tf.cast(image, tf.float32)
        image = tf.reshape(image, [50176])
        return image / 255.0

    def decode_label(label):
        label = tf.decode_raw(label, tf.uint8)  # tf.string -> [tf.uint8]
        label = tf.reshape(label, [])  # label is a scalar
        return tf.cast(label, tf.int32)

    images = tf.data.FixedLengthRecordDataset(images_file, 224 * 224, header_bytes=16)
    images = images.map(decode_image)
    labels = tf.data.FixedLengthRecordDataset(labels_file, 1, header_bytes=8).map(decode_label)

    return tf.data.Dataset.zip((images, labels))
//  [python]
def update_centroids(self, nearest_indices):
        partitions = tf.dynamic_partition(self.v_data, tf.to_int32(nearest_indices), self.n_clusters)
        return tf.concat([tf.expand_dims(tf.reduce_mean(partition, 0), 0)
                                      for partition in partitions], 0)
//  [python]
def update_centroids(self, nearest_indices):
        partitions = tf.dynamic_partition(self.v_data, tf.cast(nearest_indices, tf.int32), self.n_clusters)
        return tf.concat([tf.expand_dims(tf.reduce_mean(partition, 0), 0)
                                      for partition in partitions], 0)
//  [python]
def call(self, partioned_val, partioned_keys,
           partioned_indices, **kwargs):  # pylint: disable=unused-argument, arguments-differ
    """
    This layer is responsible for stitching a partioned layer together.

    Input:
      partioned_val:
        a list of partioned Tensors which represent the vals of the hashmap
      partioned_keys:
        a list of partioned Tensors which represent the keys of the hashmap
      partioned_indices:
        a list of partioned Tensors which represent the indices of the hashmap
    Output:
      List which contains: [output_vals, output_keys]
        output_vals:
          Values of the HashMap (float)
        output_keys:
          Keys of HashMap (float)
    """
    indices = [tf.to_int32(index) for index in partioned_indices]
    concat_keys = tf.dynamic_stitch(indices, partioned_keys)
    concat_vals = tf.dynamic_stitch(indices, partioned_val)
    return [concat_vals, concat_keys]
//  [python]
def call(self, partioned_val, partioned_keys,
           partioned_indices, **kwargs):  # pylint: disable=unused-argument, arguments-differ
    """
    This layer is responsible for stitching a partioned layer together.

    Input:
      partioned_val:
        a list of partioned Tensors which represent the vals of the hashmap
      partioned_keys:
        a list of partioned Tensors which represent the keys of the hashmap
      partioned_indices:
        a list of partioned Tensors which represent the indices of the hashmap
    Output:
      List which contains: [output_vals, output_keys]
        output_vals:
          Values of the HashMap (float)
        output_keys:
          Keys of HashMap (float)
    """
    indices = [tf.cast(index, tf.int32) for index in partioned_indices]
    concat_keys = tf.dynamic_stitch(indices, partioned_keys)
    concat_vals = tf.dynamic_stitch(indices, partioned_val)
    return [concat_vals, concat_keys]
//  [python]
def find_initial_centroids(self, k):
        r_index = tf.random_uniform([1], 0, self.n_data, dtype=tf.int32)
        r = tf.expand_dims(self.v_data[tf.squeeze(r_index)], dim=1)
        initial_centroids = []
        for i in range(k):
            dist = all_distances(self.v_data, r)
            farthest_index = tf.argmax(tf.reduce_min(dist, axis=0), 0)
            farthest_point = self.v_data[tf.to_int32(farthest_index)]
            initial_centroids.append(farthest_point)
            r = tf.stack(initial_centroids)
        return r
//  [python]
def find_initial_centroids(self, k):
        r_index = tf.random_uniform([1], 0, self.n_data, dtype=tf.int32)
        r = tf.expand_dims(self.v_data[tf.squeeze(r_index)], dim=1)
        initial_centroids = []
        for i in range(k):
            dist = all_distances(self.v_data, r)
            farthest_index = tf.argmax(tf.reduce_min(dist, axis=0), 0)
            farthest_point = self.v_data[tf.cast(farthest_index, tf.int32)]
            initial_centroids.append(farthest_point)
            r = tf.stack(initial_centroids)
        return r
//  [python]
def gaussian_blur(image, kernel_size, sigma, padding='SAME'):
  """Blurs the given image with separable convolution.
  Args:
    image: Tensor of shape [height, width, channels] and dtype float to blur.
    kernel_size: Integer Tensor for the size of the blur kernel. This is should
      be an odd number. If it is an even number, the actual kernel size will be
      size + 1.
    sigma: Sigma value for gaussian operator.
    padding: Padding to use for the convolution. Typically 'SAME' or 'VALID'.
  Returns:
    A Tensor representing the blurred image.
  """
  radius = tf.to_int32(kernel_size / 2)
  kernel_size = radius * 2 + 1
  x = tf.to_float(tf.range(-radius, radius + 1))
  blur_filter = tf.exp(
      -tf.pow(x, 2.0) / (2.0 * tf.pow(tf.to_float(sigma), 2.0)))
  blur_filter /= tf.reduce_sum(blur_filter)
  # One vertical and one horizontal filter.
  blur_v = tf.reshape(blur_filter, [kernel_size, 1, 1, 1])
  blur_h = tf.reshape(blur_filter, [1, kernel_size, 1, 1])
  num_channels = tf.shape(image)[-1]
  blur_h = tf.tile(blur_h, [1, 1, num_channels, 1])
  blur_v = tf.tile(blur_v, [1, 1, num_channels, 1])
  expand_batch_dim = image.shape.ndims == 3
  if expand_batch_dim:
    # Tensorflow requires batched input to convolutions, which we can fake with
    # an extra dimension.
    image = tf.expand_dims(image, axis=0)
  blurred = tf.nn.depthwise_conv2d(
      image, blur_h, strides=[1, 1, 1, 1], padding=padding)
  blurred = tf.nn.depthwise_conv2d(
      blurred, blur_v, strides=[1, 1, 1, 1], padding=padding)
  if expand_batch_dim:
    blurred = tf.squeeze(blurred, axis=0)
  return blurred
//  [python]
def gaussian_blur(image, kernel_size, sigma, padding='SAME'):
  """Blurs the given image with separable convolution.
  Args:
    image: Tensor of shape [height, width, channels] and dtype float to blur.
    kernel_size: Integer Tensor for the size of the blur kernel. This is should
      be an odd number. If it is an even number, the actual kernel size will be
      size + 1.
    sigma: Sigma value for gaussian operator.
    padding: Padding to use for the convolution. Typically 'SAME' or 'VALID'.
  Returns:
    A Tensor representing the blurred image.
  """
  radius = tf.cast(kernel_size / 2, tf.int32)
  kernel_size = radius * 2 + 1
  x = tf.cast(tf.range(-radius, radius + 1), tf.float32)
  blur_filter = tf.exp(
      -tf.pow(x, 2.0) / (2.0 * tf.pow(tf.cast(sigma, tf.float32), 2.0)))
  blur_filter /= tf.reduce_sum(blur_filter)
  # One vertical and one horizontal filter.
  blur_v = tf.reshape(blur_filter, [kernel_size, 1, 1, 1])
  blur_h = tf.reshape(blur_filter, [1, kernel_size, 1, 1])
  num_channels = tf.shape(image)[-1]
  blur_h = tf.tile(blur_h, [1, 1, num_channels, 1])
  blur_v = tf.tile(blur_v, [1, 1, num_channels, 1])
  expand_batch_dim = image.shape.ndims == 3
  if expand_batch_dim:
    # Tensorflow requires batched input to convolutions, which we can fake with
    # an extra dimension.
    image = tf.expand_dims(image, axis=0)
  blurred = tf.nn.depthwise_conv2d(
      image, blur_h, strides=[1, 1, 1, 1], padding=padding)
  blurred = tf.nn.depthwise_conv2d(
      blurred, blur_v, strides=[1, 1, 1, 1], padding=padding)
  if expand_batch_dim:
    blurred = tf.squeeze(blurred, axis=0)
  return blurred
//  [python]
def _using_motion_vector(depth,
                         translation,
                         rotation_angles,
                         intrinsic_mat,
                         intrinsic_mat_inv=None):
  """A helper for using_motion_vector. See docstring therein."""

  if translation.shape.ndims not in (2, 4):
    raise ValueError('\'translation\' should have rank 2 or 4, not %d' %
                     translation.shape.ndims)
  if translation.shape[-1] != 3:
    raise ValueError('translation\'s last dimension should be 3, not %d' %
                     translation.shape[1])
  if translation.shape.ndims == 2:
    translation = tf.expand_dims(tf.expand_dims(translation, 1), 1)

  _, height, width = tf.unstack(tf.shape(depth))
  grid = tf.squeeze(
      tf.stack(tf.meshgrid(tf.range(width), tf.range(height), (1,))), axis=3)
  grid = tf.to_float(grid)
  if intrinsic_mat_inv is None:
    intrinsic_mat_inv = tf.linalg.inv(intrinsic_mat)

  # Use the depth map and the inverse intrinsic matrix to generate a point
  # cloud xyz.
  xyz = tf.einsum('bij,jhw,bhw->bihw', intrinsic_mat_inv, grid, depth)

  # TPU pads aggressively tensors that have small dimensions. Therefore having
  # A rotation of the shape [....., 3, 3] would overflow the HBM memory. To
  # address this, we represnet the rotations is a 3x3 nested python tuple of
  # tf.Tensors (that is, we unroll the rotation matrix at the small dimensions).
  # The 3x3 matrix multiplication is now done in a python loop, and tensors with
  # small dimensions are avoided.
  unstacked_xyz = tf.unstack(xyz, axis=1)
  unstacked_rotation_matrix = transform_utils.unstacked_matrix_from_angles(
      *tf.unstack(rotation_angles, axis=-1))
  rank_diff = (
      unstacked_xyz[0].shape.ndims -
      unstacked_rotation_matrix[0][0].shape.ndims)

  def expand_to_needed_rank(t):
    for _ in range(rank_diff):
      t = tf.expand_dims(t, -1)
    return t

  unstacked_rotated_xyz = [0.0] * 3
  for i in range(3):
    for j in range(3):
      unstacked_rotated_xyz[i] += expand_to_needed_rank(
          unstacked_rotation_matrix[i][j]) * unstacked_xyz[j]
  rotated_xyz = tf.stack(unstacked_rotated_xyz, axis=1)

  # Project the transformed point cloud back to the camera plane.
  pcoords = tf.einsum('bij,bjhw->bihw', intrinsic_mat, rotated_xyz)

  projected_translation = tf.einsum('bij,bhwj->bihw', intrinsic_mat,
                                    translation)
  pcoords += projected_translation
  x, y, z = tf.unstack(pcoords, axis=1)
  return x / z, y / z, z
//  [python]
def _using_motion_vector(depth,
                         translation,
                         rotation_angles,
                         intrinsic_mat,
                         intrinsic_mat_inv=None):
  """A helper for using_motion_vector. See docstring therein."""

  if translation.shape.ndims not in (2, 4):
    raise ValueError('\'translation\' should have rank 2 or 4, not %d' %
                     translation.shape.ndims)
  if translation.shape[-1] != 3:
    raise ValueError('translation\'s last dimension should be 3, not %d' %
                     translation.shape[1])
  if translation.shape.ndims == 2:
    translation = tf.expand_dims(tf.expand_dims(translation, 1), 1)

  _, height, width = tf.unstack(tf.shape(depth))
  grid = tf.squeeze(
      tf.stack(tf.meshgrid(tf.range(width), tf.range(height), (1,))), axis=3)
  grid = tf.cast(grid, tf.float32)
  if intrinsic_mat_inv is None:
    intrinsic_mat_inv = tf.linalg.inv(intrinsic_mat)

  # Use the depth map and the inverse intrinsic matrix to generate a point
  # cloud xyz.
  xyz = tf.einsum('bij,jhw,bhw->bihw', intrinsic_mat_inv, grid, depth)

  # TPU pads aggressively tensors that have small dimensions. Therefore having
  # A rotation of the shape [....., 3, 3] would overflow the HBM memory. To
  # address this, we represnet the rotations is a 3x3 nested python tuple of
  # tf.Tensors (that is, we unroll the rotation matrix at the small dimensions).
  # The 3x3 matrix multiplication is now done in a python loop, and tensors with
  # small dimensions are avoided.
  unstacked_xyz = tf.unstack(xyz, axis=1)
  unstacked_rotation_matrix = transform_utils.unstacked_matrix_from_angles(
      *tf.unstack(rotation_angles, axis=-1))
  rank_diff = (
      unstacked_xyz[0].shape.ndims -
      unstacked_rotation_matrix[0][0].shape.ndims)

  def expand_to_needed_rank(t):
    for _ in range(rank_diff):
      t = tf.expand_dims(t, -1)
    return t

  unstacked_rotated_xyz = [0.0] * 3
  for i in range(3):
    for j in range(3):
      unstacked_rotated_xyz[i] += expand_to_needed_rank(
          unstacked_rotation_matrix[i][j]) * unstacked_xyz[j]
  rotated_xyz = tf.stack(unstacked_rotated_xyz, axis=1)

  # Project the transformed point cloud back to the camera plane.
  pcoords = tf.einsum('bij,bjhw->bihw', intrinsic_mat, rotated_xyz)

  projected_translation = tf.einsum('bij,bhwj->bihw', intrinsic_mat,
                                    translation)
  pcoords += projected_translation
  x, y, z = tf.unstack(pcoords, axis=1)
  return x / z, y / z, z
//  [python]
def calculate_frame_metrics(frame_labels, frame_predictions):
  """Calculate frame-based metrics."""
  frame_labels_bool = tf.cast(frame_labels, tf.bool)
  frame_predictions_bool = tf.cast(frame_predictions, tf.bool)

  frame_true_positives = tf.reduce_sum(tf.to_float(tf.logical_and(
      tf.equal(frame_labels_bool, True),
      tf.equal(frame_predictions_bool, True))))
  frame_false_positives = tf.reduce_sum(tf.to_float(tf.logical_and(
      tf.equal(frame_labels_bool, False),
      tf.equal(frame_predictions_bool, True))))
  frame_false_negatives = tf.reduce_sum(tf.to_float(tf.logical_and(
      tf.equal(frame_labels_bool, True),
      tf.equal(frame_predictions_bool, False))))
  frame_accuracy = (
      tf.reduce_sum(
          tf.to_float(tf.equal(frame_labels_bool, frame_predictions_bool))) /
      tf.cast(tf.size(frame_labels), tf.float32))

  frame_precision = tf.where(
      tf.greater(frame_true_positives + frame_false_positives, 0),
      tf.div(frame_true_positives,
             frame_true_positives + frame_false_positives),
      0)
  frame_recall = tf.where(
      tf.greater(frame_true_positives + frame_false_negatives, 0),
      tf.div(frame_true_positives,
             frame_true_positives + frame_false_negatives),
      0)
  frame_f1_score = f1_score(frame_precision, frame_recall)
  frame_accuracy_without_true_negatives = accuracy_without_true_negatives(
      frame_true_positives, frame_false_positives, frame_false_negatives)

  return {
      'true_positives': [frame_true_positives],
      'false_positives': [frame_false_positives],
      'false_negatives': [frame_false_negatives],
      'accuracy': [frame_accuracy],
      'accuracy_without_true_negatives': [
          frame_accuracy_without_true_negatives],
      'precision': [frame_precision],
      'recall': [frame_recall],
      'f1_score': [frame_f1_score],
  }
//  [python]
def calculate_frame_metrics(frame_labels, frame_predictions):
  """Calculate frame-based metrics."""
  frame_labels_bool = tf.cast(frame_labels, tf.bool)
  frame_predictions_bool = tf.cast(frame_predictions, tf.bool)

  frame_true_positives = tf.reduce_sum(tf.cast(tf.logical_and(
      tf.equal(frame_labels_bool, True),
      tf.equal(frame_predictions_bool, True)), tf.float32))
  frame_false_positives = tf.reduce_sum(tf.cast(tf.logical_and(
      tf.equal(frame_labels_bool, False),
      tf.equal(frame_predictions_bool, True)), tf.float32))
  frame_false_negatives = tf.reduce_sum(tf.cast(tf.logical_and(
      tf.equal(frame_labels_bool, True),
      tf.equal(frame_predictions_bool, False)), tf.float32))
  frame_accuracy = (
      tf.reduce_sum(
          tf.cast(tf.equal(frame_labels_bool, frame_predictions_bool), tf.float32)) /
      tf.cast(tf.size(frame_labels), tf.float32))

  frame_precision = tf.where(
      tf.greater(frame_true_positives + frame_false_positives, 0),
      tf.div(frame_true_positives,
             frame_true_positives + frame_false_positives),
      0)
  frame_recall = tf.where(
      tf.greater(frame_true_positives + frame_false_negatives, 0),
      tf.div(frame_true_positives,
             frame_true_positives + frame_false_negatives),
      0)
  frame_f1_score = f1_score(frame_precision, frame_recall)
  frame_accuracy_without_true_negatives = accuracy_without_true_negatives(
      frame_true_positives, frame_false_positives, frame_false_negatives)

  return {
      'true_positives': [frame_true_positives],
      'false_positives': [frame_false_positives],
      'false_negatives': [frame_false_negatives],
      'accuracy': [frame_accuracy],
      'accuracy_without_true_negatives': [
          frame_accuracy_without_true_negatives],
      'precision': [frame_precision],
      'recall': [frame_recall],
      'f1_score': [frame_f1_score],
  }
//  [python]
def get_global_step(hparams):
  """Returns the global optimization step."""
  step = tf.to_float(tf.train.get_or_create_global_step())
  multiplier = hparams.optimizer_multistep_accumulate_steps
  if not multiplier:
    return step
  return step / tf.to_float(multiplier)
//  [python]
def get_global_step(hparams):
  """Returns the global optimization step."""
  step = tf.cast(tf.train.get_or_create_global_step(), tf.float32)
  multiplier = hparams.optimizer_multistep_accumulate_steps
  if not multiplier:
    return step
  return step / tf.cast(multiplier, tf.float32)
//  [python]
def margin_logit_loss(model_logits, label, nb_classes=10, num_classes=None):
    """Computes difference between logit for `label` and next highest logit.

    The loss is high when `label` is unlikely (targeted by default).
    This follows the same interface as `loss_fn` for TensorOptimizer and
    projected_optimization, i.e. it returns a batch of loss values.
    """
    if num_classes is not None:
        warnings.warn(
            "`num_classes` is depreciated. Switch to `nb_classes`."
            " `num_classes` may be removed on or after 2019-04-23."
        )
        nb_classes = num_classes
        del num_classes
    if "int" in str(label.dtype):
        logit_mask = tf.one_hot(label, depth=nb_classes, axis=-1)
    else:
        logit_mask = label
    if "int" in str(logit_mask.dtype):
        logit_mask = tf.to_float(logit_mask)
    try:
        label_logits = reduce_sum(logit_mask * model_logits, axis=-1)
    except TypeError:
        raise TypeError(
            "Could not take row-wise dot product between "
            "logit mask, of dtype "
            + str(logit_mask.dtype)
            + " and model_logits, of dtype "
            + str(model_logits.dtype)
        )
    logits_with_target_label_neg_inf = model_logits - logit_mask * 99999
    highest_nonlabel_logits = reduce_max(logits_with_target_label_neg_inf, axis=-1)
    loss = highest_nonlabel_logits - label_logits
    return loss
//  [python]
def margin_logit_loss(model_logits, label, nb_classes=10, num_classes=None):
    """Computes difference between logit for `label` and next highest logit.

    The loss is high when `label` is unlikely (targeted by default).
    This follows the same interface as `loss_fn` for TensorOptimizer and
    projected_optimization, i.e. it returns a batch of loss values.
    """
    if num_classes is not None:
        warnings.warn(
            "`num_classes` is depreciated. Switch to `nb_classes`."
            " `num_classes` may be removed on or after 2019-04-23."
        )
        nb_classes = num_classes
        del num_classes
    if "int" in str(label.dtype):
        logit_mask = tf.one_hot(label, depth=nb_classes, axis=-1)
    else:
        logit_mask = label
    if "int" in str(logit_mask.dtype):
        logit_mask = tf.cast(logit_mask, tf.float32)
    try:
        label_logits = reduce_sum(logit_mask * model_logits, axis=-1)
    except TypeError:
        raise TypeError(
            "Could not take row-wise dot product between "
            "logit mask, of dtype "
            + str(logit_mask.dtype)
            + " and model_logits, of dtype "
            + str(model_logits.dtype)
        )
    logits_with_target_label_neg_inf = model_logits - logit_mask * 99999
    highest_nonlabel_logits = reduce_max(logits_with_target_label_neg_inf, axis=-1)
    loss = highest_nonlabel_logits - label_logits
    return loss
//  [python]
def learning_rate_schedule(base_learning_rate, num_examples):
  """Build learning rate schedule."""
  global_step = tf.train.get_or_create_global_step()
  warmup_steps = int(round(
      FLAGS.warmup_epochs * num_examples // FLAGS.train_batch_size))
  scaled_lr = base_learning_rate * FLAGS.train_batch_size / 256.
  learning_rate = (tf.to_float(global_step) / int(warmup_steps) * scaled_lr
                   if warmup_steps else scaled_lr)

  # Cosine decay learning rate schedule
  total_steps = get_train_steps(num_examples)
  learning_rate = tf.where(
      global_step < warmup_steps, learning_rate,
      tf.train.cosine_decay(
          scaled_lr,
          global_step - warmup_steps,
          total_steps - warmup_steps))

  return learning_rate
//  [python]
def learning_rate_schedule(base_learning_rate, num_examples):
  """Build learning rate schedule."""
  global_step = tf.train.get_or_create_global_step()
  warmup_steps = int(round(
      FLAGS.warmup_epochs * num_examples // FLAGS.train_batch_size))
  scaled_lr = base_learning_rate * FLAGS.train_batch_size / 256.
  learning_rate = (tf.cast(global_step, tf.float32) / int(warmup_steps) * scaled_lr
                   if warmup_steps else scaled_lr)

  # Cosine decay learning rate schedule
  total_steps = get_train_steps(num_examples)
  learning_rate = tf.where(
      global_step < warmup_steps, learning_rate,
      tf.train.cosine_decay(
          scaled_lr,
          global_step - warmup_steps,
          total_steps - warmup_steps))

  return learning_rate
//  [python]
def rouge_l_fscore(predictions, labels, **unused_kwargs):
  """ROUGE scores computation between labels and predictions.

  This is an approximate ROUGE scoring method since we do not glue word pieces
  or decode the ids and tokenize the output.

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

  Returns:
    rouge_l_fscore: approx rouge-l f1 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])
  rouge_l_f_score = tf.py_func(rouge_l_sentence_level, (outputs, labels),
                               tf.float32)
  return rouge_l_f_score, tf.constant(1.0)
//  [python]
def rouge_l_fscore(predictions, labels, **unused_kwargs):
  """ROUGE scores computation between labels and predictions.

  This is an approximate ROUGE scoring method since we do not glue word pieces
  or decode the ids and tokenize the output.

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

  Returns:
    rouge_l_fscore: approx rouge-l f1 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])
  rouge_l_f_score = tf.py_func(rouge_l_sentence_level, (outputs, labels),
                               tf.float32)
  return rouge_l_f_score, tf.constant(1.0)
//  [python]
def rouge_2_fscore(predictions, labels, **unused_kwargs):
  """ROUGE-2 F1 score computation between labels and predictions.

  This is an approximate ROUGE scoring method since we do not glue word pieces
  or decode the ids and tokenize the output.

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

  Returns:
    rouge2_fscore: approx rouge-2 f1 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])
  rouge_2_f_score = tf.py_func(rouge_n, (outputs, labels), tf.float32)
  return rouge_2_f_score, tf.constant(1.0)
//  [python]
def rouge_2_fscore(predictions, labels, **unused_kwargs):
  """ROUGE-2 F1 score computation between labels and predictions.

  This is an approximate ROUGE scoring method since we do not glue word pieces
  or decode the ids and tokenize the output.

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

  Returns:
    rouge2_fscore: approx rouge-2 f1 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])
  rouge_2_f_score = tf.py_func(rouge_n, (outputs, labels), tf.float32)
  return rouge_2_f_score, tf.constant(1.0)
//  [python]
def get_candidate_labels(self, candidate_starts, candidate_ends, labeled_starts, labeled_ends, labels):
        same_start = tf.equal(tf.expand_dims(labeled_starts, 1), tf.expand_dims(
            candidate_starts, 0))  # [num_labeled, num_candidates]
        same_end = tf.equal(tf.expand_dims(labeled_ends, 1), tf.expand_dims(
            candidate_ends, 0))  # [num_labeled, num_candidates]
        # [num_labeled, num_candidates]
        same_span = tf.logical_and(same_start, same_end)
        candidate_labels = tf.matmul(tf.expand_dims(
            labels, 0), tf.to_int32(same_span))  # [1, num_candidates]
        candidate_labels = tf.squeeze(candidate_labels, 0)  # [num_candidates]
        return candidate_labels
//  [python]
def get_candidate_labels(self, candidate_starts, candidate_ends, labeled_starts, labeled_ends, labels):
        same_start = tf.equal(tf.expand_dims(labeled_starts, 1), tf.expand_dims(
            candidate_starts, 0))  # [num_labeled, num_candidates]
        same_end = tf.equal(tf.expand_dims(labeled_ends, 1), tf.expand_dims(
            candidate_ends, 0))  # [num_labeled, num_candidates]
        # [num_labeled, num_candidates]
        same_span = tf.logical_and(same_start, same_end)
        candidate_labels = tf.matmul(tf.expand_dims(
            labels, 0), tf.cast(same_span, tf.int32))  # [1, num_candidates]
        candidate_labels = tf.squeeze(candidate_labels, 0)  # [num_candidates]
        return candidate_labels
//  [python]
def bucket_distance(self, distances):
        """
        Places the given values (designed for distances) into 10 semi-logscale buckets:
        [0, 1, 2, 3, 4, 5-7, 8-15, 16-31, 32-63, 64+].
        """
        logspace_idx = tf.to_int32(
            tf.floor(tf.log(tf.to_float(distances))/math.log(2))) + 3
        use_identity = tf.to_int32(distances <= 4)
        combined_idx = use_identity * distances + \
            (1 - use_identity) * logspace_idx
        return tf.clip_by_value(combined_idx, 0, 9)
//  [python]
def bucket_distance(self, distances):
        """
        Places the given values (designed for distances) into 10 semi-logscale buckets:
        [0, 1, 2, 3, 4, 5-7, 8-15, 16-31, 32-63, 64+].
        """
        logspace_idx = tf.cast(
            tf.floor(tf.math.log(tf.cast(distances, tf.float32))/math.log(2)), tf.int32) + 3
        use_identity = tf.cast(distances <= 4, tf.int32)
        combined_idx = use_identity * distances + \
            (1 - use_identity) * logspace_idx
        return tf.clip_by_value(combined_idx, 0, 9)
//  [python]
def spm(
    x,
    model,
    y=None,
    n_samples=None,
    dx_min=-0.1,
    dx_max=0.1,
    n_dxs=5,
    dy_min=-0.1,
    dy_max=0.1,
    n_dys=5,
    angle_min=-30,
    angle_max=30,
    n_angles=31,
    black_border_size=0,
):
    """
    TensorFlow implementation of the Spatial Transformation Method.
    :return: a tensor for the adversarial example
    """
    if y is None:
        preds = model.get_probs(x)
        # Using model predictions as ground truth to avoid label leaking
        preds_max = reduce_max(preds, 1, keepdims=True)
        y = tf.to_float(tf.equal(preds, preds_max))
        y = tf.stop_gradient(y)
        del preds
    y = y / reduce_sum(y, 1, keepdims=True)

    # Define the range of transformations
    dxs = np.linspace(dx_min, dx_max, n_dxs)
    dys = np.linspace(dy_min, dy_max, n_dys)
    angles = np.linspace(angle_min, angle_max, n_angles)

    if n_samples is None:
        import itertools

        transforms = list(itertools.product(*[dxs, dys, angles]))
    else:
        sampled_dxs = np.random.choice(dxs, n_samples)
        sampled_dys = np.random.choice(dys, n_samples)
        sampled_angles = np.random.choice(angles, n_samples)
        transforms = zip(sampled_dxs, sampled_dys, sampled_angles)
    transformed_ims = parallel_apply_transformations(x, transforms, black_border_size)

    def _compute_xent(x):
        preds = model.get_logits(x)
        return tf.nn.softmax_cross_entropy_with_logits_v2(labels=y, logits=preds)

    all_xents = tf.map_fn(
        _compute_xent, transformed_ims, parallel_iterations=1
    )  # Must be 1 to avoid keras race conditions

    # Return the adv_x with worst accuracy

    # all_xents is n_total_samples x batch_size (SB)
    all_xents = tf.stack(all_xents)  # SB

    # We want the worst case sample, with the largest xent_loss
    worst_sample_idx = tf.argmax(all_xents, axis=0)  # B

    batch_size = tf.shape(x)[0]
    keys = tf.stack(
        [tf.range(batch_size, dtype=tf.int32), tf.cast(worst_sample_idx, tf.int32)],
        axis=1,
    )
    transformed_ims_bshwc = tf.einsum("sbhwc->bshwc", transformed_ims)
    after_lookup = tf.gather_nd(transformed_ims_bshwc, keys)  # BHWC
    return after_lookup
//  [python]
def spm(
    x,
    model,
    y=None,
    n_samples=None,
    dx_min=-0.1,
    dx_max=0.1,
    n_dxs=5,
    dy_min=-0.1,
    dy_max=0.1,
    n_dys=5,
    angle_min=-30,
    angle_max=30,
    n_angles=31,
    black_border_size=0,
):
    """
    TensorFlow implementation of the Spatial Transformation Method.
    :return: a tensor for the adversarial example
    """
    if y is None:
        preds = model.get_probs(x)
        # Using model predictions as ground truth to avoid label leaking
        preds_max = reduce_max(preds, 1, keepdims=True)
        y = tf.cast(tf.equal(preds, preds_max), tf.float32)
        y = tf.stop_gradient(y)
        del preds
    y = y / reduce_sum(y, 1, keepdims=True)

    # Define the range of transformations
    dxs = np.linspace(dx_min, dx_max, n_dxs)
    dys = np.linspace(dy_min, dy_max, n_dys)
    angles = np.linspace(angle_min, angle_max, n_angles)

    if n_samples is None:
        import itertools

        transforms = list(itertools.product(*[dxs, dys, angles]))
    else:
        sampled_dxs = np.random.choice(dxs, n_samples)
        sampled_dys = np.random.choice(dys, n_samples)
        sampled_angles = np.random.choice(angles, n_samples)
        transforms = zip(sampled_dxs, sampled_dys, sampled_angles)
    transformed_ims = parallel_apply_transformations(x, transforms, black_border_size)

    def _compute_xent(x):
        preds = model.get_logits(x)
        return tf.nn.softmax_cross_entropy_with_logits_v2(labels=y, logits=preds)

    all_xents = tf.map_fn(
        _compute_xent, transformed_ims, parallel_iterations=1
    )  # Must be 1 to avoid keras race conditions

    # Return the adv_x with worst accuracy

    # all_xents is n_total_samples x batch_size (SB)
    all_xents = tf.stack(all_xents)  # SB

    # We want the worst case sample, with the largest xent_loss
    worst_sample_idx = tf.argmax(all_xents, axis=0)  # B

    batch_size = tf.shape(x)[0]
    keys = tf.stack(
        [tf.range(batch_size, dtype=tf.int32), tf.cast(worst_sample_idx, tf.int32)],
        axis=1,
    )
    transformed_ims_bshwc = tf.einsum("sbhwc->bshwc", transformed_ims)
    after_lookup = tf.gather_nd(transformed_ims_bshwc, keys)  # BHWC
    return after_lookup
//  [python]
def rouge_l_fscore(predictions, labels, **unused_kwargs):
  """ROUGE scores computation between labels and predictions.
  This is an approximate ROUGE scoring method since we do not glue word pieces
  or decode the ids and tokenize the output.
  Args:
    predictions: tensor, model predictions
    labels: tensor, gold output.
  Returns:
    rouge_l_fscore: approx rouge-l f1 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])
  rouge_l_f_score = tf.py_func(rouge_l_sentence_level, (outputs, labels),
                               tf.float32)
  return rouge_l_f_score, tf.constant(1.0)
//  [python]
def rouge_l_fscore(predictions, labels, **unused_kwargs):
  """ROUGE scores computation between labels and predictions.
  This is an approximate ROUGE scoring method since we do not glue word pieces
  or decode the ids and tokenize the output.
  Args:
    predictions: tensor, model predictions
    labels: tensor, gold output.
  Returns:
    rouge_l_fscore: approx rouge-l f1 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])
  rouge_l_f_score = tf.py_func(rouge_l_sentence_level, (outputs, labels),
                               tf.float32)
  return rouge_l_f_score, tf.constant(1.0)
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, top_span_mention_scores, c):
        k = util.shape(top_span_emb, 0)
        top_span_range = tf.range(k)  # [k]
        antecedent_offsets = tf.expand_dims(
            top_span_range, 1) - tf.expand_dims(top_span_range, 0)  # [k, k]
        antecedents_mask = antecedent_offsets >= 1  # [k, k]
        fast_antecedent_scores = tf.expand_dims(
            top_span_mention_scores, 1) + tf.expand_dims(top_span_mention_scores, 0)  # [k, k]
        # [k, k]
        fast_antecedent_scores += tf.log(tf.to_float(antecedents_mask))
        # [k, k]
        fast_antecedent_scores += self.get_fast_antecedent_scores(top_span_emb)

        _, top_antecedents = tf.nn.top_k(
            fast_antecedent_scores, c, sorted=False)  # [k, c]
        top_antecedents_mask = util.batch_gather(
            antecedents_mask, top_antecedents)  # [k, c]
        top_fast_antecedent_scores = util.batch_gather(
            fast_antecedent_scores, top_antecedents)  # [k, c]
        top_antecedent_offsets = util.batch_gather(
            antecedent_offsets, top_antecedents)  # [k, c]
        return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, top_span_mention_scores, c):
        k = util.shape(top_span_emb, 0)
        top_span_range = tf.range(k)  # [k]
        antecedent_offsets = tf.expand_dims(
            top_span_range, 1) - tf.expand_dims(top_span_range, 0)  # [k, k]
        antecedents_mask = antecedent_offsets >= 1  # [k, k]
        fast_antecedent_scores = tf.expand_dims(
            top_span_mention_scores, 1) + tf.expand_dims(top_span_mention_scores, 0)  # [k, k]
        # [k, k]
        fast_antecedent_scores += tf.log(tf.cast(antecedents_mask, tf.float32))
        # [k, k]
        fast_antecedent_scores += self.get_fast_antecedent_scores(top_span_emb)

        _, top_antecedents = tf.nn.top_k(
            fast_antecedent_scores, c, sorted=False)  # [k, c]
        top_antecedents_mask = util.batch_gather(
            antecedents_mask, top_antecedents)  # [k, c]
        top_fast_antecedent_scores = util.batch_gather(
            fast_antecedent_scores, top_antecedents)  # [k, c]
        top_antecedent_offsets = util.batch_gather(
            antecedent_offsets, top_antecedents)  # [k, c]
        return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def build_trainer(self, child_model):
    child_model.build_valid_rl()
    self.valid_acc = (tf.to_float(child_model.valid_shuffle_acc) /
                      tf.to_float(child_model.batch_size))
    self.reward = self.valid_acc

    if self.entropy_weight is not None:
      self.reward += self.entropy_weight * self.sample_entropy

    self.sample_log_prob = tf.reduce_sum(self.sample_log_prob)
    self.baseline = tf.Variable(0.0, dtype=tf.float32, trainable=False)
    baseline_update = tf.assign_sub(
      self.baseline, (1 - self.bl_dec) * (self.baseline - self.reward))

    with tf.control_dependencies([baseline_update]):
      self.reward = tf.identity(self.reward)

    self.loss = self.sample_log_prob * (self.reward - self.baseline)
    self.train_step = tf.Variable(0, dtype=tf.int32, trainable=False, name="train_step")

    tf_variables = [var for var in tf.trainable_variables() if var.name.startswith(self.name)]
    print("-" * 80)
    for var in tf_variables:
      print(var)

    self.train_op, self.lr, self.grad_norm, self.optimizer = get_train_ops(
      self.loss,
      tf_variables,
      self.train_step,
      clip_mode=self.clip_mode,
      grad_bound=self.grad_bound,
      l2_reg=self.l2_reg,
      lr_init=self.lr_init,
      lr_dec_start=self.lr_dec_start,
      lr_dec_every=self.lr_dec_every,
      lr_dec_rate=self.lr_dec_rate,
      optim_algo=self.optim_algo,
      sync_replicas=self.sync_replicas,
      num_aggregate=self.num_aggregate,
      num_replicas=self.num_replicas)

    self.skip_rate = tf.constant(0.0, dtype=tf.float32)
//  [python]
def build_trainer(self, child_model):
    child_model.build_valid_rl()
    self.valid_acc = (tf.cast(child_model.valid_shuffle_acc, tf.float32) /
                      tf.cast(child_model.batch_size, tf.float32))
    self.reward = self.valid_acc

    if self.entropy_weight is not None:
      self.reward += self.entropy_weight * self.sample_entropy

    self.sample_log_prob = tf.reduce_sum(self.sample_log_prob)
    self.baseline = tf.Variable(0.0, dtype=tf.float32, trainable=False)
    baseline_update = tf.assign_sub(
      self.baseline, (1 - self.bl_dec) * (self.baseline - self.reward))

    with tf.control_dependencies([baseline_update]):
      self.reward = tf.identity(self.reward)

    self.loss = self.sample_log_prob * (self.reward - self.baseline)
    self.train_step = tf.Variable(0, dtype=tf.int32, trainable=False, name="train_step")

    tf_variables = [var for var in tf.trainable_variables() if var.name.startswith(self.name)]
    print("-" * 80)
    for var in tf_variables:
      print(var)

    self.train_op, self.lr, self.grad_norm, self.optimizer = get_train_ops(
      self.loss,
      tf_variables,
      self.train_step,
      clip_mode=self.clip_mode,
      grad_bound=self.grad_bound,
      l2_reg=self.l2_reg,
      lr_init=self.lr_init,
      lr_dec_start=self.lr_dec_start,
      lr_dec_every=self.lr_dec_every,
      lr_dec_rate=self.lr_dec_rate,
      optim_algo=self.optim_algo,
      sync_replicas=self.sync_replicas,
      num_aggregate=self.num_aggregate,
      num_replicas=self.num_replicas)

    self.skip_rate = tf.constant(0.0, dtype=tf.float32)
//  [python]
def weight(w, sparsity):
  """Weight-level magnitude pruning."""
  w_shape = common_layers.shape_list(w)
  k = int(np.prod(w_shape[:-1]))
  count = tf.to_int32(k * sparsity)
  mask = common_layers.weight_targeting(w, count)
  return (1 - mask) * w
//  [python]
def weight(w, sparsity):
  """Weight-level magnitude pruning."""
  w_shape = common_layers.shape_list(w)
  k = int(np.prod(w_shape[:-1]))
  count = tf.cast(k * sparsity, tf.int32)
  mask = common_layers.weight_targeting(w, count)
  return (1 - mask) * w