jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _smallest_size_at_least(height, width, smallest_side):
  """Computes new shape with the smallest side equal to `smallest_side`.
  Computes new shape with the smallest side equal to `smallest_side` while
  preserving the original aspect ratio.
  Args:
    height: an int32 scalar tensor indicating the current height.
    width: an int32 scalar tensor indicating the current width.
    smallest_side: A python integer or scalar `Tensor` indicating the size of
      the smallest side after resize.
  Returns:
    new_height: an int32 scalar tensor indicating the new height.
    new_width: and int32 scalar tensor indicating the new width.
  """
  smallest_side = tf.convert_to_tensor(smallest_side, dtype=tf.int32)

  height = tf.to_float(height)
  width = tf.to_float(width)
  smallest_side = tf.to_float(smallest_side)

  scale = tf.cond(tf.greater(height, width),
                  lambda: smallest_side / width,
                  lambda: smallest_side / height)
  new_height = tf.to_int32(tf.rint(height * scale))
  new_width = tf.to_int32(tf.rint(width * scale))
  return new_height, new_width
//  [python]
def _smallest_size_at_least(height, width, smallest_side):
  """Computes new shape with the smallest side equal to `smallest_side`.
  Computes new shape with the smallest side equal to `smallest_side` while
  preserving the original aspect ratio.
  Args:
    height: an int32 scalar tensor indicating the current height.
    width: an int32 scalar tensor indicating the current width.
    smallest_side: A python integer or scalar `Tensor` indicating the size of
      the smallest side after resize.
  Returns:
    new_height: an int32 scalar tensor indicating the new height.
    new_width: and int32 scalar tensor indicating the new width.
  """
  smallest_side = tf.convert_to_tensor(smallest_side, dtype=tf.int32)

  height = tf.cast(height, tf.float32)
  width = tf.cast(width, tf.float32)
  smallest_side = tf.cast(smallest_side, tf.float32)

  scale = tf.cond(tf.greater(height, width),
                  lambda: smallest_side / width,
                  lambda: smallest_side / height)
  new_height = tf.to_int32(tf.rint(height * scale))
  new_width = tf.to_int32(tf.rint(width * scale))
  return new_height, new_width
//  [python]
def _parse_example(self, example_proto):
    d = tf.parse_single_example(example_proto, self._feature_description)

    img = tf.decode_raw(d['image_raw'], tf.uint8)
    img = tf.reshape(img, self._img_shape)
    d['image'] = tf.to_float(img) / 255.
    del d['image_raw']

    return d
//  [python]
def _parse_example(self, example_proto):
    d = tf.parse_single_example(example_proto, self._feature_description)

    img = tf.decode_raw(d['image_raw'], tf.uint8)
    img = tf.reshape(img, self._img_shape)
    d['image'] = tf.cast(img, tf.float32) / 255.
    del d['image_raw']

    return d
//  [python]
def create_test_input(batch, height, width, channels):
  """Create test input tensor."""
  if None in [batch, height, width, channels]:
    return tf.placeholder(tf.float32, (batch, height, width, channels))
  else:
    return tf.to_float(
        np.tile(
            np.reshape(
                np.reshape(np.arange(height), [height, 1]) +
                np.reshape(np.arange(width), [1, width]),
                [1, height, width, 1]),
            [batch, 1, 1, channels]))
//  [python]
def create_test_input(batch, height, width, channels):
  """Create test input tensor."""
  if None in [batch, height, width, channels]:
    return tf.placeholder(tf.float32, (batch, height, width, channels))
  else:
    return tf.cast(
        np.tile(
            np.reshape(
                np.reshape(np.arange(height), [height, 1]) +
                np.reshape(np.arange(width), [1, width]),
                [1, height, width, 1]),
            [batch, 1, 1, channels]), tf.float32)
//  [python]
def preprocess_for_train(image,
                         output_height,
                         output_width,
                         resize_side_min=_RESIZE_SIDE_MIN,
                         resize_side_max=_RESIZE_SIDE_MAX,
                         use_grayscale=False):
  """Preprocesses the given image for training.
  Note that the actual resizing scale is sampled from
    [`resize_size_min`, `resize_size_max`].
  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    resize_side_min: The lower bound for the smallest side of the image for
      aspect-preserving resizing.
    resize_side_max: The upper bound for the smallest side of the image for
      aspect-preserving resizing.
    use_grayscale: Whether to convert the image from RGB to grayscale.
  Returns:
    A preprocessed image.
  """
  resize_side = tf.random_uniform(
      [], minval=resize_side_min, maxval=resize_side_max+1, dtype=tf.int32)

  image = _aspect_preserving_resize(image, resize_side)
  image = _random_crop([image], output_height, output_width)[0]
  image.set_shape([output_height, output_width, 3])
  image = tf.to_float(image)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)
  image = tf.image.random_flip_left_right(image)
  return _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
