jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def ComputeChainStats(chain, target_mean, num_leapfrog_steps):
  # Chain is [num_steps, batch, num_dims]
  num_steps = tf.shape(chain)[0]
  counts = tf.to_float(tf.range(1, num_steps + 1))
  chain_mean = tf.cumsum(chain, 0) / counts[:, tf.newaxis, tf.newaxis]

  bias = target_mean - tf.reduce_mean(chain_mean, 1)
  variance = tf.reduce_mean(
      tf.square(chain_mean - tf.reduce_mean(chain_mean, 1, keep_dims=True)), 1)
  inst_bias = target_mean - tf.reduce_mean(chain, 1)
  inst_variance = tf.reduce_mean(tf.square(target_mean - chain), 1)

  def reducer(_, idx):
    chain_mean = tf.reduce_mean(chain[idx // 2:idx], 0)
    bias = tf.reduce_mean(target_mean - chain_mean, 0)
    variance = tf.reduce_mean(
        tf.square(chain_mean - tf.reduce_mean(chain_mean, 0)), 0)
    return bias, variance

  indices = 1 + tf.range(num_steps)
  warmupped_bias, warmupped_variance = tf.scan(
      reducer, indices, initializer=(chain[0, 0], chain[0, 0]))

  half_steps = num_steps // 2
  half_chain = chain[half_steps:]

  error_sq = tf.reduce_mean(
      tf.square(tf.reduce_mean(half_chain, 0) - target_mean), 0)

  ess = utils.EffectiveSampleSize(half_chain) / tf.to_float(half_steps)
  ess_per_grad = ess / tf.to_float(num_leapfrog_steps)
  rhat = tfp.mcmc.potential_scale_reduction(half_chain)
  autocorr = tf.reduce_mean(
      utils.SanitizedAutoCorrelation(half_chain, 0, max_lags=300), 1)

  return ChainStats(
      bias=bias,
      variance=variance,
      error_sq=error_sq,
      inst_bias=inst_bias,
      inst_variance=inst_variance,
      ess=ess,
      ess_per_grad=ess_per_grad,
      rhat=rhat,
      warmupped_bias=warmupped_bias,
      warmupped_variance=warmupped_variance,
      autocorr=autocorr)
//  [python]
def ComputeChainStats(chain, target_mean, num_leapfrog_steps):
  # Chain is [num_steps, batch, num_dims]
  num_steps = tf.shape(chain)[0]
  counts = tf.cast(tf.range(1, num_steps + 1), tf.float32)
  chain_mean = tf.cumsum(chain, 0) / counts[:, tf.newaxis, tf.newaxis]

  bias = target_mean - tf.reduce_mean(chain_mean, 1)
  variance = tf.reduce_mean(
      tf.square(chain_mean - tf.reduce_mean(chain_mean, 1, keep_dims=True)), 1)
  inst_bias = target_mean - tf.reduce_mean(chain, 1)
  inst_variance = tf.reduce_mean(tf.square(target_mean - chain), 1)

  def reducer(_, idx):
    chain_mean = tf.reduce_mean(chain[idx // 2:idx], 0)
    bias = tf.reduce_mean(target_mean - chain_mean, 0)
    variance = tf.reduce_mean(
        tf.square(chain_mean - tf.reduce_mean(chain_mean, 0)), 0)
    return bias, variance

  indices = 1 + tf.range(num_steps)
  warmupped_bias, warmupped_variance = tf.scan(
      reducer, indices, initializer=(chain[0, 0], chain[0, 0]))

  half_steps = num_steps // 2
  half_chain = chain[half_steps:]

  error_sq = tf.reduce_mean(
      tf.square(tf.reduce_mean(half_chain, 0) - target_mean), 0)

  ess = utils.EffectiveSampleSize(half_chain) / tf.cast(half_steps, tf.float32)
  ess_per_grad = ess / tf.cast(num_leapfrog_steps, tf.float32)
  rhat = tfp.mcmc.potential_scale_reduction(half_chain)
  autocorr = tf.reduce_mean(
      utils.SanitizedAutoCorrelation(half_chain, 0, max_lags=300), 1)

  return ChainStats(
      bias=bias,
      variance=variance,
      error_sq=error_sq,
      inst_bias=inst_bias,
      inst_variance=inst_variance,
      ess=ess,
      ess_per_grad=ess_per_grad,
      rhat=rhat,
      warmupped_bias=warmupped_bias,
      warmupped_variance=warmupped_variance,
      autocorr=autocorr)
//  [python]
def add_noise_and_normalize(self, x):
        x = tf.to_float(x) + tf.random_uniform(shape=tf.shape(x), minval=0., maxval=1.)
        x = (x - self.ob_mean) / self.ob_std
        return x
//  [python]
def add_noise_and_normalize(self, x):
        x = tf.cast(x, tf.float32) + tf.random_uniform(shape=tf.shape(x), minval=0., maxval=1.)
        x = (x - self.ob_mean) / self.ob_std
        return x
//  [python]
def softmax_loss(self, antecedent_scores, antecedent_labels):
        gold_scores = antecedent_scores + tf.log(tf.to_float(antecedent_labels)) # [k, max_ant + 1]
        marginalized_gold_scores = tf.reduce_logsumexp(gold_scores, [1]) # [k]
        log_norm = tf.reduce_logsumexp(antecedent_scores, [1]) # [k]
        return log_norm - marginalized_gold_scores
//  [python]
def softmax_loss(self, antecedent_scores, antecedent_labels):
        gold_scores = antecedent_scores + tf.log(tf.cast(antecedent_labels, dtype=tf.float32)) # [k, max_ant + 1]
        marginalized_gold_scores = tf.reduce_logsumexp(gold_scores, [1]) # [k]
        log_norm = tf.reduce_logsumexp(antecedent_scores, [1]) # [k]
        return log_norm - marginalized_gold_scores
//  [python]
def get_features(self, x, reuse):
        nl = tf.nn.leaky_relu
        x_has_timesteps = (x.get_shape().ndims == 5)
        if x_has_timesteps:
            sh = tf.shape(x)
            x = flatten_two_dims(x)
        with tf.variable_scope(self.scope + "_features", reuse=reuse):
            x = (tf.to_float(x) - self.ob_mean) / self.ob_std
            x = small_convnet(x, nl=nl, feat_dim=self.feat_dim, last_nl=None, layernormalize=self.layernormalize)
        if x_has_timesteps:
            x = unflatten_first_dim(x, sh)
        return x
//  [python]
def get_features(self, x, reuse):
        nl = tf.nn.leaky_relu
        x_has_timesteps = (x.get_shape().ndims == 5)
        if x_has_timesteps:
            sh = tf.shape(x)
            x = flatten_two_dims(x)
        with tf.variable_scope(self.scope + "_features", reuse=reuse):
            x = (tf.cast(x, tf.float32) - self.ob_mean) / self.ob_std
            x = small_convnet(x, nl=nl, feat_dim=self.feat_dim, last_nl=None, layernormalize=self.layernormalize)
        if x_has_timesteps:
            x = unflatten_first_dim(x, sh)
        return x
//  [python]
def bucket_distance(self, distances):
        """
        Places the given values (designed for distances) into 10 semi-logscale buckets:
        [0, 1, 2, 3, 4, 5-7, 8-15, 16-31, 32-63, 64+].
        """
        logspace_idx = tf.to_int32(tf.floor(tf.log(tf.to_float(distances))/math.log(2))) + 3
        use_identity = tf.to_int32(distances <= 4)
        combined_idx = use_identity * distances + (1 - use_identity) * logspace_idx
        return tf.clip_by_value(combined_idx, 0, 9)
//  [python]
def bucket_distance(self, distances):
        """
        Places the given values (designed for distances) into 10 semi-logscale buckets:
        [0, 1, 2, 3, 4, 5-7, 8-15, 16-31, 32-63, 64+].
        """
        logspace_idx = tf.cast(tf.floor(tf.math.log(tf.cast(distances, tf.float32))/math.log(2)), tf.int32) + 3
        use_identity = tf.cast(distances <= 4, tf.int32)
        combined_idx = use_identity * distances + (1 - use_identity) * logspace_idx
        return tf.clip_by_value(combined_idx, 0, 9)
//  [python]
def get_features(self, x, reuse):
        with tf.variable_scope(self.scope + "_features", reuse=reuse):
            x = (tf.to_float(x) - self.ob_mean) / self.ob_std
        return x
//  [python]
def get_features(self, x, reuse):
        with tf.variable_scope(self.scope + "_features", reuse=reuse):
            x = (tf.cast(x, tf.float32) - self.ob_mean) / self.ob_std
        return x
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, top_span_mention_scores, c):
        k = util.shape(top_span_emb, 0)
        top_span_range = tf.range(k) # [k]
        antecedent_offsets = tf.expand_dims(top_span_range, 1) - tf.expand_dims(top_span_range, 0) # [k, k]
        antecedents_mask = antecedent_offsets >= 1 # [k, k]
        fast_antecedent_scores = tf.expand_dims(top_span_mention_scores, 1) + tf.expand_dims(top_span_mention_scores, 0) # [k, k]
        fast_antecedent_scores += tf.log(tf.to_float(antecedents_mask)) # [k, k]
        fast_antecedent_scores += self.get_fast_antecedent_scores(top_span_emb) # [k, k]

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

        _, top_antecedents = tf.nn.top_k(fast_antecedent_scores, c, sorted=False) # [k, c]
        top_antecedents_mask = util.batch_gather(antecedents_mask, top_antecedents) # [k, c]
        top_fast_antecedent_scores = util.batch_gather(fast_antecedent_scores, top_antecedents) # [k, c]
        top_antecedent_offsets = util.batch_gather(antecedent_offsets, top_antecedents) # [k, c]
        return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def model_fn(features, labels, mode, params):
  """A model is called by TpuEstimator."""
  del labels
  del features

  mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
  layout_rules = mtf.convert_to_layout_rules(FLAGS.layout)

  ctx = params['context']
  num_hosts = ctx.num_hosts
  host_placement_fn = ctx.tpu_host_placement_function
  device_list = [host_placement_fn(host_id=t) for t in range(num_hosts)]
  tf.logging.info('device_list = %s' % device_list,)

  mesh_devices = [''] * mesh_shape.size
  mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(mesh_shape, layout_rules,
                                              mesh_devices,
                                              ctx.device_assignment)

  graph = mtf.Graph()
  mesh = mtf.Mesh(graph, "fft_mesh")

  with mtf.utils.outside_all_rewrites():
    fsum = benchmark_model(mesh)
  lowering = mtf.Lowering(graph, {mesh: mesh_impl})

  tf_err = tf.to_float(lowering.export_to_tf_tensor(fsum))

  with mtf.utils.outside_all_rewrites():
    return tpu_estimator.TPUEstimatorSpec(mode, loss=tf_err)
//  [python]
def model_fn(features, labels, mode, params):
  """A model is called by TpuEstimator."""
  del labels
  del features

  mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
  layout_rules = mtf.convert_to_layout_rules(FLAGS.layout)

  ctx = params['context']
  num_hosts = ctx.num_hosts
  host_placement_fn = ctx.tpu_host_placement_function
  device_list = [host_placement_fn(host_id=t) for t in range(num_hosts)]
  tf.logging.info('device_list = %s' % device_list,)

  mesh_devices = [''] * mesh_shape.size
  mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(mesh_shape, layout_rules,
                                              mesh_devices,
                                              ctx.device_assignment)

  graph = mtf.Graph()
  mesh = mtf.Mesh(graph, "fft_mesh")

  with mtf.utils.outside_all_rewrites():
    fsum = benchmark_model(mesh)
  lowering = mtf.Lowering(graph, {mesh: mesh_impl})

  tf_err = tf.cast(lowering.export_to_tf_tensor(fsum), tf.float32)

  with mtf.utils.outside_all_rewrites():
    return tpu_estimator.TPUEstimatorSpec(mode, loss=tf_err)
//  [python]
def _decode_png_instance_masks(self, keys_to_tensors: TensorDict)-> tf.Tensor:
    """Decode PNG instance segmentation masks and stack into dense tensor.

    The instance segmentation masks are reshaped to [num_instances, height,
    width].

    Args:
      keys_to_tensors: A dictionary from keys to tensors.

    Returns:
      A 3-D float tensor of shape [num_instances, height, width] with values
      in {0, 1}.
    """

    def decode_png_mask(image_buffer):
      image = tf.squeeze(
          tf.image.decode_image(image_buffer, channels=1), axis=2)
      image.set_shape([None, None])
      image = tf.to_float(tf.greater(image, 0))
      return image

    png_masks = keys_to_tensors['image/object/mask']
    height = keys_to_tensors['image/height']
    width = keys_to_tensors['image/width']
    if isinstance(png_masks, tf.SparseTensor):
      png_masks = tf.sparse_tensor_to_dense(png_masks, default_value='')
    return tf.cond(
        tf.greater(tf.size(png_masks), 0),
        lambda: tf.map_fn(decode_png_mask, png_masks, dtype=tf.float32),
        lambda: tf.zeros(tf.to_int32(tf.stack([0, height, width]))))
//  [python]
def _decode_png_instance_masks(self, keys_to_tensors: TensorDict)-> tf.Tensor:
    """Decode PNG instance segmentation masks and stack into dense tensor.

    The instance segmentation masks are reshaped to [num_instances, height,
    width].

    Args:
      keys_to_tensors: A dictionary from keys to tensors.

    Returns:
      A 3-D float tensor of shape [num_instances, height, width] with values
      in {0, 1}.
    """

    def decode_png_mask(image_buffer):
      image = tf.squeeze(
          tf.image.decode_image(image_buffer, channels=1), axis=2)
      image.set_shape([None, None])
      image = tf.cast(tf.greater(image, 0), tf.float32)
      return image

    png_masks = keys_to_tensors['image/object/mask']
    height = keys_to_tensors['image/height']
    width = keys_to_tensors['image/width']
    if isinstance(png_masks, tf.SparseTensor):
      png_masks = tf.sparse_tensor_to_dense(png_masks, default_value='')
    return tf.cond(
        tf.greater(tf.size(png_masks), 0),
        lambda: tf.map_fn(decode_png_mask, png_masks, dtype=tf.float32),
        lambda: tf.zeros(tf.to_int32(tf.stack([0, height, width]))))
//  [python]
def compute_values(edge_compatibility, v):
  """Compute values. If edge compatibilities is just adjacency, we get ggnn.

  Args:
    edge_compatibility: A tensor of shape [batch, num_transforms, length, depth]
    v: A tensor of shape [batch, num_transforms, length, depth]

  Returns:
    output: A [batch, length, depth] tensor
  """

  # Computes the incoming value vectors for each node by weighting them
  # according to the attention weights. These values are still segregated by
  # edge type.
  # Shape = [B, T, N, V].
  all_edge_values = tf.matmul(tf.to_float(edge_compatibility), v)

  # Combines the weighted value vectors together across edge types into a
  # single N x V matrix for each batch.
  output = tf.reduce_sum(all_edge_values, axis=1)  # Shape [B, N, V].
  return output
//  [python]
def compute_values(edge_compatibility, v):
  """Compute values. If edge compatibilities is just adjacency, we get ggnn.

  Args:
    edge_compatibility: A tensor of shape [batch, num_transforms, length, depth]
    v: A tensor of shape [batch, num_transforms, length, depth]

  Returns:
    output: A [batch, length, depth] tensor
  """

  # Computes the incoming value vectors for each node by weighting them
  # according to the attention weights. These values are still segregated by
  # edge type.
  # Shape = [B, T, N, V].
  all_edge_values = tf.matmul(tf.cast(edge_compatibility, tf.float32), v)

  # Combines the weighted value vectors together across edge types into a
  # single N x V matrix for each batch.
  output = tf.reduce_sum(all_edge_values, axis=1)  # Shape [B, N, V].
  return output
//  [python]
def learning_rate_schedule(base_learning_rate, num_examples):
  """Build learning rate schedule."""
  global_step = tf.train.get_or_create_global_step()
  warmup_steps = int(round(
      FLAGS.warmup_epochs * num_examples // FLAGS.train_batch_size))
  if FLAGS.learning_rate_scaling == 'linear':
    scaled_lr = base_learning_rate * FLAGS.train_batch_size / 256.
  elif FLAGS.learning_rate_scaling == 'sqrt':
    scaled_lr = base_learning_rate * math.sqrt(FLAGS.train_batch_size)
  else:
    raise ValueError('Unknown learning rate scaling {}'.format(
        FLAGS.learning_rate_scaling))
  learning_rate = (tf.to_float(global_step) / int(warmup_steps) * scaled_lr
                   if warmup_steps else scaled_lr)

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

  return learning_rate
//  [python]
def learning_rate_schedule(base_learning_rate, num_examples):
  """Build learning rate schedule."""
  global_step = tf.train.get_or_create_global_step()
  warmup_steps = int(round(
      FLAGS.warmup_epochs * num_examples // FLAGS.train_batch_size))
  if FLAGS.learning_rate_scaling == 'linear':
    scaled_lr = base_learning_rate * FLAGS.train_batch_size / 256.
  elif FLAGS.learning_rate_scaling == 'sqrt':
    scaled_lr = base_learning_rate * math.sqrt(FLAGS.train_batch_size)
  else:
    raise ValueError('Unknown learning rate scaling {}'.format(
        FLAGS.learning_rate_scaling))
  learning_rate = (tf.cast(global_step, tf.float32) / int(warmup_steps) * scaled_lr
                   if warmup_steps else scaled_lr)

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

  return learning_rate
//  [python]
def xception_exit(inputs):
  """Xception exit flow."""
  with tf.variable_scope("xception_exit"):
    x = inputs
    x_shape = x.get_shape().as_list()
    if x_shape[1] is None or x_shape[2] is None:
      length_float = tf.to_float(tf.shape(x)[1])
      length_float *= tf.to_float(tf.shape(x)[2])
      spatial_dim_float = tf.sqrt(length_float)
      spatial_dim = tf.to_int32(spatial_dim_float)
      x_depth = x_shape[3]
      x = tf.reshape(x, [-1, spatial_dim, spatial_dim, x_depth])
    elif x_shape[1] != x_shape[2]:
      spatial_dim = int(math.sqrt(float(x_shape[1] * x_shape[2])))
      if spatial_dim * spatial_dim != x_shape[1] * x_shape[2]:
        raise ValueError("Assumed inputs were square-able but they were "
                         "not. Shape: %s" % x_shape)
      x = tf.reshape(x, [-1, spatial_dim, spatial_dim, x_depth])

    x = common_layers.conv_block_downsample(x, (3, 3), (2, 2), "SAME")
    return tf.nn.relu(x)
//  [python]
def xception_exit(inputs):
  """Xception exit flow."""
  with tf.variable_scope("xception_exit"):
    x = inputs
    x_shape = x.get_shape().as_list()
    if x_shape[1] is None or x_shape[2] is None:
      length_float = tf.cast(tf.shape(x)[1], tf.float32)
      length_float *= tf.cast(tf.shape(x)[2], tf.float32)
      spatial_dim_float = tf.sqrt(length_float)
      spatial_dim = tf.cast(spatial_dim_float, tf.int32)
      x_depth = x_shape[3]
      x = tf.reshape(x, [-1, spatial_dim, spatial_dim, x_depth])
    elif x_shape[1] != x_shape[2]:
      spatial_dim = int(math.sqrt(float(x_shape[1] * x_shape[2])))
      if spatial_dim * spatial_dim != x_shape[1] * x_shape[2]:
        raise ValueError("Assumed inputs were square-able but they were "
                         "not. Shape: %s" % x_shape)
      x = tf.reshape(x, [-1, spatial_dim, spatial_dim, x_depth])

    x = common_layers.conv_block_downsample(x, (3, 3), (2, 2), "SAME")
    return tf.nn.relu(x)
//  [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, num_edge_types] tensor.
    num_edge_types: Number of different edge types
    depth: Number of channels
    name: A optional string name for scoping
  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))

    att_adj_vectors = tf.matmul(
        tf.reshape(tf.to_float(adjacency_matrix), [-1, num_edge_types]),
        adj_vectors)
    # Reshape to be [batch, num_nodes, num_nodes, depth].
    att_adj_vectors = tf.reshape(att_adj_vectors, [
        adjacency_matrix_shape[0], adjacency_matrix_shape[1],
        adjacency_matrix_shape[2], depth
    ])
    return att_adj_vectors
//  [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, num_edge_types] tensor.
    num_edge_types: Number of different edge types
    depth: Number of channels
    name: A optional string name for scoping
  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))

    att_adj_vectors = tf.matmul(
        tf.reshape(tf.cast(adjacency_matrix, dtype=tf.float32), [-1, num_edge_types]),
        adj_vectors)
    # Reshape to be [batch, num_nodes, num_nodes, depth].
    att_adj_vectors = tf.reshape(att_adj_vectors, [
        adjacency_matrix_shape[0], adjacency_matrix_shape[1],
        adjacency_matrix_shape[2], depth
    ])
    return att_adj_vectors
//  [python]
def gaussian_blobs(params, height, width, norm='sum'):
  """Creates gaussian blobs on a canvas.

  Args:

    params: [B, 4] tensor, where entries represent (y, x, y_scale, x_scale).
    height: int, height of the output.
    width: int, width of the output.
    norm: type of normalization to use; must be a postfix of some
      tf.reduce_.* method.

  Returns:
    Tensor of shape [B, height, width].
  """

  params = tf.expand_dims(params, -1)
  uy, ux, sy, sx = tf.split(params, 4, -2)

  rows = tf.range(tf.to_int32(height))
  rows = tf.to_float(rows)[tf.newaxis, :, tf.newaxis]

  cols = tf.range(tf.to_int32(width))
  cols = tf.to_float(cols)[tf.newaxis, tf.newaxis, :]

  dy = (rows - uy) / sy
  dx = (cols - ux) / sx

  z = tf.square(dy) + tf.square(dx)
  mask = tf.exp(-.5 * z)

  # normalize so that the contribution of each blob sums to one
  # change this to `tf.reduce_max` if you want max value to be one
  norm_func = getattr(tf, 'reduce_{}'.format(norm))
  mask /= norm_func(mask, (1, 2), keep_dims=True) + 1e-8  # pylint:disable=not-callable
  return mask
//  [python]
def gaussian_blobs(params, height, width, norm='sum'):
  """Creates gaussian blobs on a canvas.

  Args:

    params: [B, 4] tensor, where entries represent (y, x, y_scale, x_scale).
    height: int, height of the output.
    width: int, width of the output.
    norm: type of normalization to use; must be a postfix of some
      tf.reduce_.* method.

  Returns:
    Tensor of shape [B, height, width].
  """

  params = tf.expand_dims(params, -1)
  uy, ux, sy, sx = tf.split(params, 4, -2)

  rows = tf.range(tf.cast(height, tf.int32))
  rows = tf.to_float(rows)[tf.newaxis, :, tf.newaxis]

  cols = tf.range(tf.cast(width, tf.int32))
  cols = tf.to_float(cols)[tf.newaxis, tf.newaxis, :]

  dy = (rows - uy) / sy
  dx = (cols - ux) / sx

  z = tf.square(dy) + tf.square(dx)
  mask = tf.exp(-.5 * z)

  # normalize so that the contribution of each blob sums to one
  # change this to `tf.reduce_max` if you want max value to be one
  norm_func = getattr(tf, 'reduce_{}'.format(norm))
  mask /= norm_func(mask, (1, 2), keep_dims=True) + 1e-8  # pylint:disable=not-callable
  return mask
//  [python]
def __init__(self, config, tasks, task_weights, is_training,
               features, num_train_steps):
    # Create a shared BERT encoder
    bert_config = modeling.BertConfig.from_json_file(config.bert_config_file)
    if config.debug:
      bert_config.num_hidden_layers = 3
      bert_config.hidden_size = 144
    assert config.max_seq_length <= bert_config.max_position_embeddings
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      bert_model = modeling.BertModel(
          config=bert_config,
          is_training=is_training,
          input_ids=features["input_ids"],
          input_mask=features["input_mask"],
          token_type_ids=features["segment_ids"],
          use_one_hot_embeddings=config.use_tpu)
    percent_done = (tf.to_float(tf.train.get_or_create_global_step()) /
                    tf.to_float(num_train_steps))

    # Add specific tasks
    self.outputs = {"task_id": features["task_id"]}
    losses = []
    for task in tasks:
      with tf.variable_scope("task_specific/" + task.name):
        task_losses, task_outputs = task.get_prediction_module(
            bert_model, features, is_training, percent_done)
        losses.append(task_losses * task_weights[task.name])
        self.outputs[task.name] = task_outputs

    # For TPU-friendlyness inference for all tasks is performed on each example.
    # However, the losses from all tasks but the example's task are masked out.
    # This doesn't significantly slow down runtime as long as the task-specific
    # modules are much faster to run than the BERT encoder.
    self.loss = tf.reduce_sum(
        tf.stack(losses, -1) *
        tf.one_hot(features["task_id"], config.n_tasks))
//  [python]
def __init__(self, config, tasks, task_weights, is_training,
               features, num_train_steps):
    # Create a shared BERT encoder
    bert_config = modeling.BertConfig.from_json_file(config.bert_config_file)
    if config.debug:
      bert_config.num_hidden_layers = 3
      bert_config.hidden_size = 144
    assert config.max_seq_length <= bert_config.max_position_embeddings
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      bert_model = modeling.BertModel(
          config=bert_config,
          is_training=is_training,
          input_ids=features["input_ids"],
          input_mask=features["input_mask"],
          token_type_ids=features["segment_ids"],
          use_one_hot_embeddings=config.use_tpu)
    percent_done = (tf.cast(tf.train.get_or_create_global_step(), tf.float32) /
                    tf.cast(num_train_steps, tf.float32))

    # Add specific tasks
    self.outputs = {"task_id": features["task_id"]}
    losses = []
    for task in tasks:
      with tf.variable_scope("task_specific/" + task.name):
        task_losses, task_outputs = task.get_prediction_module(
            bert_model, features, is_training, percent_done)
        losses.append(task_losses * task_weights[task.name])
        self.outputs[task.name] = task_outputs

    # For TPU-friendlyness inference for all tasks is performed on each example.
    # However, the losses from all tasks but the example's task are masked out.
    # This doesn't significantly slow down runtime as long as the task-specific
    # modules are much faster to run than the BERT encoder.
    self.loss = tf.reduce_sum(
        tf.stack(losses, -1) *
        tf.one_hot(features["task_id"], config.n_tasks))
//  [python]
def verb_refs_to_lengths(task, verb_refs, include_eos=True):
  """Computes the length of a sequence."""
  eos_positions = tf.to_int32(tf.expand_dims(
      tf.where(tf.equal(task, 1))[:, 1], 1))
  seq_mask = tf.logical_not(tf.cast(tf.cumsum(tf.to_int32(
      tf.logical_and(
          tf.equal(verb_refs[:, :, 0], eos_positions),
          tf.equal(verb_refs[:, :, 1], eos_positions + 1))), axis=-1), tf.bool))
  lengths = tf.reduce_sum(tf.to_float(seq_mask), axis=-1)
  if include_eos:
    lengths = lengths + 1
  return lengths
//  [python]
def verb_refs_to_lengths(task, verb_refs, include_eos=True):
  """Computes the length of a sequence."""
  eos_positions = tf.cast(tf.expand_dims(
      tf.where(tf.equal(task, 1))[:, 1], 1), tf.int32)
  seq_mask = tf.logical_not(tf.cast(tf.cumsum(tf.cast(
      tf.logical_and(
          tf.equal(verb_refs[:, :, 0], eos_positions),
          tf.equal(verb_refs[:, :, 1], eos_positions + 1))), tf.int32), tf.bool))
  lengths = tf.reduce_sum(tf.cast(seq_mask, tf.float32), axis=-1)
  if include_eos:
    lengths = lengths + 1
  return lengths
//  [python]
def preprocess_for_eval(image,
                        output_height,
                        output_width,
                        add_image_summaries=True,
                        use_grayscale=False):
  """Preprocesses the given image for evaluation.
  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    add_image_summaries: Enable image summaries.
    use_grayscale: Whether to convert the image from RGB to grayscale.
  Returns:
    A preprocessed image.
  """
  if add_image_summaries:
    tf.summary.image('image', tf.expand_dims(image, 0))
  # Transform the image to floats.
  image = tf.to_float(image)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)

  # Resize and crop if needed.
  resized_image = tf.image.resize_image_with_crop_or_pad(image,
                                                         output_width,
                                                         output_height)
  if add_image_summaries:
    tf.summary.image('resized_image', tf.expand_dims(resized_image, 0))

  # Subtract off the mean and divide by the variance of the pixels.
  return tf.image.per_image_standardization(resized_image)
//  [python]
def preprocess_for_eval(image,
                        output_height,
                        output_width,
                        add_image_summaries=True,
                        use_grayscale=False):
  """Preprocesses the given image for evaluation.
  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    add_image_summaries: Enable image summaries.
    use_grayscale: Whether to convert the image from RGB to grayscale.
  Returns:
    A preprocessed image.
  """
  if add_image_summaries:
    tf.summary.image('image', tf.expand_dims(image, 0))
  # Transform the image to floats.
  image = tf.cast(image, tf.float32)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)

  # Resize and crop if needed.
  resized_image = tf.image.resize_image_with_crop_or_pad(image,
                                                         output_width,
                                                         output_height)
  if add_image_summaries:
    tf.summary.image('resized_image', tf.expand_dims(resized_image, 0))

  # Subtract off the mean and divide by the variance of the pixels.
  return tf.image.per_image_standardization(resized_image)
//  [python]
def preprocess_for_train(image,
                         output_height,
                         output_width,
                         padding=_PADDING,
                         add_image_summaries=True,
                         use_grayscale=False):
  """Preprocesses the given image for training.
  Note that the actual resizing scale is sampled from
    [`resize_size_min`, `resize_size_max`].
  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    padding: The amound of padding before and after each dimension of the image.
    add_image_summaries: Enable image summaries.
    use_grayscale: Whether to convert the image from RGB to grayscale.
  Returns:
    A preprocessed image.
  """
  if add_image_summaries:
    tf.summary.image('image', tf.expand_dims(image, 0))

  # Transform the image to floats.
  image = tf.to_float(image)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)
  if padding > 0:
    image = tf.pad(image, [[padding, padding], [padding, padding], [0, 0]])
  # Randomly crop a [height, width] section of the image.
  distorted_image = tf.random_crop(image,
                                   [output_height, output_width, 3])

  # Randomly flip the image horizontally.
  distorted_image = tf.image.random_flip_left_right(distorted_image)

  if add_image_summaries:
    tf.summary.image('distorted_image', tf.expand_dims(distorted_image, 0))

  # Because these operations are not commutative, consider randomizing
  # the order their operation.
  distorted_image = tf.image.random_brightness(distorted_image,
                                               max_delta=63)
  distorted_image = tf.image.random_contrast(distorted_image,
                                             lower=0.2, upper=1.8)
  # Subtract off the mean and divide by the variance of the pixels.
  return tf.image.per_image_standardization(distorted_image)
