jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def make_edge_vectors(adjacency_matrix, num_edge_types, depth, name=None):
  """Gets edge vectors for the edge types in the adjacency matrix.

  Args:
    adjacency_matrix: A [batch, num_nodes, num_nodes] tensor of ints.
    num_edge_types: Number of different edge types
    depth: Number of channels
    name: a string
  Returns:
    A [batch, num_nodes, num_nodes, depth] vector of tensors
  """
  with tf.variable_scope(name, default_name="edge_vectors"):
    att_adj_vectors_shape = [num_edge_types, depth]
    adjacency_matrix_shape = common_layers.shape_list(adjacency_matrix)
    adj_vectors = (
        tf.get_variable(
            "adj_vectors",
            att_adj_vectors_shape,
            initializer=tf.random_normal_initializer(0, depth**-0.5)) *
        (depth**0.5))
    # Avoiding gathers so that it works on TPUs
    # adjacency_matrix_one_hot has shape
    # [batch, num_nodes, num_nodes, num_edge_types]

    adjacency_matrix_one_hot = tf.one_hot(adjacency_matrix, num_edge_types)

    att_adj_vectors = tf.matmul(
        tf.reshape(tf.cast(adjacency_matrix_one_hot, tf.float32), [-1, num_edge_types]),
        adj_vectors)
    return tf.reshape(att_adj_vectors,
                      [adjacency_matrix_shape[0], adjacency_matrix_shape[1],
                       adjacency_matrix_shape[2], depth])
//  [python]
def add_timing_signal_1d_given_position(x,
                                        position,
                                        min_timescale=1.0,
                                        max_timescale=1.0e4):
  """Adds sinusoids of diff frequencies to a Tensor, with timing position given.

  Args:
    x: a Tensor with shape [batch, length, channels]
    position: a Tensor with shape [batch, length]
    min_timescale: a float
    max_timescale: a float

  Returns:
    a Tensor the same shape as x.
  """
  channels = common_layers.shape_list(x)[2]
  num_timescales = channels // 2
  log_timescale_increment = (
      math.log(float(max_timescale) / float(min_timescale)) /
      (tf.to_float(num_timescales) - 1))
  inv_timescales = min_timescale * tf.exp(
      tf.to_float(tf.range(num_timescales)) * -log_timescale_increment)
  scaled_time = (
      tf.expand_dims(tf.to_float(position), 2) * tf.expand_dims(
          tf.expand_dims(inv_timescales, 0), 0))
  signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=2)
  signal = tf.pad(signal, [[0, 0], [0, 0], [0, tf.mod(channels, 2)]])
  signal = common_layers.cast_like(signal, x)
  return x + signal
//  [python]
def add_timing_signal_1d_given_position(x,
                                        position,
                                        min_timescale=1.0,
                                        max_timescale=1.0e4):
  """Adds sinusoids of diff frequencies to a Tensor, with timing position given.

  Args:
    x: a Tensor with shape [batch, length, channels]
    position: a Tensor with shape [batch, length]
    min_timescale: a float
    max_timescale: a float

  Returns:
    a Tensor the same shape as x.
  """
  channels = common_layers.shape_list(x)[2]
  num_timescales = channels // 2
  log_timescale_increment = (
      math.log(float(max_timescale) / float(min_timescale)) /
      (tf.cast(num_timescales, dtype=tf.float32) - 1))
  inv_timescales = min_timescale * tf.exp(
      tf.cast(tf.range(num_timescales), dtype=tf.float32) * -log_timescale_increment)
  scaled_time = (
      tf.expand_dims(tf.cast(position, dtype=tf.float32), 2) * tf.expand_dims(
          tf.expand_dims(inv_timescales, 0), 0))
  signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=2)
  signal = tf.pad(signal, [[0, 0], [0, 0], [0, tf.mod(channels, 2)]])
  signal = common_layers.cast_like(signal, x)
  return x + signal
//  [python]
def attention_bias_same_segment(query_segment_id, memory_segment_id):
  """Create an bias tensor to be added to attention logits.

  Positions with the same segment_ids can see each other.

  Args:
    query_segment_id: a float `Tensor` with shape [batch, query_length].
    memory_segment_id: a float `Tensor` with shape [batch, memory_length].

  Returns:
    a `Tensor` with shape [batch, 1, query_length, memory_length].
  """
  ret = (tf.to_float(
      tf.not_equal(
          tf.expand_dims(query_segment_id, 2),
          tf.expand_dims(memory_segment_id, 1))) *
         large_compatible_negative(memory_segment_id.dtype))
  return tf.expand_dims(ret, axis=1)
//  [python]
def attention_bias_same_segment(query_segment_id, memory_segment_id):
  """Create an bias tensor to be added to attention logits.

  Positions with the same segment_ids can see each other.

  Args:
    query_segment_id: a float `Tensor` with shape [batch, query_length].
    memory_segment_id: a float `Tensor` with shape [batch, memory_length].

  Returns:
    a `Tensor` with shape [batch, 1, query_length, memory_length].
  """
  ret = (tf.cast(
      tf.not_equal(
          tf.expand_dims(query_segment_id, 2),
          tf.expand_dims(memory_segment_id, 1)), tf.float32) *
         large_compatible_negative(memory_segment_id.dtype))
  return tf.expand_dims(ret, axis=1)
//  [python]
def attention_bias_proximal(length):
  """Bias for self-attention to encourage attention to close positions.

  Args:
    length: an integer scalar.

  Returns:
    a Tensor with shape [1, 1, length, length]
  """
  r = tf.to_float(tf.range(length))
  diff = tf.expand_dims(r, 0) - tf.expand_dims(r, 1)
  return tf.expand_dims(tf.expand_dims(-tf.log1p(tf.abs(diff)), 0), 0)
//  [python]
def attention_bias_proximal(length):
  """Bias for self-attention to encourage attention to close positions.

  Args:
    length: an integer scalar.

  Returns:
    a Tensor with shape [1, 1, length, length]
  """
  r = tf.cast(tf.range(length), dtype=tf.float32)
  diff = tf.expand_dims(r, 0) - tf.expand_dims(r, 1)
  return tf.expand_dims(tf.expand_dims(-tf.log1p(tf.abs(diff)), 0), 0)