//  [python]
def preprocess_for_train(image,
                         output_height,
                         output_width,
                         resize_side_min=_RESIZE_SIDE_MIN,
                         resize_side_max=_RESIZE_SIDE_MAX,
                         use_grayscale=False):
  """Preprocesses the given image for training.
  Note that the actual resizing scale is sampled from
    [`resize_size_min`, `resize_size_max`].
  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    resize_side_min: The lower bound for the smallest side of the image for
      aspect-preserving resizing.
    resize_side_max: The upper bound for the smallest side of the image for
      aspect-preserving resizing.
    use_grayscale: Whether to convert the image from RGB to grayscale.
  Returns:
    A preprocessed image.
  """
  resize_side = tf.random_uniform(
      [], minval=resize_side_min, maxval=resize_side_max+1, dtype=tf.int32)

  image = _aspect_preserving_resize(image, resize_side)
  image = _random_crop([image], output_height, output_width)[0]
  image.set_shape([output_height, output_width, 3])
  image = tf.cast(image, tf.float32)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)
  image = tf.image.random_flip_left_right(image)
  return _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
//  [python]
def preprocess_example(self, example, mode, _):
    # Just resize with area.
    if self._was_reversed:
      example["inputs"] = tf.to_int64(
          tf.image.resize_images(example["inputs"], self.rescale_size,
                                 tf.image.ResizeMethod.AREA))
    else:
      example = imagenet_preprocess_example(example, mode)
      example["inputs"] = tf.to_int64(
          tf.image.resize_images(example["inputs"], self.rescale_size))
    return example
//  [python]
def preprocess_example(self, example, mode, _):
    # Just resize with area.
    if self._was_reversed:
      example["inputs"] = tf.cast(
          tf.image.resize_images(example["inputs"], self.rescale_size,
                                 tf.image.ResizeMethod.AREA), tf.int64)
    else:
      example = imagenet_preprocess_example(example, mode)
      example["inputs"] = tf.cast(
          tf.image.resize_images(example["inputs"], self.rescale_size), tf.int64)
    return example
//  [python]
def preprocess_example(self, example, mode, unused_hparams):
    example["inputs"].set_shape([_IMAGENET_SMALL_IMAGE_SIZE,
                                 _IMAGENET_SMALL_IMAGE_SIZE, 3])
    example["inputs"] = tf.to_int64(example["inputs"])
    return example
//  [python]
def preprocess_example(self, example, mode, unused_hparams):
    example["inputs"].set_shape([_IMAGENET_SMALL_IMAGE_SIZE,
                                 _IMAGENET_SMALL_IMAGE_SIZE, 3])
    example["inputs"] = tf.cast(example["inputs"], tf.int64)
    return example
//  [python]
def preprocess_example(self, example, mode, unused_hparams):
    example["inputs"].set_shape([_IMAGENET_MEDIUM_IMAGE_SIZE,
                                 _IMAGENET_MEDIUM_IMAGE_SIZE, 3])
    example["inputs"] = tf.to_int64(example["inputs"])
    return example
//  [python]
def preprocess_example(self, example, mode, unused_hparams):
    example["inputs"].set_shape([_IMAGENET_MEDIUM_IMAGE_SIZE,
                                 _IMAGENET_MEDIUM_IMAGE_SIZE, 3])
    example["inputs"] = tf.cast(example["inputs"], tf.int64)
    return example
//  [python]
def preprocess_example(self, example, mode, unused_hparams):
    example["inputs"].set_shape(
        [_IMAGENET_MEDIUM_IMAGE_SIZE, _IMAGENET_MEDIUM_IMAGE_SIZE, 3])
    example["inputs"] = tf.to_int64(example["inputs"])
    example["inputs"] = tf.reshape(example["inputs"], (-1,))

    del example["targets"]  # Ensure unconditional generation

    return example
//  [python]
def preprocess_example(self, example, mode, unused_hparams):
    example["inputs"].set_shape(
        [_IMAGENET_MEDIUM_IMAGE_SIZE, _IMAGENET_MEDIUM_IMAGE_SIZE, 3])
    example["inputs"] = tf.cast(example["inputs"], tf.int64)
    example["inputs"] = tf.reshape(example["inputs"], (-1,))

    del example["targets"]  # Ensure unconditional generation

    return example
//  [python]
def preprocess_for_train(image, image_size=224, normalize=True):
  """Preprocesses the given image for evaluation.

  Args:
    image: `Tensor` representing an image of arbitrary size.
    image_size: int, how large the output image should be.
    normalize: bool, if True the image is normalized.

  Returns:
    A preprocessed image `Tensor`.
  """
  if normalize: image = tf.to_float(image) / 255.0
  image = _random_crop(image, image_size)
  if normalize: image = _normalize(image)
  image = _flip(image)
  image = tf.reshape(image, [image_size, image_size, 3])
  return image
//  [python]
def preprocess_for_train(image, image_size=224, normalize=True):
  """Preprocesses the given image for evaluation.

  Args:
    image: `Tensor` representing an image of arbitrary size.
    image_size: int, how large the output image should be.
    normalize: bool, if True the image is normalized.

  Returns:
    A preprocessed image `Tensor`.
  """
  if normalize: image = tf.cast(image, tf.float32) / 255.0
  image = _random_crop(image, image_size)
  if normalize: image = _normalize(image)
  image = _flip(image)
  image = tf.reshape(image, [image_size, image_size, 3])
  return image