//  [python]
def preprocess_for_train(image,
                         output_height,
                         output_width,
                         padding=_PADDING,
                         add_image_summaries=True,
                         use_grayscale=False):
  """Preprocesses the given image for training.
  Note that the actual resizing scale is sampled from
    [`resize_size_min`, `resize_size_max`].
  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    padding: The amound of padding before and after each dimension of the image.
    add_image_summaries: Enable image summaries.
    use_grayscale: Whether to convert the image from RGB to grayscale.
  Returns:
    A preprocessed image.
  """
  if add_image_summaries:
    tf.summary.image('image', tf.expand_dims(image, 0))

  # Transform the image to floats.
  image = tf.cast(image, tf.float32)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)
  if padding > 0:
    image = tf.pad(image, [[padding, padding], [padding, padding], [0, 0]])
  # Randomly crop a [height, width] section of the image.
  distorted_image = tf.random_crop(image,
                                   [output_height, output_width, 3])

  # Randomly flip the image horizontally.
  distorted_image = tf.image.random_flip_left_right(distorted_image)

  if add_image_summaries:
    tf.summary.image('distorted_image', tf.expand_dims(distorted_image, 0))

  # Because these operations are not commutative, consider randomizing
  # the order their operation.
  distorted_image = tf.image.random_brightness(distorted_image,
                                               max_delta=63)
  distorted_image = tf.image.random_contrast(distorted_image,
                                             lower=0.2, upper=1.8)
  # Subtract off the mean and divide by the variance of the pixels.
  return tf.image.per_image_standardization(distorted_image)