//  [python]
def attention_bias_prepend_inputs_full_attention(padding):
  """Create a bias tensor for prepend_mode="prepend_inputs_full_attention".

  See prepend_inputs in common_hparams.py.

  Produces a bias tensor to be used in self-attention.

  This bias tensor allows for full connectivity in the "inputs" part of
  the sequence and masked connectivity in the targets part.

  Args:
    padding: a float `Tensor` with shape [batch, length] with
      ones in positions corresponding to padding.  In each row, a single
      padding position separates the input part from the target part.

  Returns:
    a `Tensor` with shape [batch, 1, length, length].
  """
  # Everything past the first padding position is part of the target.
  # This Tensor has zeros for the source portion and separator,
  # and ones for the target portion.
  in_target = tf.cumsum(padding, axis=1, exclusive=True)
  # The position within the target, or 0 if part of the source.
  target_pos = tf.cumsum(in_target, axis=1)
  # A position with a lesser target_pos cannot see a position with greater
  # target_pos.
  illegal_connections = tf.greater(
      tf.expand_dims(target_pos, 1), tf.expand_dims(target_pos, 2))
  bias = tf.to_float(illegal_connections) * -1e9
  bias = tf.expand_dims(bias, 1)
  return bias
//  [python]
def attention_bias_prepend_inputs_full_attention(padding):
  """Create a bias tensor for prepend_mode="prepend_inputs_full_attention".

  See prepend_inputs in common_hparams.py.

  Produces a bias tensor to be used in self-attention.

  This bias tensor allows for full connectivity in the "inputs" part of
  the sequence and masked connectivity in the targets part.

  Args:
    padding: a float `Tensor` with shape [batch, length] with
      ones in positions corresponding to padding.  In each row, a single
      padding position separates the input part from the target part.

  Returns:
    a `Tensor` with shape [batch, 1, length, length].
  """
  # Everything past the first padding position is part of the target.
  # This Tensor has zeros for the source portion and separator,
  # and ones for the target portion.
  in_target = tf.cumsum(padding, axis=1, exclusive=True)
  # The position within the target, or 0 if part of the source.
  target_pos = tf.cumsum(in_target, axis=1)
  # A position with a lesser target_pos cannot see a position with greater
  # target_pos.
  illegal_connections = tf.greater(
      tf.expand_dims(target_pos, 1), tf.expand_dims(target_pos, 2))
  bias = tf.cast(illegal_connections, tf.float32) * -1e9
  bias = tf.expand_dims(bias, 1)
  return bias
//  [python]
def attention_bias_batch(batch_coordinates_q,
                         batch_coordinates_k=None,
                         condition_fn=None):
  """Generate a mask to prevent the batch to attend to each others.

  Args:
    batch_coordinates_q: Int-like Tensor of shape [length_q, 1] containing the
      coordinates of the batches
    batch_coordinates_k: Int-like Tensor of shape [length_k, 1] containing the
      coordinates of the batches. If None, do self-attention.
    condition_fn: Callable defining the attention mask.

  Returns:
    Float-like Tensor of shape [length_q, length_k] containing either 0 or
    -infinity (-1e9).
  """
  if batch_coordinates_k is None:
    batch_coordinates_k = batch_coordinates_q

  # Convert to float first because of b/25387198.
  def to_float(bc):
    bc = tf.squeeze(bc, 1)
    bc = tf.to_float(bc)
    return bc

  # Broadcast to create [length_q, length_k] mask.
  bc_v = tf.expand_dims(to_float(batch_coordinates_q), 1)
  bc_h = tf.expand_dims(to_float(batch_coordinates_k), 0)
  bias_batch = bc_h - bc_v
  bias_batch = condition_fn(bias_batch)
  bias_batch *= -1e9
  return bias_batch
//  [python]
def attention_bias_batch(batch_coordinates_q,
                         batch_coordinates_k=None,
                         condition_fn=None):
  """Generate a mask to prevent the batch to attend to each others.

  Args:
    batch_coordinates_q: Int-like Tensor of shape [length_q, 1] containing the
      coordinates of the batches
    batch_coordinates_k: Int-like Tensor of shape [length_k, 1] containing the
      coordinates of the batches. If None, do self-attention.
    condition_fn: Callable defining the attention mask.

  Returns:
    Float-like Tensor of shape [length_q, length_k] containing either 0 or
    -infinity (-1e9).
  """
  if batch_coordinates_k is None:
    batch_coordinates_k = batch_coordinates_q

  # Convert to float first because of b/25387198.
  def to_float(bc):
    bc = tf.squeeze(bc, 1)
    bc = tf.cast(bc, tf.float32)
    return bc

  # Broadcast to create [length_q, length_k] mask.
  bc_v = tf.expand_dims(to_float(batch_coordinates_q), 1)
  bc_h = tf.expand_dims(to_float(batch_coordinates_k), 0)
  bias_batch = bc_h - bc_v
  bias_batch = condition_fn(bias_batch)
  bias_batch *= -1e9
  return bias_batch
//  [python]
def scaled_dot_product_attention_simple(q, k, v, bias, name=None):
  """Scaled dot-product attention. One head. One spatial dimension.

  Args:
    q: a Tensor with shape [batch, length_q, depth_k]
    k: a Tensor with shape [batch, length_kv, depth_k]
    v: a Tensor with shape [batch, length_kv, depth_v]
    bias: optional Tensor broadcastable to [batch, length_q, length_kv]
    name: an optional string

  Returns:
    A Tensor.
  """
  with tf.variable_scope(
      name, default_name="scaled_dot_product_attention_simple"):
    scalar = tf.rsqrt(tf.to_float(common_layers.shape_list(q)[2]))
    logits = tf.matmul(q * scalar, k, transpose_b=True)
    if bias is not None:
      logits += bias
    weights = tf.nn.softmax(logits, name="attention_weights")
    if common_layers.should_generate_summaries():
      tf.summary.image(
          "attention", tf.expand_dims(tf.pow(weights, 0.2), 3), max_outputs=1)
    return tf.matmul(weights, v)
//  [python]
def scaled_dot_product_attention_simple(q, k, v, bias, name=None):
  """Scaled dot-product attention. One head. One spatial dimension.

  Args:
    q: a Tensor with shape [batch, length_q, depth_k]
    k: a Tensor with shape [batch, length_kv, depth_k]
    v: a Tensor with shape [batch, length_kv, depth_v]
    bias: optional Tensor broadcastable to [batch, length_q, length_kv]
    name: an optional string

  Returns:
    A Tensor.
  """
  with tf.variable_scope(
      name, default_name="scaled_dot_product_attention_simple"):
    scalar = tf.rsqrt(tf.cast(common_layers.shape_list(q)[2], tf.float32))
    logits = tf.matmul(q * scalar, k, transpose_b=True)
    if bias is not None:
      logits += bias
    weights = tf.nn.softmax(logits, name="attention_weights")
    if common_layers.should_generate_summaries():
      tf.summary.image(
          "attention", tf.expand_dims(tf.pow(weights, 0.2), 3), max_outputs=1)
    return tf.matmul(weights, v)