//  [python]
def preprocess_for_eval(image, image_size=224, normalize=True):
  """Preprocesses the given image for evaluation.

  Args:
    image: `Tensor` representing an image of arbitrary size.
    image_size: int, how large the output image should be.
    normalize: bool, if True the image is normalized.

  Returns:
    A preprocessed image `Tensor`.
  """
  if normalize: image = tf.to_float(image) / 255.0
  image = _do_scale(image, image_size + 32)
  if normalize: image = _normalize(image)
  image = _center_crop(image, image_size)
  image = tf.reshape(image, [image_size, image_size, 3])
  return image
//  [python]
def preprocess_for_eval(image, image_size=224, normalize=True):
  """Preprocesses the given image for evaluation.

  Args:
    image: `Tensor` representing an image of arbitrary size.
    image_size: int, how large the output image should be.
    normalize: bool, if True the image is normalized.

  Returns:
    A preprocessed image `Tensor`.
  """
  if normalize: image = tf.cast(image, tf.float32) / 255.0
  image = _do_scale(image, image_size + 32)
  if normalize: image = _normalize(image)
  image = _center_crop(image, image_size)
  image = tf.reshape(image, [image_size, image_size, 3])
  return image
//  [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

    normalize = tf.to_float(self.num_layers * (self.num_layers - 1) / 2)
    self.skip_rate = tf.to_float(self.skip_count) / normalize

    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)
    if self.skip_weight is not None:
      self.loss += self.skip_weight * self.skip_penaltys

    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)
//  [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

    normalize = tf.cast(self.num_layers * (self.num_layers - 1) / 2, tf.float32)
    self.skip_rate = tf.cast(self.skip_count, tf.float32) / normalize

    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)
    if self.skip_weight is not None:
      self.loss += self.skip_weight * self.skip_penaltys

    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)
//  [python]
def _get_infer_maximum_iterations(self, hparams, source_sequence_length):
    """Maximum decoding steps at inference time."""
    if hparams.tgt_max_len_infer:
      maximum_iterations = hparams.tgt_max_len_infer
      utils.print_out("  decoding maximum_iterations %d" % maximum_iterations)
    else:
      # TODO(thangluong): add decoding_length_factor flag
      decoding_length_factor = 2.0
      max_encoder_length = tf.reduce_max(source_sequence_length)
      maximum_iterations = tf.to_int32(tf.round(
          tf.to_float(max_encoder_length) * decoding_length_factor))
    return maximum_iterations
//  [python]
def _get_infer_maximum_iterations(self, hparams, source_sequence_length):
    """Maximum decoding steps at inference time."""
    if hparams.tgt_max_len_infer:
      maximum_iterations = hparams.tgt_max_len_infer
      utils.print_out("  decoding maximum_iterations %d" % maximum_iterations)
    else:
      # TODO(thangluong): add decoding_length_factor flag
      decoding_length_factor = 2.0
      max_encoder_length = tf.reduce_max(source_sequence_length)
      maximum_iterations = tf.cast(tf.round(
          tf.cast(max_encoder_length, tf.float32) * decoding_length_factor), tf.int32)
    return maximum_iterations
//  [python]
def _compute_loss(self, logits, decoder_cell_outputs):
    """Compute optimization loss."""
    target_output = self.iterator.target_output
    if self.time_major:
      target_output = tf.transpose(target_output)
    max_time = self.get_max_time(target_output)

    crossent = self._softmax_cross_entropy_loss(
        logits, decoder_cell_outputs, target_output)

    target_weights = tf.sequence_mask(
        self.iterator.target_sequence_length, max_time, dtype=self.dtype)
    if self.time_major:
      target_weights = tf.transpose(target_weights)

    loss = tf.reduce_sum(
        crossent * target_weights) / tf.to_float(self.batch_size)
    return loss
//  [python]
def _compute_loss(self, logits, decoder_cell_outputs):
    """Compute optimization loss."""
    target_output = self.iterator.target_output
    if self.time_major:
      target_output = tf.transpose(target_output)
    max_time = self.get_max_time(target_output)

    crossent = self._softmax_cross_entropy_loss(
        logits, decoder_cell_outputs, target_output)

    target_weights = tf.sequence_mask(
        self.iterator.target_sequence_length, max_time, dtype=self.dtype)
    if self.time_major:
      target_weights = tf.transpose(target_weights)

    loss = tf.reduce_sum(
        crossent * target_weights) / tf.cast(self.batch_size, dtype=self.dtype)
    return loss
