jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def build_valid_rl(self, shuffle=False):
    print("-" * 80)
    print("Build valid graph on shuffled data")
    with tf.device("/cpu:0"):
    #with tf.device("/device:GPU:0"):
      # shuffled valid data: for choosing validation model
      if not shuffle and self.data_format == "NCHW":
        self.images["valid_original"] = np.transpose(
          self.images["valid_original"], [0, 3, 1, 2])
      x_valid_shuffle, y_valid_shuffle = tf.train.shuffle_batch(
        [self.images["valid_original"], self.labels["valid_original"]],
        batch_size=self.batch_size,
        capacity=25000,
        enqueue_many=True,
        min_after_dequeue=0,
        num_threads=16,
        seed=self.seed,
        allow_smaller_final_batch=True,
      )

      def _pre_process(x):
        x = tf.pad(x, [[4, 4], [4, 4], [0, 0]])
        x = tf.random_crop(x, [32, 32, 3], seed=self.seed)
        x = tf.image.random_flip_left_right(x, seed=self.seed)
        if self.data_format == "NCHW":
          x = tf.transpose(x, [2, 0, 1])

        return x

      if shuffle:
        x_valid_shuffle = tf.map_fn(_pre_process, x_valid_shuffle,
                                    back_prop=False)

    logits = self._model(x_valid_shuffle, False, reuse=True)
    valid_shuffle_preds = tf.argmax(logits, axis=1)
    valid_shuffle_preds = tf.cast(valid_shuffle_preds, tf.int32)
    self.valid_shuffle_acc = tf.equal(valid_shuffle_preds, y_valid_shuffle)
    self.valid_shuffle_acc = tf.cast(self.valid_shuffle_acc, tf.int32)
    self.valid_shuffle_acc = tf.reduce_sum(self.valid_shuffle_acc)
//  [python]
def _dataset(self, images_file, labels_file):

    def DecodeImage(image):
      image = tf.decode_raw(image, tf.uint8)
      image = tf.cast(image, tf.float32)
      image = tf.reshape(image, [28, 28, 1])
      return image / 255.0

    def DecodeLabel(label):
      label = tf.decode_raw(label, tf.uint8)
      label = tf.reshape(label, [])
      return tf.to_int32(label)

    images = tf.data.FixedLengthRecordDataset(
        images_file, 28 * 28, header_bytes=16).map(
            DecodeImage, num_parallel_calls=16)
    labels = tf.data.FixedLengthRecordDataset(
        labels_file, 1, header_bytes=8).map(
            DecodeLabel, num_parallel_calls=16)
    return tf.data.Dataset.zip((images, labels))
//  [python]
def _dataset(self, images_file, labels_file):

    def DecodeImage(image):
      image = tf.decode_raw(image, tf.uint8)
      image = tf.cast(image, tf.float32)
      image = tf.reshape(image, [28, 28, 1])
      return image / 255.0

    def DecodeLabel(label):
      label = tf.decode_raw(label, tf.uint8)
      label = tf.reshape(label, [])
      return tf.cast(label, tf.int32)

    images = tf.data.FixedLengthRecordDataset(
        images_file, 28 * 28, header_bytes=16).map(
            DecodeImage, num_parallel_calls=16)
    labels = tf.data.FixedLengthRecordDataset(
        labels_file, 1, header_bytes=8).map(
            DecodeLabel, num_parallel_calls=16)
    return tf.data.Dataset.zip((images, labels))
//  [python]
def bilerp_gather(img, inds):
  """Bilinear interpolation dense gather from image at query inds."""

  inds_b, _, _, = tf.meshgrid(
      tf.range(tf.shape(img)[0]),
      tf.range(tf.shape(img)[1]),
      tf.range(tf.shape(img)[2]),
      indexing='ij')

  inds_b = tf.to_float(inds_b)
  inds_x = inds[Ellipsis, 0]
  inds_y = inds[Ellipsis, 1]

  inds_x_0 = tf.floor(inds_x)
  inds_x_1 = inds_x_0 + 1
  inds_y_0 = tf.floor(inds_y)
  inds_y_1 = inds_y_0 + 1

  # store invalid indices to implement correct out-of-bounds conditions
  invalid_x = tf.logical_or(
      tf.less(inds_x_0, 0.0),
      tf.greater(inds_x_1, tf.to_float(tf.shape(img)[2] - 1)))
  invalid_y = tf.logical_or(
      tf.less(inds_y_0, 0.0),
      tf.greater(inds_y_1, tf.to_float(tf.shape(img)[1] - 1)))
  invalid_inds = tf.logical_or(invalid_x, invalid_y)

  inds_x_0 = tf.clip_by_value(inds_x_0, 0.0, tf.to_float(tf.shape(img)[2] - 2))
  inds_x_1 = tf.clip_by_value(inds_x_1, 0.0, tf.to_float(tf.shape(img)[2] - 1))
  inds_y_0 = tf.clip_by_value(inds_y_0, 0.0, tf.to_float(tf.shape(img)[1] - 2))
  inds_y_1 = tf.clip_by_value(inds_y_1, 0.0, tf.to_float(tf.shape(img)[1] - 1))

  # compute interp weights
  w_x_0 = 1.0 - (inds_x - inds_x_0)
  w_x_1 = 1.0 - w_x_0
  w_y_0 = 1.0 - (inds_y - inds_y_0)
  w_y_1 = 1.0 - w_y_0

  w_0_0 = w_y_0 * w_x_0
  w_1_0 = w_y_1 * w_x_0
  w_0_1 = w_y_0 * w_x_1
  w_1_1 = w_y_1 * w_x_1

  # gather for interp
  inds_0_0 = tf.to_int32(tf.stack([inds_b, inds_y_0, inds_x_0], axis=-1))
  inds_1_0 = tf.to_int32(tf.stack([inds_b, inds_y_1, inds_x_0], axis=-1))
  inds_0_1 = tf.to_int32(tf.stack([inds_b, inds_y_0, inds_x_1], axis=-1))
  inds_1_1 = tf.to_int32(tf.stack([inds_b, inds_y_1, inds_x_1], axis=-1))

  img_0_0 = tf.gather_nd(img, inds_0_0) * w_0_0[Ellipsis, tf.newaxis]
  img_1_0 = tf.gather_nd(img, inds_1_0) * w_1_0[Ellipsis, tf.newaxis]
  img_0_1 = tf.gather_nd(img, inds_0_1) * w_0_1[Ellipsis, tf.newaxis]
  img_1_1 = tf.gather_nd(img, inds_1_1) * w_1_1[Ellipsis, tf.newaxis]

  out_img = img_0_0 + img_1_0 + img_0_1 + img_1_1

  # boundary conditions for invalid indices
  invalid_inds = tf.tile(invalid_inds[:, :, :, tf.newaxis],
                         [1, 1, 1, tf.shape(img)[3]])

  out_img = tf.where(invalid_inds, tf.zeros_like(out_img), out_img)

  return out_img