//  [python]
def get_timing_signal_1d_given_position(channels, position, min_timescale=1.0, max_timescale=1.0e4):
    """Get sinusoids of diff frequencies, with timing position given.

    Adapted from add_timing_signal_1d_given_position in
    //third_party/py/tensor2tensor/layers/common_attention.py

    Args:
        channels: scalar, size of timing embeddings to create. The number of
            different timescales is equal to channels / 2.
        position: a Tensor with shape [batch, seq_len]
        min_timescale: a float
        max_timescale: a float

    Returns:
        a Tensor of timing signals [batch, seq_len, channels]
    """
    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)]])
    return signal
//  [python]
def get_timing_signal_1d_given_position(channels, position, min_timescale=1.0, max_timescale=1.0e4):
    """Get sinusoids of diff frequencies, with timing position given.

    Adapted from add_timing_signal_1d_given_position in
    //third_party/py/tensor2tensor/layers/common_attention.py

    Args:
        channels: scalar, size of timing embeddings to create. The number of
            different timescales is equal to channels / 2.
        position: a Tensor with shape [batch, seq_len]
        min_timescale: a float
        max_timescale: a float

    Returns:
        a Tensor of timing signals [batch, seq_len, channels]
    """
    num_timescales = channels // 2
    log_timescale_increment = (
            math.log(float(max_timescale) / float(min_timescale)) /
            (tf.cast(num_timescales, tf.float32) - 1))
    inv_timescales = min_timescale * tf.exp(
        tf.cast(tf.range(num_timescales), tf.float32) * -log_timescale_increment)
    scaled_time = (
            tf.expand_dims(tf.cast(position, 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)]])
    return signal