//  [python]
def _crop(image, offset_height, offset_width, crop_height, crop_width):
  """Crops the given image using the provided offsets and sizes.

  Note that the method doesn't assume we know the input image size but it does
  assume we know the input image rank.

  Args:
    image: `Tensor` image of shape [height, width, channels].
    offset_height: `Tensor` indicating the height offset.
    offset_width: `Tensor` indicating the width offset.
    crop_height: the height of the cropped image.
    crop_width: the width of the cropped image.

  Returns:
    the cropped (and resized) image.

  Raises:
    InvalidArgumentError: if the rank is not 3 or if the image dimensions are
      less than the crop size.
  """
  original_shape = tf.shape(image)

  rank_assertion = tf.Assert(
      tf.equal(tf.rank(image), 3), ["Rank of image must be equal to 3."])
  with tf.control_dependencies([rank_assertion]):
    cropped_shape = tf.stack([crop_height, crop_width, original_shape[2]])

  size_assertion = tf.Assert(
      tf.logical_and(
          tf.greater_equal(original_shape[0], crop_height),
          tf.greater_equal(original_shape[1], crop_width)),
      ["Crop size greater than the image size."])

  offsets = tf.to_int32(tf.stack([offset_height, offset_width, 0]))

  # Use tf.slice instead of crop_to_bounding box as it accepts tensors to
  # define the crop size.
  with tf.control_dependencies([size_assertion]):
    image = tf.slice(image, offsets, cropped_shape)
  return tf.reshape(image, cropped_shape)
//  [python]
def _crop(image, offset_height, offset_width, crop_height, crop_width):
  """Crops the given image using the provided offsets and sizes.

  Note that the method doesn't assume we know the input image size but it does
  assume we know the input image rank.

  Args:
    image: `Tensor` image of shape [height, width, channels].
    offset_height: `Tensor` indicating the height offset.
    offset_width: `Tensor` indicating the width offset.
    crop_height: the height of the cropped image.
    crop_width: the width of the cropped image.

  Returns:
    the cropped (and resized) image.

  Raises:
    InvalidArgumentError: if the rank is not 3 or if the image dimensions are
      less than the crop size.
  """
  original_shape = tf.shape(image)

  rank_assertion = tf.Assert(
      tf.equal(tf.rank(image), 3), ["Rank of image must be equal to 3."])
  with tf.control_dependencies([rank_assertion]):
    cropped_shape = tf.stack([crop_height, crop_width, original_shape[2]])

  size_assertion = tf.Assert(
      tf.logical_and(
          tf.greater_equal(original_shape[0], crop_height),
          tf.greater_equal(original_shape[1], crop_width)),
      ["Crop size greater than the image size."])

  offsets = tf.cast(tf.stack([offset_height, offset_width, 0]), tf.int32)

  # Use tf.slice instead of crop_to_bounding box as it accepts tensors to
  # define the crop size.
  with tf.control_dependencies([size_assertion]):
    image = tf.slice(image, offsets, cropped_shape)
  return tf.reshape(image, cropped_shape)
//  [python]
def build_model(self):

		with tf.device('/gpu:0'):
			# Inputs.
			self.represenation_input, self.label_input = self.model_inputs()

			# Optional NN for latent representations.
			self.interm = self.feature_extractor(inputs=self.represenation_input, use=True)

			# Attention.
			self.weights = self.attention(inputs=self.interm)

			# Classifier.
			self.prob, self.z = self.classifier(interm=self.interm, weights=self.weights)
			self.clas = tf.to_float(tf.math.greater(self.prob, 0.5))
			
			# Losses.
			self.loss = self.loss(label=self.label_input, prob=self.prob)

			# Optimizers.
			self.trainer  = self.optimization()
//  [python]
def build_model(self):

		with tf.device('/gpu:0'):
			# Inputs.
			self.represenation_input, self.label_input = self.model_inputs()

			# Optional NN for latent representations.
			self.interm = self.feature_extractor(inputs=self.represenation_input, use=True)

			# Attention.
			self.weights = self.attention(inputs=self.interm)

			# Classifier.
			self.prob, self.z = self.classifier(interm=self.interm, weights=self.weights)
			self.clas = tf.cast(tf.math.greater(self.prob, 0.5), tf.float32)
			
			# Losses.
			self.loss = self.loss(label=self.label_input, prob=self.prob)

			# Optimizers.
			self.trainer  = self.optimization()
//  [python]
def discrete_logistics_likelihood(loc, scale, obs, side_censored=True):
    from tensorflow_probability import distributions as tfd
    x_lo = obs - 1/255/2
    x_hi = obs + 1/255/2
    if side_censored:
        x_lo = tf.where(x_lo>=0, x_lo, tf.to_float(-1000)*tf.ones_like(x_lo))
        x_hi = tf.where(x_hi<=1, x_hi, tf.to_float(+1000)*tf.ones_like(x_hi))
    dist = tfd.Logistic(tf.to_double(loc), tf.to_double(scale))
    # tf.where doesn't backpropagate correctly when there's NaN in the unused branch. So instead of
    # ret = tf.where(
    #     tf.math.abs(obs - loc) <= 2*scale,
    #     tf.math.log(dist.cdf(tf.to_double(x_hi)) - dist.cdf(tf.to_double(x_lo))),
    #     dist.log_prob(tf.to_double(obs)) + tf.to_double(tf.math.log(1/255)))
    # we need
    p_exact = dist.cdf(tf.to_double(x_hi)) - dist.cdf(tf.to_double(x_lo))
    p_exact = tf.where(
        tf.math.abs(obs - loc) <= 2*scale, p_exact, tf.ones_like(p_exact))
    ret = tf.where(
        tf.math.abs(obs - loc) <= 2*scale, 
        tf.math.log(p_exact),
        dist.log_prob(tf.to_double(obs)) + tf.to_double(tf.math.log(1/255)))
    return tf.to_float(ret)