//  [python]
def bilerp_gather(img, inds):
  """Bilinear interpolation dense gather from image at query inds."""

  inds_b, _, _, = tf.meshgrid(
      tf.range(tf.shape(img)[0]),
      tf.range(tf.shape(img)[1]),
      tf.range(tf.shape(img)[2]),
      indexing='ij')

  inds_b = tf.cast(inds_b, tf.float32)
  inds_x = inds[Ellipsis, 0]
  inds_y = inds[Ellipsis, 1]

  inds_x_0 = tf.floor(inds_x)
  inds_x_1 = inds_x_0 + 1
  inds_y_0 = tf.floor(inds_y)
  inds_y_1 = inds_y_0 + 1

  # store invalid indices to implement correct out-of-bounds conditions
  invalid_x = tf.logical_or(
      tf.less(inds_x_0, 0.0),
      tf.greater(inds_x_1, tf.cast(tf.shape(img)[2] - 1, tf.float32)))
  invalid_y = tf.logical_or(
      tf.less(inds_y_0, 0.0),
      tf.greater(inds_y_1, tf.cast(tf.shape(img)[1] - 1, tf.float32)))
  invalid_inds = tf.logical_or(invalid_x, invalid_y)

  inds_x_0 = tf.clip_by_value(inds_x_0, 0.0, tf.cast(tf.shape(img)[2] - 2, tf.float32))
  inds_x_1 = tf.clip_by_value(inds_x_1, 0.0, tf.cast(tf.shape(img)[2] - 1, tf.float32))
  inds_y_0 = tf.clip_by_value(inds_y_0, 0.0, tf.cast(tf.shape(img)[1] - 2, tf.float32))
  inds_y_1 = tf.clip_by_value(inds_y_1, 0.0, tf.cast(tf.shape(img)[1] - 1, tf.float32))

  # compute interp weights
  w_x_0 = 1.0 - (inds_x - inds_x_0)
  w_x_1 = 1.0 - w_x_0
  w_y_0 = 1.0 - (inds_y - inds_y_0)
  w_y_1 = 1.0 - w_y_0

  w_0_0 = w_y_0 * w_x_0
  w_1_0 = w_y_1 * w_x_0
  w_0_1 = w_y_0 * w_x_1
  w_1_1 = w_y_1 * w_x_1

  # gather for interp
  inds_0_0 = tf.to_int32(tf.stack([inds_b, inds_y_0, inds_x_0], axis=-1))
  inds_1_0 = tf.to_int32(tf.stack([inds_b, inds_y_1, inds_x_0], axis=-1))
  inds_0_1 = tf.to_int32(tf.stack([inds_b, inds_y_0, inds_x_1], axis=-1))
  inds_1_1 = tf.to_int32(tf.stack([inds_b, inds_y_1, inds_x_1], axis=-1))

  img_0_0 = tf.gather_nd(img, inds_0_0) * w_0_0[Ellipsis, tf.newaxis]
  img_1_0 = tf.gather_nd(img, inds_1_0) * w_1_0[Ellipsis, tf.newaxis]
  img_0_1 = tf.gather_nd(img, inds_0_1) * w_0_1[Ellipsis, tf.newaxis]
  img_1_1 = tf.gather_nd(img, inds_1_1) * w_1_1[Ellipsis, tf.newaxis]

  out_img = img_0_0 + img_1_0 + img_0_1 + img_1_1

  # boundary conditions for invalid indices
  invalid_inds = tf.tile(invalid_inds[:, :, :, tf.newaxis],
                         [1, 1, 1, tf.shape(img)[3]])

  out_img = tf.where(invalid_inds, tf.zeros_like(out_img), out_img)

  return out_img
//  [python]
def _dataset(self, data_files,
               parallel_interleave = False):
    if parallel_interleave:
      data_files = tf.data.Dataset.from_tensor_slices(data_files)
      data = data_files.interleave(
          lambda f: tf.data.FixedLengthRecordDataset(f, 1 + 32 * 32 * 3),
          cycle_length=6)
    else:
      data = tf.data.FixedLengthRecordDataset(data_files, 1 + 32 * 32 * 3)

    def DecodeLabelAndImage(r):
      r = tf.decode_raw(r, tf.uint8)
      return tf.to_float(
          tf.transpose(tf.reshape(r[1:], [3, 32, 32]),
                       [1, 2, 0])) / 255.0, tf.to_int32(r[0])

    return data.map(DecodeLabelAndImage)
//  [python]
def _dataset(self, data_files,
               parallel_interleave = False):
    if parallel_interleave:
      data_files = tf.data.Dataset.from_tensor_slices(data_files)
      data = data_files.interleave(
          lambda f: tf.data.FixedLengthRecordDataset(f, 1 + 32 * 32 * 3),
          cycle_length=6)
    else:
      data = tf.data.FixedLengthRecordDataset(data_files, 1 + 32 * 32 * 3)

    def DecodeLabelAndImage(r):
      r = tf.decode_raw(r, tf.uint8)
      return tf.to_float(
          tf.transpose(tf.reshape(r[1:], [3, 32, 32]),
                       [1, 2, 0])) / 255.0, tf.cast(r[0], tf.int32)

    return data.map(DecodeLabelAndImage)