//  [python]
def fprop(self, x, **kwargs):
        alpha = 1.6732632423543772848170429916717
        scale = 1.0507009873554804934193349852946
        mask = tf.to_float(x >= 0.0)
        out = mask * x + (1.0 - mask) * (alpha * tf.exp((1.0 - mask) * x) - alpha)
        return scale * out
//  [python]
def fprop(self, x, **kwargs):
        alpha = 1.6732632423543772848170429916717
        scale = 1.0507009873554804934193349852946
        mask = tf.cast(x >= 0.0, tf.float32)
        out = mask * x + (1.0 - mask) * (alpha * tf.exp((1.0 - mask) * x) - alpha)
        return scale * out
//  [python]
def randomly_pad(image, boxes, probability=0.9):
    """
    This function makes content of the image
    smaller by scaling and padding it with zeros.
    """

    def pad(image, boxes):

        shape = tf.shape(image)
        height, width = shape[0], shape[1]

        # randomly reduce image scale
        scale = tf.random_uniform([], 0.5, 0.9)
        scaled_height = tf.to_int32(scale * tf.to_float(height))
        scaled_width = tf.to_int32(scale * tf.to_float(width))

        image = tf.image.resize_images(
            image, [scaled_height, scaled_width],
            method=RESIZE_METHOD
        )

        # randomly pad to the initial size
        offset_y = height - scaled_height
        offset_x = width - scaled_width
        offset_y = tf.random_uniform([], 0, offset_y, dtype=tf.int32)
        offset_x = tf.random_uniform([], 0, offset_x, dtype=tf.int32)
        image = tf.image.pad_to_bounding_box(image, offset_y, offset_x, height, width)

        # transform boxes
        boxes *= scale
        offset_y = tf.to_float(offset_y/height)
        offset_x = tf.to_float(offset_x/width)
        translation = tf.stack([offset_y, offset_x, offset_y, offset_x])
        boxes += translation

        return image, boxes

    do_it = tf.less(tf.random_uniform([]), probability)
    image, boxes = tf.cond(do_it, lambda: pad(image, boxes), lambda: (image, boxes))
    return image, boxes
//  [python]
def randomly_pad(image, boxes, probability=0.9):
    """
    This function makes content of the image
    smaller by scaling and padding it with zeros.
    """

    def pad(image, boxes):

        shape = tf.shape(image)
        height, width = shape[0], shape[1]

        # randomly reduce image scale
        scale = tf.random_uniform([], 0.5, 0.9)
        scaled_height = tf.cast(scale * tf.to_float(height), tf.int32)
        scaled_width = tf.cast(scale * tf.to_float(width), tf.int32)

        image = tf.image.resize_images(
            image, [scaled_height, scaled_width],
            method=RESIZE_METHOD
        )

        # randomly pad to the initial size
        offset_y = height - scaled_height
        offset_x = width - scaled_width
        offset_y = tf.random_uniform([], 0, offset_y, dtype=tf.int32)
        offset_x = tf.random_uniform([], 0, offset_x, dtype=tf.int32)
        image = tf.image.pad_to_bounding_box(image, offset_y, offset_x, height, width)

        # transform boxes
        boxes *= scale
        offset_y = tf.to_float(offset_y/height)
        offset_x = tf.to_float(offset_x/width)
        translation = tf.stack([offset_y, offset_x, offset_y, offset_x])
        boxes += translation

        return image, boxes

    do_it = tf.less(tf.random_uniform([]), probability)
    image, boxes = tf.cond(do_it, lambda: pad(image, boxes), lambda: (image, boxes))
    return image, boxes
//  [python]
def parse_and_preprocess(self, example_proto):
        """
        Returns:
            image: a float tensor with shape [height, width, 3],
                an RGB image with pixel values in the range [0, 1].
            boxes: a float tensor with shape [num_boxes, 4].
            num_boxes: an int tensor with shape [].
        """
        features = {
            'image': tf.FixedLenFeature([], tf.string),
            'num_persons': tf.FixedLenFeature([], tf.int64),
            'boxes': tf.FixedLenSequenceFeature([], tf.float32, allow_missing=True)
        }
        parsed_features = tf.parse_single_example(example_proto, features)

        # get an image
        image = tf.image.decode_jpeg(parsed_features['image'], channels=3)
        image = tf.image.convert_image_dtype(image, tf.float32)
        # now pixel values are scaled to the [0, 1] range

        # get number of people on the image
        num_boxes = tf.to_int32(parsed_features['num_persons'])
        # it is assumed that num_boxes > 0

        # get groundtruth boxes, they are in absolute coordinates
        boxes = tf.reshape(parsed_features['boxes'], [num_boxes, 4])

        # to the [0, 1] range
        height, width = tf.shape(image)[0], tf.shape(image)[1]
        scaler = tf.to_float(tf.stack([height, width, height, width]))
        boxes /= scaler

        if self.is_training:
            image, boxes = augmentation(image, boxes, self.image_size)
        else:
            image, boxes = resize_keeping_aspect_ratio(image, boxes, self.min_dimension, DIVISOR)

        # it could change after augmentations
        num_boxes = tf.shape(boxes)[0]

        features = {'images': image}
        labels = {'boxes': boxes, 'num_boxes': num_boxes}
        return features, labels
//  [python]
def parse_and_preprocess(self, example_proto):
        """
        Returns:
            image: a float tensor with shape [height, width, 3],
                an RGB image with pixel values in the range [0, 1].
            boxes: a float tensor with shape [num_boxes, 4].
            num_boxes: an int tensor with shape [].
        """
        features = {
            'image': tf.FixedLenFeature([], tf.string),
            'num_persons': tf.FixedLenFeature([], tf.int64),
            'boxes': tf.FixedLenSequenceFeature([], tf.float32, allow_missing=True)
        }
        parsed_features = tf.parse_single_example(example_proto, features)

        # get an image
        image = tf.image.decode_jpeg(parsed_features['image'], channels=3)
        image = tf.image.convert_image_dtype(image, tf.float32)
        # now pixel values are scaled to the [0, 1] range

        # get number of people on the image
        num_boxes = tf.cast(parsed_features['num_persons'], tf.int32)
        # it is assumed that num_boxes > 0

        # get groundtruth boxes, they are in absolute coordinates
        boxes = tf.reshape(parsed_features['boxes'], [num_boxes, 4])

        # to the [0, 1] range
        height, width = tf.shape(image)[0], tf.shape(image)[1]
        scaler = tf.cast(tf.stack([height, width, height, width]), tf.float32)
        boxes /= scaler

        if self.is_training:
            image, boxes = augmentation(image, boxes, self.image_size)
        else:
            image, boxes = resize_keeping_aspect_ratio(image, boxes, self.min_dimension, DIVISOR)

        # it could change after augmentations
        num_boxes = tf.shape(boxes)[0]

        features = {'images': image}
        labels = {'boxes': boxes, 'num_boxes': num_boxes}
        return features, labels