//  [python]
def discrete_logistics_likelihood(loc, scale, obs, side_censored=True):
    from tensorflow_probability import distributions as tfd
    import tensorflow as tf
    x_lo = obs - 1/255/2
    x_hi = obs + 1/255/2
    if side_censored:
        x_lo = tf.where(x_lo>=0, x_lo, tf.cast(-1000, dtype=tf.float32)*tf.ones_like(x_lo))
        x_hi = tf.where(x_hi<=1, x_hi, tf.cast(+1000, dtype=tf.float32)*tf.ones_like(x_hi))
    dist = tfd.Logistic(tf.to_double(loc), tf.to_double(scale))
    p_exact = dist.cdf(tf.to_double(x_hi)) - dist.cdf(tf.to_double(x_lo))
    p_exact = tf.where(
        tf.math.abs(obs - loc) <= 2*scale, p_exact, tf.ones_like(p_exact))
    ret = tf.where(
        tf.math.abs(obs - loc) <= 2*scale, 
        tf.math.log(p_exact),
        dist.log_prob(tf.to_double(obs)) + tf.to_double(tf.math.log(1/255)))
    return tf.cast(ret, dtype=tf.float32)
//  [python]
def compute_train_meta_loss(loss_mod, params, num=1):
  """Compute the train (inner) and validation (outer, meta) loss.

  Computes an average over num batches.
  Args:
    loss_mod: sonnet loss module
    params: loss mod params
    num: int
  Returns:
    train_loss, meta_loss
  """

  def compute_train_loss(_):
    return loss_mod.call_inner(params)

  train_loss = tf.reduce_mean(
      tf.map_fn(compute_train_loss, tf.to_float(tf.range(num))))

  def compute_meta_loss(_):
    return loss_mod.call_outer(params)

  meta_loss = tf.reduce_mean(
      tf.map_fn(compute_meta_loss, tf.to_float(tf.range(num))))
  return train_loss, meta_loss