//  [python]
def boolean_mask(boxlist, indicator, fields=None, scope=None,
                 use_static_shapes=False, indicator_sum=None):
  """Select boxes from BoxList according to indicator and return new BoxList.

  `boolean_mask` returns the subset of boxes that are marked as "True" by the
  indicator tensor. By default, `boolean_mask` returns boxes corresponding to
  the input index list, as well as all additional fields stored in the boxlist
  (indexing into the first dimension).  However one can optionally only draw
  from a subset of fields.

  Args:
    boxlist: BoxList holding N boxes
    indicator: a rank-1 boolean tensor
    fields: (optional) list of fields to also gather from.  If None (default),
      all fields are gathered from.  Pass an empty fields list to only gather
      the box coordinates.
    scope: name scope.
    use_static_shapes: Whether to use an implementation with static shape
      gurantees.
    indicator_sum: An integer containing the sum of `indicator` vector. Only
      required if `use_static_shape` is True.

  Returns:
    subboxlist: a BoxList corresponding to the subset of the input BoxList
      specified by indicator
  Raises:
    ValueError: if `indicator` is not a rank-1 boolean tensor.
  """
  with tf.name_scope(scope, 'BooleanMask'):
    if indicator.shape.ndims != 1:
      raise ValueError('indicator should have rank 1')
    if indicator.dtype != tf.bool:
      raise ValueError('indicator should be a boolean tensor')
    if use_static_shapes:
      if not (indicator_sum and isinstance(indicator_sum, int)):
        raise ValueError('`indicator_sum` must be a of type int')
      selected_positions = tf.to_float(indicator)
      indexed_positions = tf.cast(
          tf.multiply(
              tf.cumsum(selected_positions), selected_positions),
          dtype=tf.int32)
      one_hot_selector = tf.one_hot(
          indexed_positions - 1, indicator_sum, dtype=tf.float32)
      sampled_indices = tf.cast(
          tf.tensordot(
              tf.to_float(tf.range(tf.shape(indicator)[0])),
              one_hot_selector,
              axes=[0, 0]),
          dtype=tf.int32)
      return gather(boxlist, sampled_indices, use_static_shapes=True)
    else:
      subboxlist = box_list.BoxList(tf.boolean_mask(boxlist.get(), indicator))
      if fields is None:
        fields = boxlist.get_extra_fields()
      for field in fields:
        if not boxlist.has_field(field):
          raise ValueError('boxlist must contain all specified fields')
        subfieldlist = tf.boolean_mask(boxlist.get_field(field), indicator)
        subboxlist.add_field(field, subfieldlist)
      return subboxlist
//  [python]
def boolean_mask(boxlist, indicator, fields=None, scope=None,
                 use_static_shapes=False, indicator_sum=None):
  """Select boxes from BoxList according to indicator and return new BoxList.

  `boolean_mask` returns the subset of boxes that are marked as "True" by the
  indicator tensor. By default, `boolean_mask` returns boxes corresponding to
  the input index list, as well as all additional fields stored in the boxlist
  (indexing into the first dimension).  However one can optionally only draw
  from a subset of fields.

  Args:
    boxlist: BoxList holding N boxes
    indicator: a rank-1 boolean tensor
    fields: (optional) list of fields to also gather from.  If None (default),
      all fields are gathered from.  Pass an empty fields list to only gather
      the box coordinates.
    scope: name scope.
    use_static_shapes: Whether to use an implementation with static shape
      gurantees.
    indicator_sum: An integer containing the sum of `indicator` vector. Only
      required if `use_static_shape` is True.

  Returns:
    subboxlist: a BoxList corresponding to the subset of the input BoxList
      specified by indicator
  Raises:
    ValueError: if `indicator` is not a rank-1 boolean tensor.
  """
  with tf.name_scope(scope, 'BooleanMask'):
    if indicator.shape.ndims != 1:
      raise ValueError('indicator should have rank 1')
    if indicator.dtype != tf.bool:
      raise ValueError('indicator should be a boolean tensor')
    if use_static_shapes:
      if not (indicator_sum and isinstance(indicator_sum, int)):
        raise ValueError('`indicator_sum` must be a of type int')
      selected_positions = tf.cast(indicator, dtype=tf.float32)
      indexed_positions = tf.cast(
          tf.multiply(
              tf.cumsum(selected_positions), selected_positions),
          dtype=tf.int32)
      one_hot_selector = tf.one_hot(
          indexed_positions - 1, indicator_sum, dtype=tf.float32)
      sampled_indices = tf.cast(
          tf.tensordot(
              tf.cast(tf.range(tf.shape(indicator)[0]), dtype=tf.float32),
              one_hot_selector,
              axes=[0, 0]),
          dtype=tf.int32)
      return gather(boxlist, sampled_indices, use_static_shapes=True)
    else:
      subboxlist = box_list.BoxList(tf.boolean_mask(boxlist.get(), indicator))
      if fields is None:
        fields = boxlist.get_extra_fields()
      for field in fields:
        if not boxlist.has_field(field):
          raise ValueError('boxlist must contain all specified fields')
        subfieldlist = tf.boolean_mask(boxlist.get_field(field), indicator)
        subboxlist.add_field(field, subfieldlist)
      return subboxlist
//  [python]
def AISIterator(self, batch_size, shard_idx, num_workers):
    data = self._dataset(num=self.test_size)

    # Can this be done more efficiently? Need some sort of dynamic seed PRNG.
    rs = np.random.RandomState(0)
    masks = tf.to_float(rs.rand(self.test_size, 28, 28, 1))
    mask_data = tf.data.Dataset.from_tensor_slices(masks)
    data = tf.data.Dataset.zip((mask_data, data))
    data = data.shard(num_workers, shard_idx)

    def StaticBinarize(mask, image_and_label):
      image, label = image_and_label
      return tf.to_float(tf.to_float(image) > mask), label

    data = data.map(StaticBinarize)
    data = data.batch(batch_size)
    data = data.prefetch(1)
    return data.make_initializable_iterator()
//  [python]
def AISIterator(self, batch_size, shard_idx, num_workers):
    data = self._dataset(num=self.test_size)

    # Can this be done more efficiently? Need some sort of dynamic seed PRNG.
    rs = np.random.RandomState(0)
    masks = tf.cast(rs.rand(self.test_size, 28, 28, 1), dtype=tf.float32)
    mask_data = tf.data.Dataset.from_tensor_slices(masks)
    data = tf.data.Dataset.zip((mask_data, data))
    data = data.shard(num_workers, shard_idx)

    def StaticBinarize(mask, image_and_label):
      image, label = image_and_label
      return tf.cast(tf.cast(image, dtype=tf.float32) > mask, dtype=tf.float32), label

    data = data.map(StaticBinarize)
    data = data.batch(batch_size)
    data = data.prefetch(1)
    return data.make_initializable_iterator()
//  [python]
def Covariance(x):
  x = tf.to_float(x)
  mean = tf.reduce_mean(x, 0)
  x -= mean
  return tf.matmul(x, x, transpose_a=True) / tf.to_float(tf.shape(x)[0] - 1)
//  [python]
def Covariance(x):
  x = tf.cast(x, tf.float32)
  mean = tf.reduce_mean(x, 0)
  x -= mean
  return tf.matmul(x, x, transpose_a=True) / tf.cast(tf.shape(x)[0] - 1, tf.float32)