//  [python]
def apply_deltas(self,regression,img_size=None,fm_size=None):
        if len(regression.get_shape()) == 2:
            B = 1
            H = fm_size[0]
            W = fm_size[1]

        elif len(regression.get_shape()) == 4:
            B,H,W,_ = wmlt.combined_static_and_dynamic_shape(regression)
        else:
            raise NotImplementedError("Error")
        x_i,y_i = tf.meshgrid(tf.range(W),tf.range(H))
        if isinstance(img_size,tf.Tensor) and img_size.dtype != tf.float32:
            img_size = tf.to_float(img_size)
        H = tf.to_float(H)
        W = tf.to_float(W)
        y_f = tf.to_float(y_i)+0.5
        x_f = tf.to_float(x_i)+0.5
        y_delta = img_size[0]/H
        x_delta = img_size[1]/W
        y_base_value = y_f*y_delta
        x_base_value = x_f*x_delta
        base_value = tf.stack([y_base_value,x_base_value,y_base_value,x_base_value],axis=-1)
        if len(regression.get_shape())==4:
            base_value = tf.expand_dims(base_value,axis=0)
            base_value = tf.stop_gradient(tf.tile(base_value,[B,1,1,1]))
            multi = tf.convert_to_tensor([[[[-1,-1,1,1]]]],dtype=tf.float32)
        elif len(regression.get_shape()) == 2:
            base_value = tf.reshape(base_value,[-1,4])
            multi = tf.convert_to_tensor([[-1,-1,1,1]],dtype=tf.float32)


        return base_value+regression*multi