//  [python]
def compute_train_meta_loss(loss_mod, params, num=1):
  """Compute the train (inner) and validation (outer, meta) loss.

  Computes an average over num batches.
  Args:
    loss_mod: sonnet loss module
    params: loss mod params
    num: int
  Returns:
    train_loss, meta_loss
  """

  def compute_train_loss(_):
    return loss_mod.call_inner(params)

  train_loss = tf.reduce_mean(
      tf.map_fn(compute_train_loss, tf.cast(tf.range(num), tf.float32)))

  def compute_meta_loss(_):
    return loss_mod.call_outer(params)

  meta_loss = tf.reduce_mean(
      tf.map_fn(compute_meta_loss, tf.cast(tf.range(num), tf.float32)))
  return train_loss, meta_loss
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.FixedLenFeature([], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.FixedLenFeature([], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def _set_train_or_infer(self, res, reverse_target_vocab_table, hparams):
    """Set up training and inference."""
    if self.mode == tf.contrib.learn.ModeKeys.TRAIN:
      self.train_loss = res[1]
      self.word_count = tf.reduce_sum(
          self.iterator.source_sequence_length) + tf.reduce_sum(
              self.iterator.target_sequence_length)
    elif self.mode == tf.contrib.learn.ModeKeys.EVAL:
      self.eval_loss = res[1]
    elif self.mode == tf.contrib.learn.ModeKeys.INFER:
      self.infer_logits, _, self.final_context_state, self.sample_id = res
      self.sample_words = reverse_target_vocab_table.lookup(
          tf.to_int64(self.sample_id))

    if self.mode != tf.contrib.learn.ModeKeys.INFER:
      ## Count the number of predicted words for compute ppl.
      self.predict_count = tf.reduce_sum(
          self.iterator.target_sequence_length)

    params = tf.trainable_variables()

    # Gradients and SGD update operation for training the model.
    # Arrange for the embedding vars to appear at the beginning.
    if self.mode == tf.contrib.learn.ModeKeys.TRAIN:
      self.learning_rate = tf.constant(hparams.learning_rate)
      # warm-up
      self.learning_rate = self._get_learning_rate_warmup(hparams)
      # decay
      self.learning_rate = self._get_learning_rate_decay(hparams)

      # Optimizer
      if hparams.optimizer == "sgd":
        opt = tf.train.GradientDescentOptimizer(self.learning_rate)
      elif hparams.optimizer == "adam":
        opt = tf.train.AdamOptimizer(self.learning_rate)
      else:
        raise ValueError("Unknown optimizer type %s" % hparams.optimizer)

      # Gradients
      gradients = tf.gradients(
          self.train_loss,
          params,
          colocate_gradients_with_ops=hparams.colocate_gradients_with_ops)

      clipped_grads, grad_norm_summary, grad_norm = model_helper.gradient_clip(
          gradients, max_gradient_norm=hparams.max_gradient_norm)
      self.grad_norm_summary = grad_norm_summary
      self.grad_norm = grad_norm

      self.update = opt.apply_gradients(
          zip(clipped_grads, params), global_step=self.global_step)

      # Summary
      self.train_summary = self._get_train_summary()
    elif self.mode == tf.contrib.learn.ModeKeys.INFER:
      self.infer_summary = self._get_infer_summary(hparams)

    # Print trainable variables
    utils.print_out("# Trainable variables")
    utils.print_out("Format: <name>, <shape>, <(soft) device placement>")
    for param in params:
      utils.print_out("  %s, %s, %s" % (param.name, str(param.get_shape()),
                                        param.op.device))
//  [python]
def _set_train_or_infer(self, res, reverse_target_vocab_table, hparams):
    """Set up training and inference."""
    if self.mode == tf.contrib.learn.ModeKeys.TRAIN:
      self.train_loss = res[1]
      self.word_count = tf.reduce_sum(
          self.iterator.source_sequence_length) + tf.reduce_sum(
              self.iterator.target_sequence_length)
    elif self.mode == tf.contrib.learn.ModeKeys.EVAL:
      self.eval_loss = res[1]
    elif self.mode == tf.contrib.learn.ModeKeys.INFER:
      self.infer_logits, _, self.final_context_state, self.sample_id = res
      self.sample_words = reverse_target_vocab_table.lookup(
          tf.cast(self.sample_id, tf.int64))

    if self.mode != tf.contrib.learn.ModeKeys.INFER:
      ## Count the number of predicted words for compute ppl.
      self.predict_count = tf.reduce_sum(
          self.iterator.target_sequence_length)

    params = tf.trainable_variables()

    # Gradients and SGD update operation for training the model.
    # Arrange for the embedding vars to appear at the beginning.
    if self.mode == tf.contrib.learn.ModeKeys.TRAIN:
      self.learning_rate = tf.constant(hparams.learning_rate)
      # warm-up
      self.learning_rate = self._get_learning_rate_warmup(hparams)
      # decay
      self.learning_rate = self._get_learning_rate_decay(hparams)

      # Optimizer
      if hparams.optimizer == "sgd":
        opt = tf.train.GradientDescentOptimizer(self.learning_rate)
      elif hparams.optimizer == "adam":
        opt = tf.train.AdamOptimizer(self.learning_rate)
      else:
        raise ValueError("Unknown optimizer type %s" % hparams.optimizer)

      # Gradients
      gradients = tf.gradients(
          self.train_loss,
          params,
          colocate_gradients_with_ops=hparams.colocate_gradients_with_ops)

      clipped_grads, grad_norm_summary, grad_norm = model_helper.gradient_clip(
          gradients, max_gradient_norm=hparams.max_gradient_norm)
      self.grad_norm_summary = grad_norm_summary
      self.grad_norm = grad_norm

      self.update = opt.apply_gradients(
          zip(clipped_grads, params), global_step=self.global_step)

      # Summary
      self.train_summary = self._get_train_summary()
    elif self.mode == tf.contrib.learn.ModeKeys.INFER:
      self.infer_summary = self._get_infer_summary(hparams)

    # Print trainable variables
    utils.print_out("# Trainable variables")
    utils.print_out("Format: <name>, <shape>, <(soft) device placement>")
    for param in params:
      utils.print_out("  %s, %s, %s" % (param.name, str(param.get_shape()),
                                        param.op.device))
//  [python]
def get_filled_rect(box, fill_vals, output_size, mode):
  """Returns a feature map with the values in the box filled in.

  This can either work with a single box or a batch of boxes.

  Args:
    box: The box to fill, in x,y,width,height format normalized
      between 0 and 1. shape: [num_boxes, 4]
    fill_vals: The vector to tile over each bounding box. This could be the
      embedding for the element.
    output_size: The length and width of the output. Assumes length and width
      are the same.
    mode: Method to compute values in the box. 'step': Pixels change immediately
      at box border. 'cont': Pixels change gradually at box border and increase
      towards the center.

  Returns:
    The tensor with the boxes filled.
  """

  with tf.variable_scope('get_filled_rect'):
    axis = tf.to_float(tf.range(output_size))

    disp_box = box
    disp_box *= tf.to_float(output_size)
    disp_box += [-1.5, -1.5, 3.0, 3.0]
    disp_box *= [1, 1, .5, .5]

    if mode == 'step':
      x_vals = tf.nn.relu(
          tf.nn.relu(axis - disp_box[:, 0:1]) / 2 -
          tf.nn.relu(axis - disp_box[:, 0:1] - disp_box[:, 2:3]))
      y_vals = tf.nn.relu(
          tf.nn.relu(axis - disp_box[:, 1:2]) / 2 -
          tf.nn.relu(axis - disp_box[:, 1:2] - disp_box[:, 3:4]))
    else:
      x_vals = tf.nn.relu(-tf.abs(axis - disp_box[:, 0:1] - disp_box[:, 2:3]) /
                          (disp_box[:, 2:3]) + 1)
      y_vals = tf.nn.relu(-tf.abs(axis - disp_box[:, 1:2] - disp_box[:, 3:4]) /
                          (disp_box[:, 3:4]) + 1)

    x_vals = tf.expand_dims(x_vals, 1)
    y_vals = tf.expand_dims(y_vals, 2)

    filled_rect = x_vals * y_vals
    if mode == 'step':
      filled_rect = tf.minimum(filled_rect, 1.0)

    fill_vals = tf.reshape(
        fill_vals, [tf.shape(fill_vals)[0], 1, 1,
                    tf.shape(fill_vals)[1]])
    filled_rect = tf.expand_dims(filled_rect, 3)
    filled_rect *= fill_vals

  return filled_rect
//  [python]
def get_filled_rect(box, fill_vals, output_size, mode):
  """Returns a feature map with the values in the box filled in.

  This can either work with a single box or a batch of boxes.

  Args:
    box: The box to fill, in x,y,width,height format normalized
      between 0 and 1. shape: [num_boxes, 4]
    fill_vals: The vector to tile over each bounding box. This could be the
      embedding for the element.
    output_size: The length and width of the output. Assumes length and width
      are the same.
    mode: Method to compute values in the box. 'step': Pixels change immediately
      at box border. 'cont': Pixels change gradually at box border and increase
      towards the center.

  Returns:
    The tensor with the boxes filled.
  """

  with tf.variable_scope('get_filled_rect'):
    axis = tf.cast(tf.range(output_size), tf.float32)

    disp_box = box
    disp_box *= tf.cast(output_size, tf.float32)
    disp_box += [-1.5, -1.5, 3.0, 3.0]
    disp_box *= [1, 1, .5, .5]

    if mode == 'step':
      x_vals = tf.nn.relu(
          tf.nn.relu(axis - disp_box[:, 0:1]) / 2 -
          tf.nn.relu(axis - disp_box[:, 0:1] - disp_box[:, 2:3]))
      y_vals = tf.nn.relu(
          tf.nn.relu(axis - disp_box[:, 1:2]) / 2 -
          tf.nn.relu(axis - disp_box[:, 1:2] - disp_box[:, 3:4]))
    else:
      x_vals = tf.nn.relu(-tf.abs(axis - disp_box[:, 0:1] - disp_box[:, 2:3]) /
                          (disp_box[:, 2:3]) + 1)
      y_vals = tf.nn.relu(-tf.abs(axis - disp_box[:, 1:2] - disp_box[:, 3:4]) /
                          (disp_box[:, 3:4]) + 1)

    x_vals = tf.expand_dims(x_vals, 1)
    y_vals = tf.expand_dims(y_vals, 2)

    filled_rect = x_vals * y_vals
    if mode == 'step':
      filled_rect = tf.minimum(filled_rect, 1.0)

    fill_vals = tf.reshape(
        fill_vals, [tf.shape(fill_vals)[0], 1, 1,
                    tf.shape(fill_vals)[1]])
    filled_rect = tf.expand_dims(filled_rect, 3)
    filled_rect *= fill_vals

  return filled_rect
//  [python]
def compute_policy_entropy(available_actions, policy, actions):
  """Compute total policy entropy.

  Args: (same as compute_policy_log_probs)

  Returns:
    entropy: a scalar float tensor.
  """

  def compute_entropy(probs):
    return -tf.reduce_sum(safe_log(probs) * probs, axis=-1)

  _, arg_ids = actions

  fn_pi, arg_pis = policy
  fn_pi = mask_unavailable_actions(available_actions, fn_pi)
  entropy = tf.reduce_mean(compute_entropy(fn_pi))
  tf.summary.scalar('entropy/fn', entropy)

  for arg_type in arg_ids.keys():
    arg_id = arg_ids[arg_type]
    arg_pi = arg_pis[arg_type]
    batch_mask = tf.to_float(tf.not_equal(arg_id, -1))
    arg_entropy = safe_div(
        tf.reduce_sum(compute_entropy(arg_pi) * batch_mask),
        tf.reduce_sum(batch_mask))
    entropy += arg_entropy
    tf.summary.scalar('used/arg/%s' % arg_type.name,
                      tf.reduce_mean(batch_mask))
    tf.summary.scalar('entropy/arg/%s' % arg_type.name, arg_entropy)

  return entropy
//  [python]
def compute_policy_entropy(available_actions, policy, actions):
  """Compute total policy entropy.

  Args: (same as compute_policy_log_probs)

  Returns:
    entropy: a scalar float tensor.
  """

  def compute_entropy(probs):
    return -tf.reduce_sum(safe_log(probs) * probs, axis=-1)

  _, arg_ids = actions

  fn_pi, arg_pis = policy
  fn_pi = mask_unavailable_actions(available_actions, fn_pi)
  entropy = tf.reduce_mean(compute_entropy(fn_pi))
  tf.summary.scalar('entropy/fn', entropy)

  for arg_type in arg_ids.keys():
    arg_id = arg_ids[arg_type]
    arg_pi = arg_pis[arg_type]
    batch_mask = tf.cast(tf.not_equal(arg_id, -1), tf.float32)
    arg_entropy = safe_div(
        tf.reduce_sum(compute_entropy(arg_pi) * batch_mask),
        tf.reduce_sum(batch_mask))
    entropy += arg_entropy
    tf.summary.scalar('used/arg/%s' % arg_type.name,
                      tf.reduce_mean(batch_mask))
    tf.summary.scalar('entropy/arg/%s' % arg_type.name, arg_entropy)

  return entropy
//  [python]
def to_float(x, name=None):
  x, b, lib = get_lib(x)
  if is_tf_backend(b):
    return tf.to_float(x, name=name)
  elif is_numpy_backend(b):
    return x.astype(np.float32)
  elif is_torch_backend(b):
    return x.float()
  else:
    raise NotImplementedError()
//  [python]
def to_float(x, name=None):
  x, b, lib = get_lib(x)
  if is_tf_backend(b):
    return tf.cast(x, tf.float32, name=name)
  elif is_numpy_backend(b):
    return x.astype(np.float32)
  elif is_torch_backend(b):
    return x.float()
  else:
    raise NotImplementedError()
//  [python]
def _using_transform_matrix(depth,
                            transform,
                            intrinsic_mat,
                            intrinsic_mat_inv=None):
  """A helper for using_transform_matrix. See docstring therein."""
  with tf.name_scope('Transform', values=[depth, transform, intrinsic_mat]):
    _, 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)
    cam_coords = tf.einsum('bij,jhw,bhw->bihw', intrinsic_mat_inv, grid, depth)

    rotation = transform[:, :3, :3]
    translation = transform[:, :3, 3]

    xyz = (
        tf.einsum('bij,bjk,bkhw->bihw', intrinsic_mat, rotation, cam_coords) +
        _expand_last_dim_twice(
            tf.einsum('bij,bj->bi', intrinsic_mat, translation)))

    x, y, z = tf.unstack(xyz, axis=1)
    pixel_x = x / z
    pixel_y = y / z
    return pixel_x, pixel_y, z
//  [python]
def _using_transform_matrix(depth,
                            transform,
                            intrinsic_mat,
                            intrinsic_mat_inv=None):
  """A helper for using_transform_matrix. See docstring therein."""
  with tf.name_scope('Transform', values=[depth, transform, intrinsic_mat]):
    _, 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)
    cam_coords = tf.einsum('bij,jhw,bhw->bihw', intrinsic_mat_inv, grid, depth)

    rotation = transform[:, :3, :3]
    translation = transform[:, :3, 3]

    xyz = (
        tf.einsum('bij,bjk,bkhw->bihw', intrinsic_mat, rotation, cam_coords) +
        _expand_last_dim_twice(
            tf.einsum('bij,bj->bi', intrinsic_mat, translation)))

    x, y, z = tf.unstack(xyz, axis=1)
    pixel_x = x / z
    pixel_y = y / z
    return pixel_x, pixel_y, z