//  [python]
def AISIterator(self, batch_size, shard_idx, num_workers):
    if self._test_is_valid:
      data = self._dataset(
          os.path.join(self._data_dir, "train-images-idx3-ubyte"),
          os.path.join(self._data_dir, "train-labels-idx1-ubyte"))
      data = data.take(self.test_size)
    else:
      data = self._dataset(
          os.path.join(self._data_dir, "t10k-images-idx3-ubyte"),
          os.path.join(self._data_dir, "t10k-labels-idx1-ubyte"))

    # Can this be done more efficiently? Need some sort of dynamic seed PRNG.
    rs = np.random.RandomState(0)
    masks = tf.to_float(rs.rand(self.test_size, 28, 28, 1))
    mask_data = tf.data.Dataset.from_tensor_slices(masks)
    data = tf.data.Dataset.zip((mask_data, data))
    data = data.shard(num_workers, shard_idx)

    def StaticBinarize(mask, image_and_label):
      image, label = image_and_label
      return tf.to_float(tf.to_float(image) > mask), label

    data = data.map(StaticBinarize)
    data = data.batch(batch_size)
    data = data.prefetch(1)
    return data.make_initializable_iterator()
//  [python]
def AISIterator(self, batch_size, shard_idx, num_workers):
    if self._test_is_valid:
      data = self._dataset(
          os.path.join(self._data_dir, "train-images-idx3-ubyte"),
          os.path.join(self._data_dir, "train-labels-idx1-ubyte"))
      data = data.take(self.test_size)
    else:
      data = self._dataset(
          os.path.join(self._data_dir, "t10k-images-idx3-ubyte"),
          os.path.join(self._data_dir, "t10k-labels-idx1-ubyte"))

    # Can this be done more efficiently? Need some sort of dynamic seed PRNG.
    rs = np.random.RandomState(0)
    masks = tf.cast(rs.rand(self.test_size, 28, 28, 1), dtype=tf.float32)
    mask_data = tf.data.Dataset.from_tensor_slices(masks)
    data = tf.data.Dataset.zip((mask_data, data))
    data = data.shard(num_workers, shard_idx)

    def StaticBinarize(mask, image_and_label):
      image, label = image_and_label
      return tf.cast(tf.cast(image, dtype=tf.float32) > mask, dtype=tf.float32), label

    data = data.map(StaticBinarize)
    data = data.batch(batch_size)
    data = data.prefetch(1)
    return data.make_initializable_iterator()
//  [python]
def Expectation(dist, num_steps, batch_size, f=tf.identity):

  def Mean():
    return tf.reduce_mean(f(dist.sample(batch_size)), 0)

  def Cond(step, _):
    return step < num_steps

  def Body(step, running_mean):
    new_mean = Mean()

    running_mean = (
        running_mean * tf.to_float(step) + new_mean) / tf.to_float(step + 1)

    return step + 1, running_mean

  return tf.while_loop(Cond, Body, [0, Mean()])[1]
//  [python]
def Expectation(dist, num_steps, batch_size, f=tf.identity):

  def Mean():
    return tf.reduce_mean(f(dist.sample(batch_size)), 0)

  def Cond(step, _):
    return step < num_steps

  def Body(step, running_mean):
    new_mean = Mean()

    running_mean = (
        running_mean * tf.cast(step, tf.float32) + new_mean) / tf.cast(step + 1, tf.float32)

    return step + 1, running_mean

  return tf.while_loop(Cond, Body, [0, Mean()])[1]
//  [python]
def StitchImages(images):
  # images is [batch, x, y, c]
  batch, width, _, channels = tf.unstack(tf.shape(images))
  num_per_side = tf.to_int32(tf.ceil(tf.sqrt(tf.to_float(batch))))
  new_width = num_per_side * width
  paddings = tf.concat([tf.zeros([4, 1], dtype=tf.int32), tf.stack([num_per_side * num_per_side - batch, 0, 0, 0])[Ellipsis, tf.newaxis]], -1)
  images = tf.pad(images, paddings)

  images = tf.transpose(images, [1, 0, 2, 3])
  images = tf.reshape(images, [width, num_per_side, new_width, channels])
  images = tf.transpose(images, [1, 0, 2, 3])
  images = tf.reshape(images, [1, new_width, new_width, channels])

  return images
//  [python]
def StitchImages(images):
  # images is [batch, x, y, c]
  batch, width, _, channels = tf.unstack(tf.shape(images))
  num_per_side = tf.cast(tf.ceil(tf.cast(batch, tf.float32)), tf.int32)
  new_width = num_per_side * width
  paddings = tf.concat([tf.zeros([4, 1], dtype=tf.int32), tf.stack([num_per_side * num_per_side - batch, 0, 0, 0])[Ellipsis, tf.newaxis]], -1)
  images = tf.pad(images, paddings)

  images = tf.transpose(images, [1, 0, 2, 3])
  images = tf.reshape(images, [width, num_per_side, new_width, channels])
  images = tf.transpose(images, [1, 0, 2, 3])
  images = tf.reshape(images, [1, new_width, new_width, channels])

  return images
//  [python]
def log_prob(self, actions):
    return tf.to_float(self._parameter * actions)
//  [python]
def log_prob(self, actions):
    return tf.cast(self._parameter * actions, tf.float32)
//  [python]
def _prng_key(i, key):
  return tf.stack([tf.to_int32(key), tf.to_int32(i)], axis=0)
//  [python]
def _prng_key(i, key):
  return tf.stack([tf.cast(key, tf.int32), tf.cast(i, tf.int32)], axis=0)
//  [python]
def _distributional_to_value(value_d, size, subscale, threshold):
  """Get a scalar value out of a value distribution in distributional RL."""
  half = size // 2
  value_range = (tf.to_float(tf.range(-half, half)) + 0.5) * subscale
  probs = tf.nn.softmax(value_d)

  if threshold == 0.0:
    return tf.reduce_sum(probs * value_range, axis=-1)

  # accumulated_probs[..., i] is the sum of probabilities in buckets upto i
  # so it is the probability that value <= i'th bucket value
  accumulated_probs = tf.cumsum(probs, axis=-1)
  # New probs are 0 on all lower buckets, until the threshold
  probs = tf.where(accumulated_probs < threshold, tf.zeros_like(probs), probs)
  probs /= tf.reduce_sum(probs, axis=-1, keepdims=True)  # Re-normalize.
  return tf.reduce_sum(probs * value_range, axis=-1)
//  [python]
def _distributional_to_value(value_d, size, subscale, threshold):
  """Get a scalar value out of a value distribution in distributional RL."""
  half = size // 2
  value_range = (tf.cast(tf.range(-half, half), tf.float32) + 0.5) * subscale
  probs = tf.nn.softmax(value_d)

  if threshold == 0.0:
    return tf.reduce_sum(probs * value_range, axis=-1)

  # accumulated_probs[..., i] is the sum of probabilities in buckets upto i
  # so it is the probability that value <= i'th bucket value
  accumulated_probs = tf.cumsum(probs, axis=-1)
  # New probs are 0 on all lower buckets, until the threshold
  probs = tf.where(accumulated_probs < threshold, tf.zeros_like(probs), probs)
  probs /= tf.reduce_sum(probs, axis=-1, keepdims=True)  # Re-normalize.
  return tf.reduce_sum(probs * value_range, axis=-1)