//  [python]
def apply_deltas(self,regression,img_size=None,fm_size=None):
        if len(regression.get_shape()) == 2:
            B = 1
            H = fm_size[0]
            W = fm_size[1]

        elif len(regression.get_shape()) == 4:
            B,H,W,_ = wmlt.combined_static_and_dynamic_shape(regression)
        else:
            raise NotImplementedError("Error")
        x_i,y_i = tf.meshgrid(tf.range(W),tf.range(H))
        if isinstance(img_size,tf.Tensor) and img_size.dtype != tf.float32:
            img_size = tf.cast(img_size, tf.float32)
        H = tf.cast(H, tf.float32)
        W = tf.cast(W, tf.float32)
        y_f = tf.cast(y_i, tf.float32)+0.5
        x_f = tf.cast(x_i, tf.float32)+0.5
        y_delta = img_size[0]/H
        x_delta = img_size[1]/W
        y_base_value = y_f*y_delta
        x_base_value = x_f*x_delta
        base_value = tf.stack([y_base_value,x_base_value,y_base_value,x_base_value],axis=-1)
        if len(regression.get_shape())==4:
            base_value = tf.expand_dims(base_value,axis=0)
            base_value = tf.stop_gradient(tf.tile(base_value,[B,1,1,1]))
            multi = tf.convert_to_tensor([[[[-1,-1,1,1]]]],dtype=tf.float32)
        elif len(regression.get_shape()) == 2:
            base_value = tf.reshape(base_value,[-1,4])
            multi = tf.convert_to_tensor([[-1,-1,1,1]],dtype=tf.float32)


        return base_value+regression*multi