//  [python]
def test_visualize_boxes_in_image(self):
    image = tf.zeros((6, 4, 3))
    corners = tf.constant([[0, 0, 5, 3],
                           [0, 0, 3, 2]], tf.float32)
    boxes = box_list.BoxList(corners)
    image_and_boxes = box_list_ops.visualize_boxes_in_image(image, boxes)
    image_and_boxes_bw = tf.to_float(
        tf.greater(tf.reduce_sum(image_and_boxes, 2), 0.0))
    exp_result = [[1, 1, 1, 0],
                  [1, 1, 1, 0],
                  [1, 1, 1, 0],
                  [1, 0, 1, 0],
                  [1, 1, 1, 0],
                  [0, 0, 0, 0]]
    with self.test_session() as sess:
      output = sess.run(image_and_boxes_bw)
      self.assertAllEqual(output.astype(int), exp_result)
//  [python]
def test_visualize_boxes_in_image(self):
    image = tf.zeros((6, 4, 3))
    corners = tf.constant([[0, 0, 5, 3],
                           [0, 0, 3, 2]], tf.float32)
    boxes = box_list.BoxList(corners)
    image_and_boxes = box_list_ops.visualize_boxes_in_image(image, boxes)
    image_and_boxes_bw = tf.cast(
        tf.greater(tf.reduce_sum(image_and_boxes, 2), 0.0), tf.float32)
    exp_result = [[1, 1, 1, 0],
                  [1, 1, 1, 0],
                  [1, 1, 1, 0],
                  [1, 0, 1, 0],
                  [1, 1, 1, 0],
                  [0, 0, 0, 0]]
    with self.test_session() as sess:
      output = sess.run(image_and_boxes_bw)
      self.assertAllEqual(output.astype(int), exp_result)
//  [python]
def resize_keeping_aspect_ratio(image, boxes, min_dimension, divisor):
    """
    This function resizes and possibly pads with zeros.
    When using a usual FPN, divisor must be equal to 128.

    Arguments:
        image: a float tensor with shape [height, width, 3].
        boxes: a float tensor with shape [n, 4].
        min_dimension: an integer.
        divisor: an integer.
    Returns:
        image: a float tensor with shape [h, w, 3],
            where `min_dimension = min(h, w)`,
            `h` and `w` are divisible by `divisor`.
        boxes: a float tensor with shape [n, 4].
    """
    assert min_dimension % divisor == 0

    min_dimension = tf.constant(min_dimension, dtype=tf.int32)
    divisor = tf.constant(divisor, dtype=tf.int32)

    shape = tf.shape(image)
    height, width = shape[0], shape[1]

    original_min_dim = tf.minimum(height, width)
    scale_factor = tf.to_float(min_dimension / original_min_dim)

    def scale(x):
        unpadded_x = tf.to_int32(tf.round(tf.to_float(x) * scale_factor))
        x = tf.to_int32(tf.ceil(unpadded_x / divisor))
        pad = divisor * x - unpadded_x
        return (unpadded_x, pad)

    zero = tf.constant(0, dtype=tf.int32)
    new_height, pad_height, new_width, pad_width = tf.cond(
        tf.greater_equal(height, width),
        lambda: scale(height) + (min_dimension, zero),
        lambda: (min_dimension, zero) + scale(width)
    )

    # resize keeping aspect ratio
    image = tf.image.resize_images(image, [new_height, new_width], method=RESIZE_METHOD)

    h = new_height + pad_height
    w = new_width + pad_width

    image = tf.image.pad_to_bounding_box(
        image, offset_height=0, offset_width=0,
        target_height=h, target_width=w
    )
    # it pads image at the bottom or at the right

    # we need to rescale bounding box coordinates
    box_scaler = tf.to_float(tf.stack([
        new_height/h, new_width/w,
        new_height/h, new_width/w
    ]))

    boxes *= box_scaler
    return image, boxes
//  [python]
def resize_keeping_aspect_ratio(image, boxes, min_dimension, divisor):
    """
    This function resizes and possibly pads with zeros.
    When using a usual FPN, divisor must be equal to 128.

    Arguments:
        image: a float tensor with shape [height, width, 3].
        boxes: a float tensor with shape [n, 4].
        min_dimension: an integer.
        divisor: an integer.
    Returns:
        image: a float tensor with shape [h, w, 3],
            where `min_dimension = min(h, w)`,
            `h` and `w` are divisible by `divisor`.
        boxes: a float tensor with shape [n, 4].
    """
    assert min_dimension % divisor == 0

    min_dimension = tf.constant(min_dimension, dtype=tf.int32)
    divisor = tf.constant(divisor, dtype=tf.int32)

    shape = tf.shape(image)
    height, width = shape[0], shape[1]

    original_min_dim = tf.minimum(height, width)
    scale_factor = tf.cast(min_dimension / original_min_dim, tf.float32)

    def scale(x):
        unpadded_x = tf.to_int32(tf.round(tf.cast(x, tf.float32) * scale_factor))
        x = tf.to_int32(tf.ceil(unpadded_x / divisor))
        pad = divisor * x - unpadded_x
        return (unpadded_x, pad)

    zero = tf.constant(0, dtype=tf.int32)
    new_height, pad_height, new_width, pad_width = tf.cond(
        tf.greater_equal(height, width),
        lambda: scale(height) + (min_dimension, zero),
        lambda: (min_dimension, zero) + scale(width)
    )

    # resize keeping aspect ratio
    image = tf.image.resize_images(image, [new_height, new_width], method=RESIZE_METHOD)

    h = new_height + pad_height
    w = new_width + pad_width

    image = tf.image.pad_to_bounding_box(
        image, offset_height=0, offset_width=0,
        target_height=h, target_width=w
    )
    # it pads image at the bottom or at the right

    # we need to rescale bounding box coordinates
    box_scaler = tf.cast(tf.stack([
        new_height/h, new_width/w,
        new_height/h, new_width/w
    ]), tf.float32)

    boxes *= box_scaler
    return image, boxes