//  [python]
def to_byte(d: dict):
    return tf.to_int32(tf.round(127.5 * (d['image'] + 1)))
//  [python]
def to_byte(d: dict):
    return tf.cast(tf.round(127.5 * (d['image'] + 1)), tf.int32)
//  [python]
def resize_img_labels(image, label, resized_h, resized_w):

    new_shape = tf.stack([tf.to_int32(resized_h), tf.to_int32(resized_w)])
    img = tf.image.resize_images(image, new_shape)
    label = tf.image.resize_nearest_neighbor(tf.expand_dims(label, 0), new_shape)
    label = tf.squeeze(label, squeeze_dims=[0])
    return img, label
//  [python]
def resize_img_labels(image, label, resized_h, resized_w):

    new_shape = tf.stack([tf.cast(resized_h, tf.int32), tf.cast(resized_w, tf.int32)])
    img = tf.image.resize_images(image, new_shape)
    label = tf.image.resize_nearest_neighbor(tf.expand_dims(label, 0), new_shape)
    label = tf.squeeze(label, squeeze_dims=[0])
    return img, label
//  [python]
def features_to_nonpadding(features, inputs_or_targets="inputs"):
  key = inputs_or_targets + "_segmentation"
  if features and key in features:
    return tf.minimum(tf.to_float(features[key]), 1.0)
  return None
//  [python]
def features_to_nonpadding(features, inputs_or_targets="inputs"):
  key = inputs_or_targets + "_segmentation"
  if features and key in features:
    return tf.minimum(tf.cast(features[key], tf.float32), 1.0)
  return None
//  [python]
def cast_ints_to_int32(features):
  f = {}
  for k, v in sorted(six.iteritems(features)):
    if v.dtype in [tf.int64, tf.uint8]:
      v = tf.to_int32(v)
    f[k] = v
  return f
//  [python]
def cast_ints_to_int32(features):
  f = {}
  for k, v in sorted(six.iteritems(features)):
    if v.dtype in [tf.int64, tf.uint8]:
      v = tf.cast(v, tf.int32)
    f[k] = v
  return f
//  [python]
def random_resize_img_labels(image, label, resized_h, resized_w):

    scale = tf.random_uniform([1], minval=0.75, maxval=1.25, dtype=tf.float32, seed=None)
    h_new = tf.to_int32(tf.multiply(tf.to_float(resized_h), scale))
    w_new = tf.to_int32(tf.multiply(tf.to_float(resized_w), scale))

    new_shape = tf.squeeze(tf.stack([h_new, w_new]), squeeze_dims=[1])
    img = tf.image.resize_images(image, new_shape)
    label = tf.image.resize_nearest_neighbor(tf.expand_dims(label, 0), new_shape)
    label = tf.squeeze(label, squeeze_dims=[0])
    return img, label
//  [python]
def random_resize_img_labels(image, label, resized_h, resized_w):

    scale = tf.random_uniform([1], minval=0.75, maxval=1.25, dtype=tf.float32, seed=None)
    h_new = tf.cast(tf.multiply(tf.to_float(resized_h), scale), tf.int32)
    w_new = tf.cast(tf.multiply(tf.to_float(resized_w), scale), tf.int32)

    new_shape = tf.squeeze(tf.stack([h_new, w_new]), squeeze_dims=[1])
    img = tf.image.resize_images(image, new_shape)
    label = tf.image.resize_nearest_neighbor(tf.expand_dims(label, 0), new_shape)
    label = tf.squeeze(label, squeeze_dims=[0])
    return img, label
//  [python]
def image_scaling(img, label, edge):
    """
    Randomly scales the images between 0.5 to 1.5 times the original size.
    Args:
      img: Training image to scale.
      label: Segmentation mask to scale.
    """
    
    scale = tf.random_uniform([1], minval=0.5, maxval=2.0, dtype=tf.float32, seed=None)
    h_new = tf.to_int32(tf.multiply(tf.to_float(tf.shape(img)[0]), scale))
    w_new = tf.to_int32(tf.multiply(tf.to_float(tf.shape(img)[1]), scale))
    new_shape = tf.squeeze(tf.stack([h_new, w_new]), squeeze_dims=[1])
    img = tf.image.resize_images(img, new_shape)
    #label = tf.image.resize_nearest_neighbor(tf.expand_dims(label, 0), new_shape)
    #label = tf.squeeze(label, squeeze_dims=[0])
    #edge = tf.image.resize_nearest_neighbor(tf.expand_dims(edge, 0), new_shape)
    #edge = tf.squeeze(edge, squeeze_dims=[0])
   
    #return img, label, edge
    return img
//  [python]
def image_scaling(img, label, edge):
    """
    Randomly scales the images between 0.5 to 1.5 times the original size.
    Args:
      img: Training image to scale.
      label: Segmentation mask to scale.
    """
    
    scale = tf.random_uniform([1], minval=0.5, maxval=2.0, dtype=tf.float32, seed=None)
    h_new = tf.cast(tf.multiply(tf.cast(tf.shape(img)[0], tf.float32), scale), tf.int32)
    w_new = tf.cast(tf.multiply(tf.cast(tf.shape(img)[1], tf.float32), scale), tf.int32)
    new_shape = tf.squeeze(tf.stack([h_new, w_new]), squeeze_dims=[1])
    img = tf.image.resize_images(img, new_shape)
    #label = tf.image.resize_nearest_neighbor(tf.expand_dims(label, 0), new_shape)
    #label = tf.squeeze(label, squeeze_dims=[0])
    #edge = tf.image.resize_nearest_neighbor(tf.expand_dims(edge, 0), new_shape)
    #edge = tf.squeeze(edge, squeeze_dims=[0])
   
    #return img, label, edge
    return img