//  [python]
def __init__(self, reuse=False, trainable=True):
    # 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")

    X = tf.to_float(self.states) / 255.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("value_net"):
      self.logits = tf.contrib.layers.fully_connected(
        inputs=fc1,
        num_outputs=1,
        activation_fn=None)
      self.logits = tf.squeeze(self.logits, squeeze_dims=[1], name="logits")

      self.losses = tf.squared_difference(self.logits, self.targets)
      self.loss = tf.reduce_sum(self.losses, name="loss")

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

      # Summaries
      prefix = tf.get_variable_scope().name
      tf.summary.scalar(self.loss.name, self.loss)
      tf.summary.scalar("{}/max_value".format(prefix), tf.reduce_max(self.logits))
      tf.summary.scalar("{}/min_value".format(prefix), tf.reduce_min(self.logits))
      tf.summary.scalar("{}/mean_value".format(prefix), tf.reduce_mean(self.logits))
      tf.summary.scalar("{}/reward_max".format(prefix), tf.reduce_max(self.targets))
      tf.summary.scalar("{}/reward_min".format(prefix), tf.reduce_min(self.targets))
      tf.summary.scalar("{}/reward_mean".format(prefix), tf.reduce_mean(self.targets))
      tf.summary.histogram("{}/reward_targets".format(prefix), self.targets)
      tf.summary.histogram("{}/values".format(prefix), self.logits)

      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())

    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, reuse=False, trainable=True):
    # 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")

    X = tf.cast(self.states, tf.float32) / 255.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("value_net"):
      self.logits = tf.contrib.layers.fully_connected(
        inputs=fc1,
        num_outputs=1,
        activation_fn=None)
      self.logits = tf.squeeze(self.logits, squeeze_dims=[1], name="logits")

      self.losses = tf.squared_difference(self.logits, self.targets)
      self.loss = tf.reduce_sum(self.losses, name="loss")

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

      # Summaries
      prefix = tf.get_variable_scope().name
      tf.summary.scalar(self.loss.name, self.loss)
      tf.summary.scalar("{}/max_value".format(prefix), tf.reduce_max(self.logits))
      tf.summary.scalar("{}/min_value".format(prefix), tf.reduce_min(self.logits))
      tf.summary.scalar("{}/mean_value".format(prefix), tf.reduce_mean(self.logits))
      tf.summary.scalar("{}/reward_max".format(prefix), tf.reduce_max(self.targets))
      tf.summary.scalar("{}/reward_min".format(prefix), tf.reduce_min(self.targets))
      tf.summary.scalar("{}/reward_mean".format(prefix), tf.reduce_mean(self.targets))
      tf.summary.histogram("{}/reward_targets".format(prefix), self.targets)
      tf.summary.histogram("{}/values".format(prefix), self.logits)

      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())

    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)