//  [python]
def build_feature_extractor(input_):
  # We only want to create the weights once
  # In all future calls we should set reuse = True

  # scale the inputs from 0..255 to 0..1
  input_ = tf.to_float(input_) / 255.0

  # conv layers
  conv1 = tf.contrib.layers.conv2d(
    input_,
    16, # num output feature maps
    8,  # kernel size
    4,  # stride
    activation_fn=tf.nn.relu,
    scope="conv1")
  conv2 = tf.contrib.layers.conv2d(
    conv1,
    32, # num output feature maps
    4,  # kernel size
    2,  # stride
    activation_fn=tf.nn.relu,
    scope="conv2")

  # image -> feature vector
  flat = tf.contrib.layers.flatten(conv2)

  # dense layer
  fc1 = tf.contrib.layers.fully_connected(
    inputs=flat,
    num_outputs=256,
    scope="fc1")

  return fc1
//  [python]
def build_feature_extractor(input_):
  # We only want to create the weights once
  # In all future calls we should set reuse = True

  # scale the inputs from 0..255 to 0..1
  input_ = tf.cast(input_, tf.float32) / 255.0

  # conv layers
  conv1 = tf.contrib.layers.conv2d(
    input_,
    16, # num output feature maps
    8,  # kernel size
    4,  # stride
    activation_fn=tf.nn.relu,
    scope="conv1")
  conv2 = tf.contrib.layers.conv2d(
    conv1,
    32, # num output feature maps
    4,  # kernel size
    2,  # stride
    activation_fn=tf.nn.relu,
    scope="conv2")

  # image -> feature vector
  flat = tf.contrib.layers.flatten(conv2)

  # dense layer
  fc1 = tf.contrib.layers.fully_connected(
    inputs=flat,
    num_outputs=256,
    scope="fc1")

  return fc1
//  [python]
def dot_product_batched_head(q, k, v, gates_q, gates_k, mask_right=False):
  """Perform a dot product attention on a single sequence on a single head.

  This function dispatch the q, k, v and loop over the buckets to compute the
  attention dot product on each subsequences.

  Args:
    q (tf.Tensor): [batch*heads, length_q, depth_q]
    k (tf.Tensor): [batch*heads, length_k, depth_q]
    v (tf.Tensor): [batch*heads, length_k, depth_v]
    gates_q (tf.Tensor): One-hot of shape [batch*heads, length_q, nb_buckets]
    gates_k (tf.Tensor): One-hot of shape [batch*heads, length_k, nb_buckets]
    mask_right (bool): Add a bias to prevent attention to the future

  Returns:
    tf.Tensor: [length_q, depth_v]
  """
  nb_buckets = common_layers.shape_list(gates_q)[-1]

  @expert_utils.add_name_scope()
  def get_dispatcher(gates):
    """Construct dispatcher for gates."""
    length = common_layers.shape_list(gates)[1]
    # Count the number of ones per batch (and keep the max value)
    nb_elems_to_dispatch = tf.reduce_sum(gates, axis=[1, 2])
    nb_elems_to_dispatch = tf.reduce_max(nb_elems_to_dispatch)
    nb_elems_to_dispatch = tf.to_int32(nb_elems_to_dispatch)
    capacity = nb_elems_to_dispatch // nb_buckets * 2  # Capacity is hardcoded
    capacity = tf.minimum(length, capacity)
    tf.summary.scalar("dispatch_capacity", capacity, family="lsh")
    return expert_utils.TruncatingDispatcher(gates, capacity)

  def add_summary_capacity(x, prefix):
    # Monitor if capacity overflow
    x = x[0, ...]  # Take first batch/head
    x = tf.reduce_sum(x, axis=0)
    tf.summary.scalar(prefix + "_min", tf.reduce_min(x), family="lsh")
    tf.summary.scalar(prefix + "_max", tf.reduce_max(x), family="lsh")
    tf.summary.histogram(prefix + "capacity_distribution", x, family="lsh")
    for i in range(3):  # Show the first 3 buckets
      tf.summary.scalar("{}_{}".format(prefix, i), x[i], family="lsh")

  add_summary_capacity(gates_q, "q")
  add_summary_capacity(gates_k, "k")

  q_dispatcher = get_dispatcher(gates_q)
  k_dispatcher = get_dispatcher(gates_k)

  q = q_dispatcher.dispatch(q)
  k = k_dispatcher.dispatch(k)
  v = k_dispatcher.dispatch(v)

  # Bias of shape [batch*heads, nb_buckets, 1, capacity] broadcasted to every
  # queries
  bias = tf.expand_dims((k_dispatcher.nonpadding() - 1.0) * 1e9, 2)
  if mask_right:
    q_coordinate = tf.to_float(
        tf.expand_dims(q_dispatcher.length_coordinate(), 3))
    k_coordinate = tf.to_float(
        tf.expand_dims(k_dispatcher.length_coordinate(), 2))
    bias += tf.to_float(tf.greater(k_coordinate, q_coordinate)) * -1e9
  # The sequence padding is not masked but is ignored on the next layers

  # q, k, v now have shape [batch*heads, nb_bucket, capacity, depth]
  # The buckets can be seen as different heads
  v_out = dot_product_attention(q, k, v, bias=bias)

  # Combine all buckets together to restore the original length
  return q_dispatcher.combine(v_out)