//  [python]
def _summarize_features(features, num_shards=1):
  with tf.name_scope("input_stats"):
    for (k, v) in six.iteritems(features):
      if isinstance(v, tf.Tensor) and v.get_shape().ndims > 1:
        tf.summary.scalar("%s_batch" % k, tf.shape(v)[0] // num_shards)
        tf.summary.scalar("%s_length" % k, tf.shape(v)[1])
        nonpadding = tf.to_float(tf.not_equal(v, 0))
        nonpadding_tokens = tf.reduce_sum(nonpadding)
        tf.summary.scalar("%s_nonpadding_tokens" % k, nonpadding_tokens)
        tf.summary.scalar("%s_nonpadding_fraction" % k,
                          tf.reduce_mean(nonpadding))
//  [python]
def _summarize_features(features, num_shards=1):
  with tf.name_scope("input_stats"):
    for (k, v) in six.iteritems(features):
      if isinstance(v, tf.Tensor) and v.get_shape().ndims > 1:
        tf.summary.scalar("%s_batch" % k, tf.shape(v)[0] // num_shards)
        tf.summary.scalar("%s_length" % k, tf.shape(v)[1])
        nonpadding = tf.cast(tf.not_equal(v, 0), tf.float32)
        nonpadding_tokens = tf.reduce_sum(nonpadding)
        tf.summary.scalar("%s_nonpadding_tokens" % k, nonpadding_tokens)
        tf.summary.scalar("%s_nonpadding_fraction" % k,
                          tf.reduce_mean(nonpadding))
//  [python]
def masked_apply(tensor, op, mask, set_outside_zero=True):
    """
    Appliy the function op to tensor only at locations indicated by mask. If set_outside_zero == True, set the
    locations outside the mask to zero, otherwise keep original value of tensor.
    :param tensor: The tensor on which op is applied.
    :param op: The operation.
    :param mask: The boolean mask.
    :param set_outside_zero: If True, set the locations outside the mask to zero, otherwise keep original values of tensor.
    :return: Tensor with applied function.
    """
    chosen = tf.boolean_mask(tensor, mask)
    applied = op(chosen)
    idx = tf.to_int32(tf.where(mask))
    result = tf.scatter_nd(idx, applied, tf.shape(tensor))
    if not set_outside_zero:
        result = tf.where(mask, result, tensor)
    return result
//  [python]
def masked_apply(tensor, op, mask, set_outside_zero=True):
    """
    Appliy the function op to tensor only at locations indicated by mask. If set_outside_zero == True, set the
    locations outside the mask to zero, otherwise keep original value of tensor.
    :param tensor: The tensor on which op is applied.
    :param op: The operation.
    :param mask: The boolean mask.
    :param set_outside_zero: If True, set the locations outside the mask to zero, otherwise keep original values of tensor.
    :return: Tensor with applied function.
    """
    chosen = tf.boolean_mask(tensor, mask)
    applied = op(chosen)
    idx = tf.cast(tf.where(mask), tf.int32)
    result = tf.scatter_nd(idx, applied, tf.shape(tensor))
    if not set_outside_zero:
        result = tf.where(mask, result, tensor)
    return result
//  [python]
def pixels_from_softmax(frame_logits, pure_sampling=False,
                        temperature=1.0, gumbel_noise_factor=0.2):
  """Given frame_logits from a per-pixel softmax, generate colors."""
  # If we're purely sampling, just sample each pixel.
  if pure_sampling or temperature == 0.0:
    return common_layers.sample_with_temperature(frame_logits, temperature)

  # Gumbel-sample from the pixel sofmax and average by pixel values.
  pixel_range = tf.to_float(tf.range(256))
  for _ in range(len(frame_logits.get_shape().as_list()) - 1):
    pixel_range = tf.expand_dims(pixel_range, axis=0)

  frame_logits = tf.nn.log_softmax(frame_logits)
  gumbel_samples = discretization.gumbel_sample(
      common_layers.shape_list(frame_logits)) * gumbel_noise_factor

  frame = tf.nn.softmax((frame_logits + gumbel_samples) / temperature, axis=-1)
  result = tf.reduce_sum(frame * pixel_range, axis=-1)
  # Round on the forward pass, not on the backward one.
  return result + tf.stop_gradient(tf.round(result) - result)
//  [python]
def pixels_from_softmax(frame_logits, pure_sampling=False,
                        temperature=1.0, gumbel_noise_factor=0.2):
  """Given frame_logits from a per-pixel softmax, generate colors."""
  # If we're purely sampling, just sample each pixel.
  if pure_sampling or temperature == 0.0:
    return common_layers.sample_with_temperature(frame_logits, temperature)

  # Gumbel-sample from the pixel sofmax and average by pixel values.
  pixel_range = tf.cast(tf.range(256), tf.float32)
  for _ in range(len(frame_logits.get_shape().as_list()) - 1):
    pixel_range = tf.expand_dims(pixel_range, axis=0)

  frame_logits = tf.nn.log_softmax(frame_logits)
  gumbel_samples = discretization.gumbel_sample(
      common_layers.shape_list(frame_logits)) * gumbel_noise_factor

  frame = tf.nn.softmax((frame_logits + gumbel_samples) / temperature, axis=-1)
  result = tf.reduce_sum(frame * pixel_range, axis=-1)
  # Round on the forward pass, not on the backward one.
  return result + tf.stop_gradient(tf.round(result) - result)
//  [python]
def get_sampled_frame(self, pred_frame):
    """Samples the frame based on modality.

      if the modality is L2/L1 then the next predicted frame is the
      next frame and there is no sampling but in case of Softmax loss
      the next actual frame should be sampled from predicted frame.

      This enables multi-frame target prediction with Softmax loss.

    Args:
      pred_frame: predicted frame.

    Returns:
      sampled frame.

    """
    # TODO(lukaszkaiser): the logic below heavily depend on the current
    # (a bit strange) video modalities - we should change that.

    sampled_frame = pred_frame
    if self.is_per_pixel_softmax:
      frame_shape = common_layers.shape_list(pred_frame)
      target_shape = frame_shape[:-1] + [self.hparams.problem.num_channels]
      sampled_frame = tf.reshape(pred_frame, target_shape + [256])
      sampled_frame = pixels_from_softmax(
          sampled_frame, temperature=self.hparams.pixel_sampling_temperature)
      # TODO(lukaszkaiser): this should be consistent with modality.bottom()
      # sampled_frame = common_layers.standardize_images(sampled_frame)
    return tf.to_float(sampled_frame)
//  [python]
def get_sampled_frame(self, pred_frame):
    """Samples the frame based on modality.

      if the modality is L2/L1 then the next predicted frame is the
      next frame and there is no sampling but in case of Softmax loss
      the next actual frame should be sampled from predicted frame.

      This enables multi-frame target prediction with Softmax loss.

    Args:
      pred_frame: predicted frame.

    Returns:
      sampled frame.

    """
    # TODO(lukaszkaiser): the logic below heavily depend on the current
    # (a bit strange) video modalities - we should change that.

    sampled_frame = pred_frame
    if self.is_per_pixel_softmax:
      frame_shape = common_layers.shape_list(pred_frame)
      target_shape = frame_shape[:-1] + [self.hparams.problem.num_channels]
      sampled_frame = tf.reshape(pred_frame, target_shape + [256])
      sampled_frame = pixels_from_softmax(
          sampled_frame, temperature=self.hparams.pixel_sampling_temperature)
      # TODO(lukaszkaiser): this should be consistent with modality.bottom()
      # sampled_frame = common_layers.standardize_images(sampled_frame)
    return tf.cast(sampled_frame, tf.float32)
//  [python]
def _pairwise_distances(embeddings, squared=False):
    """Compute the 2D matrix of distances between all the embeddings.
    Args:
        embeddings: tensor of shape (batch_size, embed_dim)
        squared: Boolean. If true, output is the pairwise squared euclidean distance matrix.
                 If false, output is the pairwise euclidean distance matrix.
    Returns:
        pairwise_distances: tensor of shape (batch_size, batch_size)
    """
    # Get the dot product between all embeddings
    # shape (batch_size, batch_size)
    dot_product = tf.matmul(embeddings, tf.transpose(embeddings))

    # Get squared L2 norm for each embedding. We can just take the diagonal of `dot_product`.
    # This also provides more numerical stability (the diagonal of the result will be exactly 0).
    # shape (batch_size,)
    square_norm = tf.diag_part(dot_product)

    # Compute the pairwise distance matrix as we have:
    # ||a - b||^2 = ||a||^2  - 2 <a, b> + ||b||^2
    # shape (batch_size, batch_size)
    distances = tf.expand_dims(square_norm, 1) - 2.0 * dot_product + tf.expand_dims(square_norm, 0)

    # Because of computation errors, some distances might be negative so we put everything >= 0.0
    distances = tf.maximum(distances, 0.0)

    if not squared:
        # Because the gradient of sqrt is infinite when distances == 0.0 (ex: on the diagonal)
        # we need to add a small epsilon where distances == 0.0
        mask = tf.to_float(tf.equal(distances, 0.0))
        distances = distances + mask * 1e-16

        distances = tf.sqrt(distances)

        # Correct the epsilon added: set the distances on the mask to be exactly 0.0
        distances = distances * (1.0 - mask)

    return distances
//  [python]
def _pairwise_distances(embeddings, squared=False):
    """Compute the 2D matrix of distances between all the embeddings.
    Args:
        embeddings: tensor of shape (batch_size, embed_dim)
        squared: Boolean. If true, output is the pairwise squared euclidean distance matrix.
                 If false, output is the pairwise euclidean distance matrix.
    Returns:
        pairwise_distances: tensor of shape (batch_size, batch_size)
    """
    # Get the dot product between all embeddings
    # shape (batch_size, batch_size)
    dot_product = tf.matmul(embeddings, tf.transpose(embeddings))

    # Get squared L2 norm for each embedding. We can just take the diagonal of `dot_product`.
    # This also provides more numerical stability (the diagonal of the result will be exactly 0).
    # shape (batch_size,)
    square_norm = tf.diag_part(dot_product)

    # Compute the pairwise distance matrix as we have:
    # ||a - b||^2 = ||a||^2  - 2 <a, b> + ||b||^2
    # shape (batch_size, batch_size)
    distances = tf.expand_dims(square_norm, 1) - 2.0 * dot_product + tf.expand_dims(square_norm, 0)

    # Because of computation errors, some distances might be negative so we put everything >= 0.0
    distances = tf.maximum(distances, 0.0)

    if not squared:
        # Because the gradient of sqrt is infinite when distances == 0.0 (ex: on the diagonal)
        # we need to add a small epsilon where distances == 0.0
        mask = tf.cast(tf.equal(distances, 0.0), tf.float32)
        distances = distances + mask * 1e-16

        distances = tf.sqrt(distances)

        # Correct the epsilon added: set the distances on the mask to be exactly 0.0
        distances = distances * (1.0 - mask)

    return distances
//  [python]
def batch_hard_triplet_loss(labels, embeddings, margin, squared=False):
    """Build the triplet loss over a batch of embeddings.
    For each anchor, we get the hardest positive and hardest negative to form a triplet.
    Args:
        labels: labels of the batch, of size (batch_size,)
        embeddings: tensor of shape (batch_size, embed_dim)
        margin: margin for triplet loss
        squared: Boolean. If true, output is the pairwise squared euclidean distance matrix.
                 If false, output is the pairwise euclidean distance matrix.
    Returns:
        triplet_loss: scalar tensor containing the triplet loss
    """
    # Get the pairwise distance matrix
    pairwise_dist = _pairwise_distances(embeddings, squared=squared)

    # For each anchor, get the hardest positive
    # First, we need to get a mask for every valid positive (they should have same label)
    mask_anchor_positive = _get_anchor_positive_triplet_mask(labels)
    mask_anchor_positive = tf.to_float(mask_anchor_positive)

    # We put to 0 any element where (a, p) is not valid (valid if a != p and label(a) == label(p))
    anchor_positive_dist = tf.multiply(mask_anchor_positive, pairwise_dist)

    # shape (batch_size, 1)
    hardest_positive_dist = tf.reduce_max(anchor_positive_dist, axis=1, keepdims=True)
    tf.summary.scalar("hardest_positive_dist", tf.reduce_mean(hardest_positive_dist))

    # For each anchor, get the hardest negative
    # First, we need to get a mask for every valid negative (they should have different labels)
    mask_anchor_negative = _get_anchor_negative_triplet_mask(labels)
    mask_anchor_negative = tf.to_float(mask_anchor_negative)

    # We add the maximum value in each row to the invalid negatives (label(a) == label(n))
    max_anchor_negative_dist = tf.reduce_max(pairwise_dist, axis=1, keepdims=True)
    anchor_negative_dist = pairwise_dist + max_anchor_negative_dist * (1.0 - mask_anchor_negative)

    # shape (batch_size,)
    hardest_negative_dist = tf.reduce_min(anchor_negative_dist, axis=1, keepdims=True)
    tf.summary.scalar("hardest_negative_dist", tf.reduce_mean(hardest_negative_dist))

    # Combine biggest d(a, p) and smallest d(a, n) into final triplet loss
    triplet_loss = tf.maximum(hardest_positive_dist - hardest_negative_dist + margin, 0.0)

    # Get final mean triplet loss
    triplet_loss = tf.reduce_mean(triplet_loss)

    return triplet_loss
//  [python]
def batch_hard_triplet_loss(labels, embeddings, margin, squared=False):
    """Build the triplet loss over a batch of embeddings.
    For each anchor, we get the hardest positive and hardest negative to form a triplet.
    Args:
        labels: labels of the batch, of size (batch_size,)
        embeddings: tensor of shape (batch_size, embed_dim)
        margin: margin for triplet loss
        squared: Boolean. If true, output is the pairwise squared euclidean distance matrix.
                 If false, output is the pairwise euclidean distance matrix.
    Returns:
        triplet_loss: scalar tensor containing the triplet loss
    """
    # Get the pairwise distance matrix
    pairwise_dist = _pairwise_distances(embeddings, squared=squared)

    # For each anchor, get the hardest positive
    # First, we need to get a mask for every valid positive (they should have same label)
    mask_anchor_positive = _get_anchor_positive_triplet_mask(labels)
    mask_anchor_positive = tf.cast(mask_anchor_positive, tf.float32)

    # We put to 0 any element where (a, p) is not valid (valid if a != p and label(a) == label(p))
    anchor_positive_dist = tf.multiply(mask_anchor_positive, pairwise_dist)

    # shape (batch_size, 1)
    hardest_positive_dist = tf.reduce_max(anchor_positive_dist, axis=1, keepdims=True)
    tf.summary.scalar("hardest_positive_dist", tf.reduce_mean(hardest_positive_dist))

    # For each anchor, get the hardest negative
    # First, we need to get a mask for every valid negative (they should have different labels)
    mask_anchor_negative = _get_anchor_negative_triplet_mask(labels)
    mask_anchor_negative = tf.cast(mask_anchor_negative, tf.float32)

    # We add the maximum value in each row to the invalid negatives (label(a) == label(n))
    max_anchor_negative_dist = tf.reduce_max(pairwise_dist, axis=1, keepdims=True)
    anchor_negative_dist = pairwise_dist + max_anchor_negative_dist * (1.0 - mask_anchor_negative)

    # shape (batch_size,)
    hardest_negative_dist = tf.reduce_min(anchor_negative_dist, axis=1, keepdims=True)
    tf.summary.scalar("hardest_negative_dist", tf.reduce_mean(hardest_negative_dist))

    # Combine biggest d(a, p) and smallest d(a, n) into final triplet loss
    triplet_loss = tf.maximum(hardest_positive_dist - hardest_negative_dist + margin, 0.0)

    # Get final mean triplet loss
    triplet_loss = tf.reduce_mean(triplet_loss)

    return triplet_loss
//  [python]
def compute_weight_by_frequency(labels):
  """Computes weights to keep positive/negative labels in balance.

   For instance if labels contains 1 positive out of 9 negative, the positive
   one will have weight 0.9, while negative will have weight 0.1

  Args:
    labels: +1/-1 tensor of shape [..., num_channels]

  Returns:
    tensor of the same shape as labels.
  """
  p = tf.greater(labels, 0)
  pf = tf.to_float(p)
  positives = tf.reduce_sum(pf, axis=-1, keepdims=True) + tf.zeros_like(pf)
  negatives = tf.reduce_sum(1 - pf, axis=-1, keepdims=True) + tf.zeros_like(pf)
  total = positives + negatives
  weights = tf.where(p, negatives / total, positives / total)
  return weights
//  [python]
def compute_weight_by_frequency(labels):
  """Computes weights to keep positive/negative labels in balance.

   For instance if labels contains 1 positive out of 9 negative, the positive
   one will have weight 0.9, while negative will have weight 0.1

  Args:
    labels: +1/-1 tensor of shape [..., num_channels]

  Returns:
    tensor of the same shape as labels.
  """
  p = tf.greater(labels, 0)
  pf = tf.cast(p, tf.float32)
  positives = tf.reduce_sum(pf, axis=-1, keepdims=True) + tf.zeros_like(pf)
  negatives = tf.reduce_sum(1 - pf, axis=-1, keepdims=True) + tf.zeros_like(pf)
  total = positives + negatives
  weights = tf.where(p, negatives / total, positives / total)
  return weights
//  [python]
def blend(image1, image2, factor):
  """Blend image1 and image2 using 'factor'.

  Factor can be above 0.0.  A value of 0.0 means only image1 is used.
  A value of 1.0 means only image2 is used.  A value between 0.0 and
  1.0 means we linearly interpolate the pixel values between the two
  images.  A value greater than 1.0 "extrapolates" the difference
  between the two pixel values, and we clip the results to values
  between 0 and 255.
  Args:
    image1: An image Tensor of type uint8.
    image2: An image Tensor of type uint8.
    factor: A floating point value above 0.0.
  Returns:
    A blended image Tensor of type uint8.
  """
  if factor == 0.0:
    return tf.convert_to_tensor(image1)
  if factor == 1.0:
    return tf.convert_to_tensor(image2)

  image1 = tf.to_float(image1)
  image2 = tf.to_float(image2)

  difference = image2 - image1
  scaled = factor * difference

  # Do addition in float.
  temp = tf.to_float(image1) + scaled

  # Interpolate
  if factor > 0.0 and factor < 1.0:
    # Interpolation means we always stay within 0 and 255.
    return tf.cast(temp, tf.uint8)

  # Extrapolate:
  #
  # We need to clip and then cast.
  return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
//  [python]
def blend(image1, image2, factor):
  """Blend image1 and image2 using 'factor'.

  Factor can be above 0.0.  A value of 0.0 means only image1 is used.
  A value of 1.0 means only image2 is used.  A value between 0.0 and
  1.0 means we linearly interpolate the pixel values between the two
  images.  A value greater than 1.0 "extrapolates" the difference
  between the two pixel values, and we clip the results to values
  between 0 and 255.
  Args:
    image1: An image Tensor of type uint8.
    image2: An image Tensor of type uint8.
    factor: A floating point value above 0.0.
  Returns:
    A blended image Tensor of type uint8.
  """
  if factor == 0.0:
    return tf.convert_to_tensor(image1)
  if factor == 1.0:
    return tf.convert_to_tensor(image2)

  image1 = tf.cast(image1, tf.float32)
  image2 = tf.cast(image2, tf.float32)

  difference = image2 - image1
  scaled = factor * difference

  # Do addition in float.
  temp = tf.cast(image1, tf.float32) + scaled

  # Interpolate
  if factor > 0.0 and factor < 1.0:
    # Interpolation means we always stay within 0 and 255.
    return tf.cast(temp, tf.uint8)

  # Extrapolate:
  #
  # We need to clip and then cast.
  return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
//  [python]
def autocontrast(image):
  """Implements Autocontrast function from PIL using TF ops.

  Args:
    image: A 3D uint8 tensor.
  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """

  def scale_channel(image):
    """Scale the 2D image using the autocontrast rule."""
    # A possibly cheaper version can be done using cumsum/unique_with_counts
    # over the histogram values, rather than iterating over the entire image.
    # to compute mins and maxes.
    lo = tf.to_float(tf.reduce_min(image))
    hi = tf.to_float(tf.reduce_max(image))

    # Scale the image, making the lowest value 0 and the highest value 255.
    def scale_values(im):
      scale = 255.0 / (hi - lo)
      offset = -lo * scale
      im = tf.to_float(im) * scale + offset
      im = tf.clip_by_value(im, 0.0, 255.0)
      return tf.cast(im, tf.uint8)

    result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)
    return result

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = tf.stack([s1, s2, s3], 2)
  return image