//  [python]
def create_input_pipeline(input_queue, image_size, nrof_preprocess_threads, batch_size_placeholder):
    images_and_labels_list = []
    for _ in range(nrof_preprocess_threads):
        filenames, label, control = input_queue.dequeue()
        images = []
        for filename in tf.unstack(filenames):
            file_contents = tf.read_file(filename)
            image = tf.image.decode_image(file_contents, 3)
            image = tf.cond(get_control_flag(control[0], RANDOM_ROTATE),
                            lambda:tf.py_func(random_rotate_image, [image], tf.uint8), 
                            lambda:tf.identity(image))
            image = tf.cond(get_control_flag(control[0], RANDOM_CROP), 
                            lambda:tf.random_crop(image, image_size + (3,)), 
                            lambda:tf.image.resize_image_with_crop_or_pad(image, image_size[0], image_size[1]))
            image = tf.cond(get_control_flag(control[0], RANDOM_FLIP),
                            lambda:tf.image.random_flip_left_right(image),
                            lambda:tf.identity(image))
            image = tf.cond(get_control_flag(control[0], FIXED_STANDARDIZATION),
                            lambda:(tf.cast(image, tf.float32) - 127.5)/128.0,
                            lambda:tf.image.per_image_standardization(tf.to_float(image)))
            image = tf.cond(get_control_flag(control[0], FLIP),
                            lambda:tf.image.flip_left_right(image),
                            lambda:tf.identity(image))
            #pylint: disable=no-member
            image.set_shape(image_size + (3,))
            images.append(image)
        images_and_labels_list.append([images, label])

    image_batch, label_batch = tf.train.batch_join(
        images_and_labels_list, batch_size=batch_size_placeholder, 
        shapes=[image_size + (3,), ()], enqueue_many=True,
        capacity=4 * nrof_preprocess_threads * 100,
        allow_smaller_final_batch=True)
    
    return image_batch, label_batch