//  [python]
def dot_product_batched_head(q, k, v, gates_q, gates_k, mask_right=False):
  """Perform a dot product attention on a single sequence on a single head.

  This function dispatch the q, k, v and loop over the buckets to compute the
  attention dot product on each subsequences.

  Args:
    q (tf.Tensor): [batch*heads, length_q, depth_q]
    k (tf.Tensor): [batch*heads, length_k, depth_q]
    v (tf.Tensor): [batch*heads, length_k, depth_v]
    gates_q (tf.Tensor): One-hot of shape [batch*heads, length_q, nb_buckets]
    gates_k (tf.Tensor): One-hot of shape [batch*heads, length_k, nb_buckets]
    mask_right (bool): Add a bias to prevent attention to the future

  Returns:
    tf.Tensor: [length_q, depth_v]
  """
  nb_buckets = common_layers.shape_list(gates_q)[-1]

  @expert_utils.add_name_scope()
  def get_dispatcher(gates):
    """Construct dispatcher for gates."""
    length = common_layers.shape_list(gates)[1]
    # Count the number of ones per batch (and keep the max value)
    nb_elems_to_dispatch = tf.reduce_sum(gates, axis=[1, 2])
    nb_elems_to_dispatch = tf.reduce_max(nb_elems_to_dispatch)
    nb_elems_to_dispatch = tf.cast(nb_elems_to_dispatch, tf.int32)
    capacity = nb_elems_to_dispatch // nb_buckets * 2  # Capacity is hardcoded
    capacity = tf.minimum(length, capacity)
    tf.summary.scalar("dispatch_capacity", capacity, family="lsh")
    return expert_utils.TruncatingDispatcher(gates, capacity)

  def add_summary_capacity(x, prefix):
    # Monitor if capacity overflow
    x = x[0, ...]  # Take first batch/head
    x = tf.reduce_sum(x, axis=0)
    tf.summary.scalar(prefix + "_min", tf.reduce_min(x), family="lsh")
    tf.summary.scalar(prefix + "_max", tf.reduce_max(x), family="lsh")
    tf.summary.histogram(prefix + "capacity_distribution", x, family="lsh")
    for i in range(3):  # Show the first 3 buckets
      tf.summary.scalar("{}_{}".format(prefix, i), x[i], family="lsh")

  add_summary_capacity(gates_q, "q")
  add_summary_capacity(gates_k, "k")

  q_dispatcher = get_dispatcher(gates_q)
  k_dispatcher = get_dispatcher(gates_k)

  q = q_dispatcher.dispatch(q)
  k = k_dispatcher.dispatch(k)
  v = k_dispatcher.dispatch(v)

  # Bias of shape [batch*heads, nb_buckets, 1, capacity] broadcasted to every
  # queries
  bias = tf.expand_dims((k_dispatcher.nonpadding() - 1.0) * 1e9, 2)
  if mask_right:
    q_coordinate = tf.cast(
        tf.expand_dims(q_dispatcher.length_coordinate(), 3), tf.float32)
    k_coordinate = tf.cast(
        tf.expand_dims(k_dispatcher.length_coordinate(), 2), tf.float32)
    bias += tf.cast(tf.greater(k_coordinate, q_coordinate), tf.float32) * -1e9
  # The sequence padding is not masked but is ignored on the next layers

  # q, k, v now have shape [batch*heads, nb_bucket, capacity, depth]
  # The buckets can be seen as different heads
  v_out = dot_product_attention(q, k, v, bias=bias)

  # Combine all buckets together to restore the original length
  return q_dispatcher.combine(v_out)
//  [python]
def bottleneck(self, x):
    hparams = self.hparams
    x = tf.tanh(tf.layers.dense(x, hparams.bottleneck_bits, name="bottleneck"))
    d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
      noise = tf.random_uniform(common_layers.shape_list(x))
      noise = 2.0 * tf.to_float(tf.less(hparams.bottleneck_noise, noise)) - 1.0
      d *= noise
    x = common_layers.mix(d, x, hparams.discretize_warmup_steps,
                          hparams.mode == tf.estimator.ModeKeys.TRAIN)
    return x, 0.0
//  [python]
def bottleneck(self, x):
    hparams = self.hparams
    x = tf.tanh(tf.layers.dense(x, hparams.bottleneck_bits, name="bottleneck"))
    d = x + tf.stop_gradient(2.0 * tf.cast(tf.less(0.0, x), tf.float32) - 1.0 - x)
    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
      noise = tf.random_uniform(common_layers.shape_list(x))
      noise = 2.0 * tf.cast(tf.less(hparams.bottleneck_noise, noise), tf.float32) - 1.0
      d *= noise
    x = common_layers.mix(d, x, hparams.discretize_warmup_steps,
                          hparams.mode == tf.estimator.ModeKeys.TRAIN)
    return x, 0.0
//  [python]
def _aspect_preserving_resize(self, image, target_smaller_side):
        """"Resize image such that the smaller size has size
        ``target_smaller_sider`` while preserving the aspect ratio.

        Args:
            image (tf.Tensor): Tensor containing the image to resize.
            target_smaller_side (int): Target size for the smaller side in pixel.

        Returns:
            tf.Tensor: The resized image, with the same aspect ratio as the input.

        """

        shape = tf.shape(image)
        height = tf.to_float(shape[0])
        width = tf.to_float(shape[1])
        smaller_side = tf.reduce_min(shape[0:2])
        scale = tf.divide(target_smaller_side, tf.to_float(smaller_side))
        new_height = tf.to_int32(tf.round(scale * height))
        new_width = tf.to_int32(tf.round(scale * width))
        resized_image = tf.image.resize_images(image, [new_height, new_width])

        return resized_image
//  [python]
def _aspect_preserving_resize(self, image, target_smaller_side):
        """"Resize image such that the smaller size has size
        ``target_smaller_sider`` while preserving the aspect ratio.

        Args:
            image (tf.Tensor): Tensor containing the image to resize.
            target_smaller_side (int): Target size for the smaller side in pixel.

        Returns:
            tf.Tensor: The resized image, with the same aspect ratio as the input.

        """

        shape = tf.shape(image)
        height = tf.cast(shape[0], tf.float32)
        width = tf.cast(shape[1], tf.float32)
        smaller_side = tf.reduce_min(shape[0:2])
        scale = tf.divide(target_smaller_side, tf.cast(smaller_side, tf.float32))
        new_height = tf.to_int32(tf.round(scale * height))
        new_width = tf.to_int32(tf.round(scale * width))
        resized_image = tf.image.resize_images(image, [new_height, new_width])

        return resized_image
//  [python]
def variance_loss(self, b):
    part = tf.random_uniform(common_layers.shape_list(b))
    selection = tf.to_float(tf.less(part, tf.random_uniform([])))
    selection_size = tf.reduce_sum(selection)
    part_avg = tf.abs(tf.reduce_sum(b * selection)) / (selection_size + 1)
    return part_avg
//  [python]
def variance_loss(self, b):
    part = tf.random_uniform(common_layers.shape_list(b))
    selection = tf.cast(tf.less(part, tf.random_uniform([])), tf.float32)
    selection_size = tf.reduce_sum(selection)
    part_avg = tf.abs(tf.reduce_sum(b * selection)) / (selection_size + 1)
    return part_avg
//  [python]
def sample(self, features=None, shape=None):
    del features
    hp = self.hparams
    div_x = 2**hp.num_hidden_layers
    div_y = 1 if self.is1d else 2**hp.num_hidden_layers
    size = [
        hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
        hp.bottleneck_bits
    ]
    size = size if shape is None else shape
    rand = tf.random_uniform(size)
    return 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
//  [python]
def sample(self, features=None, shape=None):
    del features
    hp = self.hparams
    div_x = 2**hp.num_hidden_layers
    div_y = 1 if self.is1d else 2**hp.num_hidden_layers
    size = [
        hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
        hp.bottleneck_bits
    ]
    size = size if shape is None else shape
    rand = tf.random_uniform(size)
    return 2.0 * tf.cast(tf.less(0.5, rand), tf.float32) - 1.0
//  [python]
def stack_loss(self, b, b_pred, name):
    with tf.variable_scope(name):
      labels_discrete = tf.to_int32((b + 1.0) * 0.5)
      loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
          labels=labels_discrete, logits=b_pred)
      return tf.reduce_mean(loss)
//  [python]
def stack_loss(self, b, b_pred, name):
    with tf.variable_scope(name):
      labels_discrete = tf.cast((b + 1.0) * 0.5, tf.int32)
      loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
          labels=labels_discrete, logits=b_pred)
      return tf.reduce_mean(loss)
//  [python]
def sample(self, features=None, shape=None):
    del features
    hp = self.hparams
    div_x = 2**hp.num_hidden_layers
    div_y = 1 if self.is1d else 2**hp.num_hidden_layers
    size = [
        hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
        hp.bottleneck_bits
    ]
    size = size if shape is None else shape
    rand = tf.random_uniform(size)
    res = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
    # If you want to set some first bits to a fixed value, do this:
    # fixed = tf.zeros_like(rand) - 1.0
    # nbits = 3
    # res = tf.concat([fixed[:, :, :, :nbits], res[:, :, :, nbits:]], axis=-1)
    return res
//  [python]
def sample(self, features=None, shape=None):
    del features
    hp = self.hparams
    div_x = 2**hp.num_hidden_layers
    div_y = 1 if self.is1d else 2**hp.num_hidden_layers
    size = [
        hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
        hp.bottleneck_bits
    ]
    size = size if shape is None else shape
    rand = tf.random_uniform(size)
    res = 2.0 * tf.cast(tf.less(0.5, rand), tf.float32) - 1.0
    # If you want to set some first bits to a fixed value, do this:
    # fixed = tf.zeros_like(rand) - 1.0
    # nbits = 3
    # res = tf.concat([fixed[:, :, :, :nbits], res[:, :, :, nbits:]], axis=-1)
    return res
//  [python]
def apply_line_prediction(inputs,
                          features,
                          blur_steps,
                          learn_alpha=True,
                          name=None):
  """Applies "Line Prediction" layer to input images."""
  inputs.shape.assert_is_compatible_with([None, None, None, 6])

  with tf.name_scope(name, 'blur_prediction', values=[inputs, features]):

    with tf.name_scope(None, 'input_frames', values=[inputs]):
      frames = [inputs[:, :, :, :3], inputs[:, :, :, 3:]]

    with tf.name_scope(None, 'frame_size', values=[inputs, features]):
      shape = tf.shape(inputs)
      height = shape[1]
      width = shape[2]

    with tf.name_scope(None, 'identity_warp', values=[]):
      x_idx, y_idx = tf.meshgrid(tf.range(width), tf.range(height))
      identity_warp = tf.to_float(tf.stack([x_idx, y_idx], axis=-1))
      identity_warp = identity_warp[tf.newaxis, :, :, tf.newaxis, :]

      warp_steps = tf.to_float(tf.range(blur_steps - 1) + 1) / (blur_steps - 1)
      warp_steps = warp_steps[tf.newaxis, tf.newaxis, tf.newaxis, :, tf.newaxis]

      max_warps = tf.to_float(tf.stack([width - 1, height - 1]))
      max_warps = max_warps[tf.newaxis, tf.newaxis, tf.newaxis, tf.newaxis, :]

    output_frames = []
    for frame in frames:
      with tf.name_scope(None, 'predict_blurs', values=[features]):
        flow = tf.layers.conv2d(features, 2, 1, padding='same')

        if learn_alpha:
          alpha = tf.layers.conv2d(
              features, blur_steps, 1, padding='same', activation=tf.nn.softmax)

      with tf.name_scope(None, 'apply_blurs', values=[]):
        with tf.name_scope(None, 'warp', values=[frame, flow]):
          warps = identity_warp + flow[:, :, :, tf.newaxis, :] * warp_steps
          warps = tf.clip_by_value(warps, 0.0, max_warps)
          warped = contrib_resampler.resampler(frame, warps)
          warped = tf.concat([frame[:, :, :, tf.newaxis, :], warped], axis=3)

        with tf.name_scope(None, 'apply_alpha', values=[frame, flow]):
          if learn_alpha:
            mask = alpha[:, :, :, :, tf.newaxis]
          else:
            mask = 1.0 / blur_steps
          output_frames.append(tf.reduce_sum(warped * mask, axis=3))

    with tf.name_scope(None, 'outputs', values=[output_frames]):
      output = tf.add_n(output_frames) / len(frames)
      return output
//  [python]
def apply_line_prediction(inputs,
                          features,
                          blur_steps,
                          learn_alpha=True,
                          name=None):
  """Applies "Line Prediction" layer to input images."""
  inputs.shape.assert_is_compatible_with([None, None, None, 6])

  with tf.name_scope(name, 'blur_prediction', values=[inputs, features]):

    with tf.name_scope(None, 'input_frames', values=[inputs]):
      frames = [inputs[:, :, :, :3], inputs[:, :, :, 3:]]

    with tf.name_scope(None, 'frame_size', values=[inputs, features]):
      shape = tf.shape(inputs)
      height = shape[1]
      width = shape[2]

    with tf.name_scope(None, 'identity_warp', values=[]):
      x_idx, y_idx = tf.meshgrid(tf.range(width), tf.range(height))
      identity_warp = tf.cast(tf.stack([x_idx, y_idx], axis=-1), tf.float32)
      identity_warp = identity_warp[tf.newaxis, :, :, tf.newaxis, :]

      warp_steps = tf.cast(tf.range(blur_steps - 1) + 1, tf.float32) / (blur_steps - 1)
      warp_steps = warp_steps[tf.newaxis, tf.newaxis, tf.newaxis, :, tf.newaxis]

      max_warps = tf.cast(tf.stack([width - 1, height - 1]), tf.float32)
      max_warps = max_warps[tf.newaxis, tf.newaxis, tf.newaxis, tf.newaxis, :]

    output_frames = []
    for frame in frames:
      with tf.name_scope(None, 'predict_blurs', values=[features]):
        flow = tf.layers.conv2d(features, 2, 1, padding='same')

        if learn_alpha:
          alpha = tf.layers.conv2d(
              features, blur_steps, 1, padding='same', activation=tf.nn.softmax)

      with tf.name_scope(None, 'apply_blurs', values=[]):
        with tf.name_scope(None, 'warp', values=[frame, flow]):
          warps = identity_warp + flow[:, :, :, tf.newaxis, :] * warp_steps
          warps = tf.clip_by_value(warps, 0.0, max_warps)
          warped = contrib_resampler.resampler(frame, warps)
          warped = tf.concat([frame[:, :, :, tf.newaxis, :], warped], axis=3)

        with tf.name_scope(None, 'apply_alpha', values=[frame, flow]):
          if learn_alpha:
            mask = alpha[:, :, :, :, tf.newaxis]
          else:
            mask = 1.0 / blur_steps
          output_frames.append(tf.reduce_sum(warped * mask, axis=3))

    with tf.name_scope(None, 'outputs', values=[output_frames]):
      output = tf.add_n(output_frames) / len(frames)
      return output