//  [python]
def create_input_pipeline(input_queue, image_size, nrof_preprocess_threads, batch_size_placeholder):
    images_and_labels_list = []
    for _ in range(nrof_preprocess_threads):
        filenames, label, control = input_queue.dequeue()
        images = []
        for filename in tf.unstack(filenames):
            file_contents = tf.read_file(filename)
            image = tf.image.decode_image(file_contents, 3)
            image = tf.cond(get_control_flag(control[0], RANDOM_ROTATE),
                            lambda:tf.py_func(random_rotate_image, [image], tf.uint8), 
                            lambda:tf.identity(image))
            image = tf.cond(get_control_flag(control[0], RANDOM_CROP), 
                            lambda:tf.random_crop(image, image_size + (3,)), 
                            lambda:tf.image.resize_image_with_crop_or_pad(image, image_size[0], image_size[1]))
            image = tf.cond(get_control_flag(control[0], RANDOM_FLIP),
                            lambda:tf.image.random_flip_left_right(image),
                            lambda:tf.identity(image))
            image = tf.cond(get_control_flag(control[0], FIXED_STANDARDIZATION),
                            lambda:(tf.cast(image, tf.float32) - 127.5)/128.0,
                            lambda:tf.image.per_image_standardization(tf.cast(image, tf.float32)))
            image = tf.cond(get_control_flag(control[0], FLIP),
                            lambda:tf.image.flip_left_right(image),
                            lambda:tf.identity(image))
            #pylint: disable=no-member
            image.set_shape(image_size + (3,))
            images.append(image)
        images_and_labels_list.append([images, label])

    image_batch, label_batch = tf.train.batch_join(
        images_and_labels_list, batch_size=batch_size_placeholder, 
        shapes=[image_size + (3,), ()], enqueue_many=True,
        capacity=4 * nrof_preprocess_threads * 100,
        allow_smaller_final_batch=True)
    
    return image_batch, label_batch
//  [python]
def weights_concatenated(labels):
  """Assign weight 1.0 to the "target" part of the concatenated labels.

  The labels look like:
    source English I love you . ID1 target French Je t'aime . ID1 source
      English the cat ID1 target French le chat ID1 source English ...

  We want to assign weight 1.0 to all words in the target text (including the
  ID1 end symbol), but not to the source text or the boilerplate.  In the
  above example, the target words that get positive weight are:
    Je t'aime . ID1 le chat ID1

  Args:
    labels: a Tensor
  Returns:
    a Tensor
  """
  eos_mask = tf.to_int32(tf.equal(labels, 1))
  sentence_num = tf.cumsum(eos_mask, axis=1, exclusive=True)
  in_target = tf.equal(tf.mod(sentence_num, 2), 1)
  # first two tokens of each sentence are boilerplate.
  sentence_num_plus_one = sentence_num + 1
  shifted = tf.pad(sentence_num_plus_one,
                   [[0, 0], [2, 0], [0, 0], [0, 0]])[:, :-2, :, :]
  nonboilerplate = tf.equal(sentence_num_plus_one, shifted)
  ret = to_float(tf.logical_and(nonboilerplate, in_target))
  return ret
//  [python]
def weights_concatenated(labels):
  """Assign weight 1.0 to the "target" part of the concatenated labels.

  The labels look like:
    source English I love you . ID1 target French Je t'aime . ID1 source
      English the cat ID1 target French le chat ID1 source English ...

  We want to assign weight 1.0 to all words in the target text (including the
  ID1 end symbol), but not to the source text or the boilerplate.  In the
  above example, the target words that get positive weight are:
    Je t'aime . ID1 le chat ID1

  Args:
    labels: a Tensor
  Returns:
    a Tensor
  """
  eos_mask = tf.cast(tf.equal(labels, 1), tf.int32)
  sentence_num = tf.cumsum(eos_mask, axis=1, exclusive=True)
  in_target = tf.equal(tf.mod(sentence_num, 2), 1)
  # first two tokens of each sentence are boilerplate.
  sentence_num_plus_one = sentence_num + 1
  shifted = tf.pad(sentence_num_plus_one,
                   [[0, 0], [2, 0], [0, 0], [0, 0]])[:, :-2, :, :]
  nonboilerplate = tf.equal(sentence_num_plus_one, shifted)
  ret = tf.cast(tf.logical_and(nonboilerplate, in_target), tf.float32)
  return ret