//  [python]
def unstack(self, b, size, bottleneck_bits, name):
    with tf.variable_scope(name + "_unstack"):
      unb = self.unbottleneck(b, size)
      dec = self.decoder(unb)
      pred = tf.layers.dense(dec, bottleneck_bits, name="pred")
      pred_shape = common_layers.shape_list(pred)
      pred1 = tf.reshape(pred, pred_shape[:-1] + [-1, 2])
      x, y = tf.split(pred1, 2, axis=-1)
      x = tf.squeeze(x, axis=[-1])
      y = tf.squeeze(y, axis=[-1])
      gt = 2.0 * tf.to_float(tf.less(x, y)) - 1.0
      gtc = tf.tanh(y - x)
      gt += gtc - tf.stop_gradient(gtc)
      return gt, pred1
//  [python]
def unstack(self, b, size, bottleneck_bits, name):
    with tf.variable_scope(name + "_unstack"):
      unb = self.unbottleneck(b, size)
      dec = self.decoder(unb)
      pred = tf.layers.dense(dec, bottleneck_bits, name="pred")
      pred_shape = common_layers.shape_list(pred)
      pred1 = tf.reshape(pred, pred_shape[:-1] + [-1, 2])
      x, y = tf.split(pred1, 2, axis=-1)
      x = tf.squeeze(x, axis=[-1])
      y = tf.squeeze(y, axis=[-1])
      gt = 2.0 * tf.cast(tf.less(x, y), tf.float32) - 1.0
      gtc = tf.tanh(y - x)
      gt += gtc - tf.stop_gradient(gtc)
      return gt, pred1
//  [python]
def distance_prnuing_wo_mention_score(self, top_span_emb, c):
    k = util.shape(top_span_emb, 0)
    top_antecedent_offsets = tf.tile(tf.expand_dims(tf.range(c) + 1, 0), [k, 1]) # [k, c]
    raw_top_antecedents = tf.expand_dims(tf.range(k), 1) - top_antecedent_offsets # [k, c]
    top_antecedents_mask = raw_top_antecedents >= 0 # [k, c]
    top_antecedents = tf.maximum(raw_top_antecedents, 0) # [k, c]
    
    top_fast_antecedent_scores = tf.log(tf.to_float(top_antecedents_mask)) # [k, c]
    return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def distance_prnuing_wo_mention_score(self, top_span_emb, c):
    k = util.shape(top_span_emb, 0)
    top_antecedent_offsets = tf.tile(tf.expand_dims(tf.range(c) + 1, 0), [k, 1]) # [k, c]
    raw_top_antecedents = tf.expand_dims(tf.range(k), 1) - top_antecedent_offsets # [k, c]
    top_antecedents_mask = raw_top_antecedents >= 0 # [k, c]
    top_antecedents = tf.maximum(raw_top_antecedents, 0) # [k, c]
    
    top_fast_antecedent_scores = tf.log(tf.cast(top_antecedents_mask, tf.float32)) # [k, c]
    return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def bottleneck(self, x):  # pylint: disable=arguments-differ
    hparams = self.hparams
    if hparams.unordered:
      return super(AutoencoderOrderedDiscrete, self).bottleneck(x)
    noise = hparams.bottleneck_noise
    hparams.bottleneck_noise = 0.0  # We'll add noise below.
    x, loss = discretization.parametrized_bottleneck(x, hparams)
    hparams.bottleneck_noise = noise
    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
      # We want a number p such that p^bottleneck_bits = 1 - noise.
      # So log(p) * bottleneck_bits = log(noise)
      log_p = tf.log1p(-float(noise) / 2) / float(hparams.bottleneck_bits)
      # Probabilities of flipping are p, p^2, p^3, ..., p^bottleneck_bits.
      noise_mask = 1.0 - tf.exp(tf.cumsum(tf.zeros_like(x) + log_p, axis=-1))
      # Having the no-noise mask, we can make noise just uniformly at random.
      ordered_noise = tf.random_uniform(tf.shape(x))
      # We want our noise to be 1s at the start and random {-1, 1} bits later.
      ordered_noise = tf.to_float(tf.less(noise_mask, ordered_noise))
      # Now we flip the bits of x on the noisy positions (ordered and normal).
      x *= 2.0 * ordered_noise - 1
    return x, loss
//  [python]
def bottleneck(self, x):  # pylint: disable=arguments-differ
    hparams = self.hparams
    if hparams.unordered:
      return super(AutoencoderOrderedDiscrete, self).bottleneck(x)
    noise = hparams.bottleneck_noise
    hparams.bottleneck_noise = 0.0  # We'll add noise below.
    x, loss = discretization.parametrized_bottleneck(x, hparams)
    hparams.bottleneck_noise = noise
    if hparams.mode == tf.estimator.ModeKeys.TRAIN:
      # We want a number p such that p^bottleneck_bits = 1 - noise.
      # So log(p) * bottleneck_bits = log(noise)
      log_p = tf.log1p(-float(noise) / 2) / float(hparams.bottleneck_bits)
      # Probabilities of flipping are p, p^2, p^3, ..., p^bottleneck_bits.
      noise_mask = 1.0 - tf.exp(tf.cumsum(tf.zeros_like(x) + log_p, axis=-1))
      # Having the no-noise mask, we can make noise just uniformly at random.
      ordered_noise = tf.random_uniform(tf.shape(x))
      # We want our noise to be 1s at the start and random {-1, 1} bits later.
      ordered_noise = tf.cast(tf.less(noise_mask, ordered_noise), tf.float32)
      # Now we flip the bits of x on the noisy positions (ordered and normal).
      x *= 2.0 * ordered_noise - 1
    return x, loss