jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def compute_average_embedding(input_embeddings, input_lengths):
  """Computes bag-of-words embedding.

  Args:
    input_embeddings: <tf.float32>[bs, max_seq_len, emb_dim]
    input_lengths: <tf.int64>[bs, 1]

  Returns:
    bow_embedding: <tf.float32>[bs, emb_dim]
  """
  max_seq_len = tf.shape(input_embeddings)[1]
  # <tf.float32>[bs, 1, max_seq_len]
  mask = tf.sequence_mask(input_lengths, max_seq_len, dtype=tf.float32)
  # <tf.float32>[bs, 1, emb_dim]
  sum_embedding = tf.matmul(mask, input_embeddings)
  # <tf.float32>[bs, 1, emb_dim]
  avg_embedding = sum_embedding / tf.to_float(tf.expand_dims(input_lengths, 2))
  # <tf.float32>[bs, dim]
  return tf.squeeze(avg_embedding, 1)
//  [python]
def compute_average_embedding(input_embeddings, input_lengths):
  """Computes bag-of-words embedding.

  Args:
    input_embeddings: <tf.float32>[bs, max_seq_len, emb_dim]
    input_lengths: <tf.int64>[bs, 1]

  Returns:
    bow_embedding: <tf.float32>[bs, emb_dim]
  """
  max_seq_len = tf.shape(input_embeddings)[1]
  # <tf.float32>[bs, 1, max_seq_len]
  mask = tf.sequence_mask(input_lengths, max_seq_len, dtype=tf.float32)
  # <tf.float32>[bs, 1, emb_dim]
  sum_embedding = tf.matmul(mask, input_embeddings)
  # <tf.float32>[bs, 1, emb_dim]
  avg_embedding = sum_embedding / tf.cast(tf.expand_dims(input_lengths, 2), tf.float32)
  # <tf.float32>[bs, dim]
  return tf.squeeze(avg_embedding, 1)
//  [python]
def build_graph_dist_strategy(self, features, labels, mode, params):
    """Model function."""
    del labels, params
    misc_utils.print_out("Running dist_strategy mode_fn")

    hparams = self.hparams

    # Create a GNMT model for training.
    # assert (hparams.encoder_type == "gnmt" or
    #        hparams.attention_architecture in ["gnmt", "gnmt_v2"])
    with mixed_precision_scope():
      model = gnmt_model.GNMTModel(hparams, mode=mode, features=features)
      if mode == tf.contrib.learn.ModeKeys.INFER:
        sample_ids = model.sample_id
        reverse_target_vocab_table = lookup_ops.index_to_string_table_from_file(
            hparams.tgt_vocab_file, default_value=vocab_utils.UNK)
        sample_words = reverse_target_vocab_table.lookup(
            tf.to_int64(sample_ids))
        # make sure outputs is of shape [batch_size, time] or [beam_width,
        # batch_size, time] when using beam search.
        if hparams.time_major:
          sample_words = tf.transpose(sample_words)
        elif sample_words.shape.ndims == 3:
          # beam search output in [batch_size, time, beam_width] shape.
          sample_words = tf.transpose(sample_words, [2, 0, 1])
        predictions = {"predictions": sample_words}
        # return loss, vars, grads, predictions, train_op, scaffold
        return None, None, None, predictions, None, None
      elif mode == tf.contrib.learn.ModeKeys.TRAIN:
        loss = model.train_loss
        train_op = model.update
        return loss, model.params, model.grads, None, train_op, None
      else:
        raise ValueError("Unknown mode in model_fn: %s" % mode)
//  [python]
def build_graph_dist_strategy(self, features, labels, mode, params):
    """Model function."""
    del labels, params
    misc_utils.print_out("Running dist_strategy mode_fn")

    hparams = self.hparams

    # Create a GNMT model for training.
    # assert (hparams.encoder_type == "gnmt" or
    #        hparams.attention_architecture in ["gnmt", "gnmt_v2"])
    with mixed_precision_scope():
      model = gnmt_model.GNMTModel(hparams, mode=mode, features=features)
      if mode == tf.contrib.learn.ModeKeys.INFER:
        sample_ids = model.sample_id
        reverse_target_vocab_table = lookup_ops.index_to_string_table_from_file(
            hparams.tgt_vocab_file, default_value=vocab_utils.UNK)
        sample_words = reverse_target_vocab_table.lookup(
            tf.cast(sample_ids, tf.int64))
        # make sure outputs is of shape [batch_size, time] or [beam_width,
        # batch_size, time] when using beam search.
        if hparams.time_major:
          sample_words = tf.transpose(sample_words)
        elif sample_words.shape.ndims == 3:
          # beam search output in [batch_size, time, beam_width] shape.
          sample_words = tf.transpose(sample_words, [2, 0, 1])
        predictions = {"predictions": sample_words}
        # return loss, vars, grads, predictions, train_op, scaffold
        return None, None, None, predictions, None, None
      elif mode == tf.contrib.learn.ModeKeys.TRAIN:
        loss = model.train_loss
        train_op = model.update
        return loss, model.params, model.grads, None, train_op, None
      else:
        raise ValueError("Unknown mode in model_fn: %s" % mode)
//  [python]
def compute_last_embedding(input_embeddings, input_lengths, hparams):
  """Computes average of last K embedding.

  Args:
    input_embeddings: <tf.float32>[bs, max_seq_len, emb_dim]
    input_lengths: <tf.int64>[bs, 1]
    hparams: model hparams

  Returns:
    last_k_embedding: <tf.float32>[bs, emb_dim]
  """
  max_seq_len = tf.shape(input_embeddings)[1]
  # <tf.float32>[bs, 1, max_seq_len]
  mask = tf.sequence_mask(input_lengths, max_seq_len, dtype=tf.float32)
  del_mask = tf.sequence_mask(
      input_lengths - hparams.last_k, max_seq_len, dtype=tf.float32)
  final_mask = mask - del_mask
  # <tf.float32>[bs, 1, emb_dim]
  sum_embedding = tf.matmul(final_mask, input_embeddings)
  # <tf.float32>[bs, 1, emb_dim]
  last_k_embedding = sum_embedding / tf.to_float(
      tf.expand_dims(
          tf.ones([tf.shape(input_embeddings)[0], 1]) * hparams.last_k, 2))
  # <tf.float32>[bs, dim]
  return tf.squeeze(last_k_embedding, 1)
//  [python]
def compute_last_embedding(input_embeddings, input_lengths, hparams):
  """Computes average of last K embedding.

  Args:
    input_embeddings: <tf.float32>[bs, max_seq_len, emb_dim]
    input_lengths: <tf.int64>[bs, 1]
    hparams: model hparams

  Returns:
    last_k_embedding: <tf.float32>[bs, emb_dim]
  """
  max_seq_len = tf.shape(input_embeddings)[1]
  # <tf.float32>[bs, 1, max_seq_len]
  mask = tf.sequence_mask(input_lengths, max_seq_len, dtype=tf.float32)
  del_mask = tf.sequence_mask(
      input_lengths - hparams.last_k, max_seq_len, dtype=tf.float32)
  final_mask = mask - del_mask
  # <tf.float32>[bs, 1, emb_dim]
  sum_embedding = tf.matmul(final_mask, input_embeddings)
  # <tf.float32>[bs, 1, emb_dim]
  last_k_embedding = sum_embedding / tf.cast(
      tf.expand_dims(
          tf.ones([tf.shape(input_embeddings)[0], 1]) * hparams.last_k, 2), tf.float32)
  # <tf.float32>[bs, dim]
  return tf.squeeze(last_k_embedding, 1)
//  [python]
def masked_apply(tensor, op, mask, set_outside_zero=True):
    """
    Apply the function op to tensor only at locations indicated by mask. If set_outside_zero == True, set the
    locations outside the mask to zero, otherwise keep original value of tensor.
    :param tensor: The tensor on which op is applied.
    :param op: The operation.
    :param mask: The boolean mask.
    :param set_outside_zero: If True, set the locations outside the mask to zero, otherwise keep original values of tensor.
    :return: Tensor with applied function.
    """
    chosen = tf.boolean_mask(tensor, mask)
    applied = op(chosen)
    idx = tf.to_int32(tf.where(mask))
    result = tf.scatter_nd(idx, applied, tf.shape(tensor))
    if not set_outside_zero:
        result = tf.where(mask, result, tensor)
    return result
//  [python]
def masked_apply(tensor, op, mask, set_outside_zero=True):
    """
    Apply the function op to tensor only at locations indicated by mask. If set_outside_zero == True, set the
    locations outside the mask to zero, otherwise keep original value of tensor.
    :param tensor: The tensor on which op is applied.
    :param op: The operation.
    :param mask: The boolean mask.
    :param set_outside_zero: If True, set the locations outside the mask to zero, otherwise keep original values of tensor.
    :return: Tensor with applied function.
    """
    chosen = tf.boolean_mask(tensor, mask)
    applied = op(chosen)
    idx = tf.cast(tf.where(mask), tf.int32)
    result = tf.scatter_nd(idx, applied, tf.shape(tensor))
    if not set_outside_zero:
        result = tf.where(mask, result, tensor)
    return result
//  [python]
def bottom_simple(x, model_hparams, vocab_size, name, reuse):
  """Internal bottom transformation."""
  with tf.variable_scope(name, reuse=reuse):
    var = _get_weights(model_hparams, vocab_size)
    x = common_layers.dropout_no_scaling(
        x, 1.0 - model_hparams.symbol_dropout)
    # Add together the embeddings for each tuple position.
    ret = tf.add_n([
        tf.gather(var, x[:, :, :, i] + sum(vocab_size[:i])) *
        tf.expand_dims(tf.to_float(tf.not_equal(x[:, :, :, i], 0)), -1)
        for i in range(len(vocab_size))
    ])
    if model_hparams.multiply_embedding_mode == 'sqrt_depth':
      ret *= model_hparams.hidden_size**0.5
    return ret
//  [python]
def bottom_simple(x, model_hparams, vocab_size, name, reuse):
  """Internal bottom transformation."""
  with tf.variable_scope(name, reuse=reuse):
    var = _get_weights(model_hparams, vocab_size)
    x = common_layers.dropout_no_scaling(
        x, 1.0 - model_hparams.symbol_dropout)
    # Add together the embeddings for each tuple position.
    ret = tf.add_n([
        tf.gather(var, x[:, :, :, i] + sum(vocab_size[:i])) *
        tf.expand_dims(tf.cast(tf.not_equal(x[:, :, :, i], 0), tf.float32), -1)
        for i in range(len(vocab_size))
    ])
    if model_hparams.multiply_embedding_mode == 'sqrt_depth':
      ret *= model_hparams.hidden_size**0.5
    return ret
//  [python]
def flip_randomly(inputs, horizontally, vertically, is_training, name=None):
    """Flip images randomly. Make separate flipping decision for each image.

    Args:
        inputs (4-D tensor): Input images (batch size, height, width, channels).
        horizontally (bool): If True, flip horizontally with 50% probability. Otherwise, don't.
        vertically (bool): If True, flip vertically with 50% probability. Otherwise, don't.
        is_training (bool): If False, no flip is performed.
        scope: A name for the operation.
    """
    with tf.name_scope(name, "flip_randomly") as scope:
        batch_size, height, width, _ = tf.unstack(tf.shape(inputs))
        vertical_choices = (tf.random_uniform([batch_size], 0, 2, tf.int32) *
                            tf.to_int32(vertically) *
                            tf.to_int32(is_training))
        horizontal_choices = (tf.random_uniform([batch_size], 0, 2, tf.int32) *
                              tf.to_int32(horizontally) *
                              tf.to_int32(is_training))
        vertically_flipped = tf.reverse_sequence(inputs, vertical_choices * height, 1)
        both_flipped = tf.reverse_sequence(vertically_flipped, horizontal_choices * width, 2)
        return tf.identity(both_flipped, name=scope)
//  [python]
def flip_randomly(inputs, horizontally, vertically, is_training, name=None):
    """Flip images randomly. Make separate flipping decision for each image.

    Args:
        inputs (4-D tensor): Input images (batch size, height, width, channels).
        horizontally (bool): If True, flip horizontally with 50% probability. Otherwise, don't.
        vertically (bool): If True, flip vertically with 50% probability. Otherwise, don't.
        is_training (bool): If False, no flip is performed.
        scope: A name for the operation.
    """
    with tf.name_scope(name, "flip_randomly") as scope:
        batch_size, height, width, _ = tf.unstack(tf.shape(inputs))
        vertical_choices = (tf.random_uniform([batch_size], 0, 2, tf.int32) *
                            tf.cast(vertically, tf.int32) *
                            tf.cast(is_training, tf.int32))
        horizontal_choices = (tf.random_uniform([batch_size], 0, 2, tf.int32) *
                              tf.cast(horizontally, tf.int32) *
                              tf.cast(is_training, tf.int32))
        vertically_flipped = tf.reverse_sequence(inputs, vertical_choices * height, 1)
        both_flipped = tf.reverse_sequence(vertically_flipped, horizontal_choices * width, 2)
        return tf.identity(both_flipped, name=scope)
//  [python]
def anneal(c_max, step, iteration_threshold):
  """Anneal function for anneal_vae (https://arxiv.org/abs/1804.03599).

  Args:
    c_max: Maximum capacity.
    step: Current step.
    iteration_threshold: How many iterations to reach c_max.

  Returns:
    Capacity annealed linearly until c_max.
  """
  return tf.math.minimum(c_max * 1.,
                         c_max * 1. * tf.to_float(step) / iteration_threshold)
//  [python]
def anneal(c_max, step, iteration_threshold):
  """Anneal function for anneal_vae (https://arxiv.org/abs/1804.03599).

  Args:
    c_max: Maximum capacity.
    step: Current step.
    iteration_threshold: How many iterations to reach c_max.

  Returns:
    Capacity annealed linearly until c_max.
  """
  return tf.math.minimum(c_max * 1.,
                         c_max * 1. * tf.cast(step, tf.float32) / iteration_threshold)
//  [python]
def resize_by_area(img, size):
  """image resize function used by quite a few image problems."""
  return tf.to_int64(
      tf.image.resize_images(img, [size, size], tf.image.ResizeMethod.AREA))
//  [python]
def resize_by_area(img, size):
  """image resize function used by quite a few image problems."""
  return tf.cast(
      tf.image.resize_images(img, [size, size], tf.image.ResizeMethod.AREA), tf.int64)
//  [python]
def make_multiscale(image, resolutions,
                    resize_method=tf.image.ResizeMethod.BICUBIC,
                    num_channels=3):
  """Returns list of scaled images, one for each resolution.

  Args:
    image: Tensor of shape [height, height, num_channels].
    resolutions: List of heights that image's height is resized to.
    resize_method: tf.image.ResizeMethod.
    num_channels: Number of channels in image.

  Returns:
    List of Tensors, one for each resolution with shape given by
    [resolutions[i], resolutions[i], num_channels].
  """
  scaled_images = []
  for height in resolutions:
    scaled_image = tf.image.resize_images(
        image,
        size=[height, height],  # assuming that height = width
        method=resize_method)
    scaled_image = tf.to_int64(scaled_image)
    scaled_image.set_shape([height, height, num_channels])
    scaled_images.append(scaled_image)

  return scaled_images
//  [python]
def make_multiscale(image, resolutions,
                    resize_method=tf.image.ResizeMethod.BICUBIC,
                    num_channels=3):
  """Returns list of scaled images, one for each resolution.

  Args:
    image: Tensor of shape [height, height, num_channels].
    resolutions: List of heights that image's height is resized to.
    resize_method: tf.image.ResizeMethod.
    num_channels: Number of channels in image.

  Returns:
    List of Tensors, one for each resolution with shape given by
    [resolutions[i], resolutions[i], num_channels].
  """
  scaled_images = []
  for height in resolutions:
    scaled_image = tf.image.resize_images(
        image,
        size=[height, height],  # assuming that height = width
        method=resize_method)
    scaled_image = tf.cast(scaled_image, tf.int64)
    scaled_image.set_shape([height, height, num_channels])
    scaled_images.append(scaled_image)

  return scaled_images
//  [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.compat.v1.to_float(num_timescales) - 1))
  inv_timescales = min_timescale * tf.exp(
      tf.compat.v1.to_float(tf.range(num_timescales)) * -log_timescale_increment)
  scaled_time = (
      tf.expand_dims(tf.compat.v1.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.compat.v1.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.math.mod(channels, 2)]])
  return signal
//  [python]
def make_multiscale_dilated(image, resolutions, num_channels=3):
  """Returns list of scaled images, one for each resolution.

  Resizes by skipping every nth pixel.

  Args:
    image: Tensor of shape [height, height, num_channels].
    resolutions: List of heights that image's height is resized to. The function
      assumes VALID padding, so the original image's height must be divisible
      by each resolution's height to return the exact resolution size.
    num_channels: Number of channels in image.

  Returns:
    List of Tensors, one for each resolution with shape given by
    [resolutions[i], resolutions[i], num_channels] if resolutions properly
    divide the original image's height; otherwise shape height and width is up
    to valid skips.
  """
  image_height = common_layers.shape_list(image)[0]
  scaled_images = []
  for height in resolutions:
    dilation_rate = image_height // height  # assuming height = width
    scaled_image = image[::dilation_rate, ::dilation_rate]
    scaled_image = tf.to_int64(scaled_image)
    scaled_image.set_shape([None, None, num_channels])
    scaled_images.append(scaled_image)
  return scaled_images
//  [python]
def make_multiscale_dilated(image, resolutions, num_channels=3):
  """Returns list of scaled images, one for each resolution.

  Resizes by skipping every nth pixel.

  Args:
    image: Tensor of shape [height, height, num_channels].
    resolutions: List of heights that image's height is resized to. The function
      assumes VALID padding, so the original image's height must be divisible
      by each resolution's height to return the exact resolution size.
    num_channels: Number of channels in image.

  Returns:
    List of Tensors, one for each resolution with shape given by
    [resolutions[i], resolutions[i], num_channels] if resolutions properly
    divide the original image's height; otherwise shape height and width is up
    to valid skips.
  """
  image_height = common_layers.shape_list(image)[0]
  scaled_images = []
  for height in resolutions:
    dilation_rate = image_height // height  # assuming height = width
    scaled_image = image[::dilation_rate, ::dilation_rate]
    scaled_image = tf.cast(scaled_image, tf.int64)
    scaled_image.set_shape([None, None, num_channels])
    scaled_images.append(scaled_image)
  return scaled_images
//  [python]
def call(self, inputs: tuple, **kwargs):
        """
        :param inputs: (input_tensor, seq_len)
        input_tensor with shape [batch_size, max_seq_len, x]
        seq_len with shape [batch_size]
        """
        input_tensor, seq_len = inputs
        mask = tf.expand_dims(tf.sequence_mask(seq_len, dtype=tf.float32), -1)
        # zeroed padding vectors
        input_tensor = mask * input_tensor

        return tf.reduce_sum(input_tensor, axis=-2) / tf.expand_dims(tf.to_float(seq_len), axis=-1)
//  [python]
def call(self, inputs: tuple, **kwargs):
        """
        :param inputs: (input_tensor, seq_len)
        input_tensor with shape [batch_size, max_seq_len, x]
        seq_len with shape [batch_size]
        """
        input_tensor, seq_len = inputs
        mask = tf.expand_dims(tf.sequence_mask(seq_len, dtype=tf.float32), -1)
        # zeroed padding vectors
        input_tensor = mask * input_tensor

        return tf.reduce_sum(input_tensor, axis=-2) / tf.expand_dims(tf.cast(seq_len, tf.float32), axis=-1)
//  [python]
def get_anneal_mask(hparams):
  """Get anneal and kl mask."""
  startup = hparams.kl_startup_steps
  anneal = hparams.kl_anneal_steps
  global_step = tf.train.get_global_step()
  min_value = hparams.anneal_min_value
  step = tf.maximum(global_step - startup, 0)
  anneal = common_layers.inverse_lin_decay(
      anneal, min_value=min_value, step=step)
  kl_mask = tf.less(startup, tf.to_int32(global_step))
  kl_mask = tf.cast(kl_mask, tf.float32)
  return anneal, kl_mask
//  [python]
def get_anneal_mask(hparams):
  """Get anneal and kl mask."""
  startup = hparams.kl_startup_steps
  anneal = hparams.kl_anneal_steps
  global_step = tf.train.get_global_step()
  min_value = hparams.anneal_min_value
  step = tf.maximum(global_step - startup, 0)
  anneal = common_layers.inverse_lin_decay(
      anneal, min_value=min_value, step=step)
  kl_mask = tf.less(startup, tf.cast(global_step, tf.int32))
  kl_mask = tf.cast(kl_mask, tf.float32)
  return anneal, kl_mask
//  [python]
def prepare_data(problem, hparams, params, config):
  """Construct input pipeline."""
  input_fn = problem.make_estimator_input_fn(
      tf.estimator.ModeKeys.EVAL, hparams, force_repeat=True)
  dataset = input_fn(params, config)
  features, _ = dataset.make_one_shot_iterator().get_next()
  inputs, labels = features["targets"], features["inputs"]
  inputs = tf.to_float(inputs)
  input_shape = inputs.shape.as_list()
  inputs = tf.reshape(inputs, [hparams.batch_size] + input_shape[1:])
  labels = tf.reshape(labels, [hparams.batch_size])
  return inputs, labels, features
//  [python]
def prepare_data(problem, hparams, params, config):
  """Construct input pipeline."""
  input_fn = problem.make_estimator_input_fn(
      tf.estimator.ModeKeys.EVAL, hparams, force_repeat=True)
  dataset = input_fn(params, config)
  features, _ = dataset.make_one_shot_iterator().get_next()
  inputs, labels = features["targets"], features["inputs"]
  inputs = tf.cast(inputs, tf.float32)
  input_shape = inputs.shape.as_list()
  inputs = tf.reshape(inputs, [hparams.batch_size] + input_shape[1:])
  labels = tf.reshape(labels, [hparams.batch_size])
  return inputs, labels, features
//  [python]
def _smallest_size_at_least(height, width, smallest_side):
  """Computes new shape with the smallest side equal to `smallest_side`.

    Computes new shape with the smallest side equal to `smallest_side` while
    preserving the original aspect ratio.

  Args:
    height: an int32 scalar tensor indicating the current height.
    width: an int32 scalar tensor indicating the current width.
    smallest_side: A python integer or scalar `Tensor` indicating the size of
    the smallest side after resize.

  Returns:
    new_height: an int32 scalar tensor indicating the new height.
    new_width: and int32 scalar tensor indicating the new width.
  """
  smallest_side = tf.convert_to_tensor(smallest_side, dtype=tf.int32)

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

  scale = tf.cond(
      tf.greater(height, width), lambda: smallest_side / width,
      lambda: smallest_side / height)
  new_height = tf.to_int32(height * scale)
  new_width = tf.to_int32(width * scale)
  return new_height, new_width
//  [python]
def _smallest_size_at_least(height, width, smallest_side):
  """Computes new shape with the smallest side equal to `smallest_side`.

    Computes new shape with the smallest side equal to `smallest_side` while
    preserving the original aspect ratio.

  Args:
    height: an int32 scalar tensor indicating the current height.
    width: an int32 scalar tensor indicating the current width.
    smallest_side: A python integer or scalar `Tensor` indicating the size of
    the smallest side after resize.

  Returns:
    new_height: an int32 scalar tensor indicating the new height.
    new_width: and int32 scalar tensor indicating the new width.
  """
  smallest_side = tf.convert_to_tensor(smallest_side, dtype=tf.int32)

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

  scale = tf.cond(
      tf.greater(height, width), lambda: smallest_side / width,
      lambda: smallest_side / height)
  new_height = tf.to_int32(height * scale)
  new_width = tf.to_int32(width * scale)
  return new_height, new_width
//  [python]
def sari_score(predictions, labels, features, **unused_kwargs):
  """Computes the SARI scores from the given source, prediction and targets.

  An approximate SARI scoring method since we do not glue word pieces or
  decode the ids and tokenize the output. By default, we use ngram order of 4.
  Also, this does not have beam search.

  Args:
    predictions: tensor, model predictions.
    labels: tensor, gold output.
    features: dict, containing inputs.

  Returns:
    sari: int, approx sari score
  """
  if "inputs" not in features:
    raise ValueError("sari_score requires inputs feature")

  # Convert the inputs and outputs to a [batch_size, sequence_length] tensor.
  inputs = tf.squeeze(features["inputs"], axis=[-1, -2])
  outputs = tf.to_int32(tf.argmax(predictions, axis=-1))
  outputs = tf.squeeze(outputs, axis=[-1, -2])

  # Convert the labels to a [batch_size, 1, sequence_length] tensor.
  labels = tf.squeeze(labels, axis=[-1, -2])
  labels = tf.expand_dims(labels, axis=1)

  score, _, _, _ = get_sari(inputs, outputs, labels)
  return score, tf.constant(1.0)
//  [python]
def sari_score(predictions, labels, features, **unused_kwargs):
  """Computes the SARI scores from the given source, prediction and targets.

  An approximate SARI scoring method since we do not glue word pieces or
  decode the ids and tokenize the output. By default, we use ngram order of 4.
  Also, this does not have beam search.

  Args:
    predictions: tensor, model predictions.
    labels: tensor, gold output.
    features: dict, containing inputs.

  Returns:
    sari: int, approx sari score
  """
  if "inputs" not in features:
    raise ValueError("sari_score requires inputs feature")

  # Convert the inputs and outputs to a [batch_size, sequence_length] tensor.
  inputs = tf.squeeze(features["inputs"], axis=[-1, -2])
  outputs = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)
  outputs = tf.squeeze(outputs, axis=[-1, -2])

  # Convert the labels to a [batch_size, 1, sequence_length] tensor.
  labels = tf.squeeze(labels, axis=[-1, -2])
  labels = tf.expand_dims(labels, axis=1)

  score, _, _, _ = get_sari(inputs, outputs, labels)
  return score, tf.constant(1.0)
//  [python]
def cross_replica_concat(value, replica_id, num_replicas):
  """Reduce a concatenation of the `value` across TPU replicas.

  Args:
    value: Tensor to concatenate.
    replica_id: Integer tensor that indicates the index of the replica.
    num_replicas: Python integer, total number of replicas.

  Returns:
    Tensor of the same rank as value with first dimension `num_replicas`
    times larger.

  Raises:
    ValueError: If `value` is a scalar.
  """
  if value.shape.ndims < 1:
    raise ValueError("Value must have at least rank 1 but got {}.".format(
        value.shape.ndims))
  if num_replicas <= 1:
    return value
  with tf.name_scope(None, "tpu_cross_replica_concat"):
    # Mask is one hot encoded position of the core_index.
    mask = tf.to_float(tf.equal(tf.range(num_replicas), replica_id))
    # Expand dims with 1's to match rank of value.
    mask = tf.reshape(mask, [num_replicas] + [1] * value.shape.ndims)
    if value.dtype in {tf.bfloat16, tf.float32}:
      result = mask * value
    else:
      result = mask * tf.to_float(value)
    # Thanks to broadcasting now result is set only in the position pointed by
    # replica_id, the rest of the vector is set to 0's.
    # All these steps are basically implementing tf.scatter_nd which is missing
    # in TPU's backend since it doesn't support sparse operations.

    # Merge first 2 dimensions.
    # This is equivalent to (value.shape[0].value * num_replicas).
    # Using [-1] trick to support also scalar input.
    result = tf.reshape(result, [-1] + result.shape.as_list()[2:])
    # Each core set the "results" in position pointed by replica_id. When we now
    # sum across replicas we exchange the information and fill in local 0's with
    # values from other cores.
    result = tf.tpu.cross_replica_sum(result)
    # Now all the cores see exactly the same data.
    return tf.cast(result, dtype=value.dtype)
//  [python]
def cross_replica_concat(value, replica_id, num_replicas):
  """Reduce a concatenation of the `value` across TPU replicas.

  Args:
    value: Tensor to concatenate.
    replica_id: Integer tensor that indicates the index of the replica.
    num_replicas: Python integer, total number of replicas.

  Returns:
    Tensor of the same rank as value with first dimension `num_replicas`
    times larger.

  Raises:
    ValueError: If `value` is a scalar.
  """
  if value.shape.ndims < 1:
    raise ValueError("Value must have at least rank 1 but got {}.".format(
        value.shape.ndims))
  if num_replicas <= 1:
    return value
  with tf.name_scope(None, "tpu_cross_replica_concat"):
    # Mask is one hot encoded position of the core_index.
    mask = tf.cast(tf.equal(tf.range(num_replicas), replica_id), tf.float32)
    # Expand dims with 1's to match rank of value.
    mask = tf.reshape(mask, [num_replicas] + [1] * value.shape.ndims)
    if value.dtype in {tf.bfloat16, tf.float32}:
      result = mask * value
    else:
      result = mask * tf.cast(value, tf.float32)
    # Thanks to broadcasting now result is set only in the position pointed by
    # replica_id, the rest of the vector is set to 0's.
    # All these steps are basically implementing tf.scatter_nd which is missing
    # in TPU's backend since it doesn't support sparse operations.

    # Merge first 2 dimensions.
    # This is equivalent to (value.shape[0].value * num_replicas).
    # Using [-1] trick to support also scalar input.
    result = tf.reshape(result, [-1] + result.shape.as_list()[2:])
    # Each core set the "results" in position pointed by replica_id. When we now
    # sum across replicas we exchange the information and fill in local 0's with
    # values from other cores.
    result = tf.tpu.cross_replica_sum(result)
    # Now all the cores see exactly the same data.
    return tf.cast(result, dtype=value.dtype)
//  [python]
def attention_2(self,M):
        #######这里采用的注意力打分机制是使用的是出去实体的其它部分原文本#####################
        #https://zhuanlan.zhihu.com/p/53682800
        #M:[batch_size,max_len,2*hidden],
        e = tf.to_float(tf.expand_dims(self.input_segment, -1))
        #M: [batch_size,max_len,2*hidden] e1:[batch_size,max_len,1],相乘后：[batch_size,max_len,2*hidden]
        e1 = tf.where(tf.equal(e,1),e,e-e)
        e1 = tf.multiply(M,e1)
        e2 = tf.where(tf.equal(e,-1),0-e,e-e)
        e2 = tf.multiply(M,e2)
        
        #e1,e2,text:[batch_size,max_len,1]
        e1 = tf.tanh(tf.nn.conv1d(e1,self.e1_w,1,'SAME'))
        e2 = tf.tanh(tf.nn.conv1d(e2,self.e2_w,1,'SAME'))
        text = tf.tanh(tf.nn.conv1d(M,self.text_w,1,'SAME'))
        score = tf.nn.softmax(tf.subtract(text,tf.add(e1,e2)),1)
        M = tf.transpose(M, [0,2,1])
        return tf.matmul(M,score)
//  [python]
def attention_2(self,M):
        #######这里采用的注意力打分机制是使用的是出去实体的其它部分原文本#####################
        #https://zhuanlan.zhihu.com/p/53682800
        #M:[batch_size,max_len,2*hidden],
        e = tf.cast(tf.expand_dims(self.input_segment, -1),dtype=tf.float32)
        #M: [batch_size,max_len,2*hidden] e1:[batch_size,max_len,1],相乘后：[batch_size,max_len,2*hidden]
        e1 = tf.where(tf.equal(e,1),e,e-e)
        e1 = tf.multiply(M,e1)
        e2 = tf.where(tf.equal(e,-1),0-e,e-e)
        e2 = tf.multiply(M,e2)
        
        #e1,e2,text:[batch_size,max_len,1]
        e1 = tf.tanh(tf.nn.conv1d(e1,self.e1_w,1,'SAME'))
        e2 = tf.tanh(tf.nn.conv1d(e2,self.e2_w,1,'SAME'))
        text = tf.tanh(tf.nn.conv1d(M,self.text_w,1,'SAME'))
        score = tf.nn.softmax(tf.subtract(text,tf.add(e1,e2)),1)
        M = tf.transpose(M, [0,2,1])
        return tf.matmul(M,score)
//  [python]
def testMultiheadSelfAttentionMemoryEfficient(self):
    num_heads = 4
    io_size = 16
    batch = 2
    length = 7
    head_size = 5
    x = np.random.rand(batch, length, io_size)
    dy = np.random.rand(batch, length, io_size)
    with self.test_session() as session:
      x = tf.to_float(x)
      dy = tf.to_float(dy)
      bias = common_attention.attention_bias_lower_triangle(length)
      wqkv = tf.get_variable(
          "wqkv", [num_heads, 1, io_size, 3 * head_size],
          initializer=tf.random_normal_initializer(stddev=io_size**-0.5))
      wo = tf.get_variable(
          "wo", [num_heads, 1, head_size, io_size],
          initializer=tf.random_normal_initializer(
              stddev=(head_size * num_heads)**-0.5))
      norm_scale, norm_bias = common_layers.layer_norm_vars(io_size)
      y = common_attention.multihead_self_attention_memory_efficient(
          x, bias, num_heads, head_size=head_size, forget=False,
          test_vars=(wqkv, wo, norm_scale, norm_bias))
      y_forget = common_attention.multihead_self_attention_memory_efficient(
          x, bias, num_heads, head_size=head_size, forget=True,
          test_vars=(wqkv, wo, norm_scale, norm_bias))
      dx, dwqkv, dwo, dnorm_scale, dnorm_bias = tf.gradients(
          ys=[y], xs=[x, wqkv, wo, norm_scale, norm_bias], grad_ys=[dy])
      dx_f, dwqkv_f, dwo_f, dnorm_scale_f, dnorm_bias_f = tf.gradients(
          ys=[y_forget], xs=[x, wqkv, wo, norm_scale, norm_bias], grad_ys=[dy])
      session.run(tf.global_variables_initializer())
      (y, y_forget,
       dx, dwqkv, dwo, dnorm_scale, dnorm_bias,
       dx_f, dwqkv_f, dwo_f, dnorm_scale_f, dnorm_bias_f) = session.run(
           [y, y_forget,
            dx, dwqkv, dwo, dnorm_scale, dnorm_bias,
            dx_f, dwqkv_f, dwo_f, dnorm_scale_f, dnorm_bias_f])
    self.assertAllClose(y, y_forget)
    self.assertAllClose(dwo, dwo_f)
    self.assertAllClose(dwqkv, dwqkv_f)
    self.assertAllClose(dnorm_scale, dnorm_scale_f)
    self.assertAllClose(dnorm_bias, dnorm_bias_f)
    self.assertAllClose(dx, dx_f)
//  [python]
def testMultiheadSelfAttentionMemoryEfficient(self):
    num_heads = 4
    io_size = 16
    batch = 2
    length = 7
    head_size = 5
    x = np.random.rand(batch, length, io_size)
    dy = np.random.rand(batch, length, io_size)
    with self.test_session() as session:
      x = tf.cast(x, tf.float32)
      dy = tf.cast(dy, tf.float32)
      bias = common_attention.attention_bias_lower_triangle(length)
      wqkv = tf.get_variable(
          "wqkv", [num_heads, 1, io_size, 3 * head_size],
          initializer=tf.random_normal_initializer(stddev=io_size**-0.5))
      wo = tf.get_variable(
          "wo", [num_heads, 1, head_size, io_size],
          initializer=tf.random_normal_initializer(
              stddev=(head_size * num_heads)**-0.5))
      norm_scale, norm_bias = common_layers.layer_norm_vars(io_size)
      y = common_attention.multihead_self_attention_memory_efficient(
          x, bias, num_heads, head_size=head_size, forget=False,
          test_vars=(wqkv, wo, norm_scale, norm_bias))
      y_forget = common_attention.multihead_self_attention_memory_efficient(
          x, bias, num_heads, head_size=head_size, forget=True,
          test_vars=(wqkv, wo, norm_scale, norm_bias))
      dx, dwqkv, dwo, dnorm_scale, dnorm_bias = tf.gradients(
          ys=[y], xs=[x, wqkv, wo, norm_scale, norm_bias], grad_ys=[dy])
      dx_f, dwqkv_f, dwo_f, dnorm_scale_f, dnorm_bias_f = tf.gradients(
          ys=[y_forget], xs=[x, wqkv, wo, norm_scale, norm_bias], grad_ys=[dy])
      session.run(tf.global_variables_initializer())
      (y, y_forget,
       dx, dwqkv, dwo, dnorm_scale, dnorm_bias,
       dx_f, dwqkv_f, dwo_f, dnorm_scale_f, dnorm_bias_f) = session.run(
           [y, y_forget,
            dx, dwqkv, dwo, dnorm_scale, dnorm_bias,
            dx_f, dwqkv_f, dwo_f, dnorm_scale_f, dnorm_bias_f])
    self.assertAllClose(y, y_forget)
    self.assertAllClose(dwo, dwo_f)
    self.assertAllClose(dwqkv, dwqkv_f)
    self.assertAllClose(dnorm_scale, dnorm_scale_f)
    self.assertAllClose(dnorm_bias, dnorm_bias_f)
    self.assertAllClose(dx, dx_f)
//  [python]
def _process_inputs(self, x, observed_mask, missing_mask):
        observed_mask = t.cast(observed_mask, x.dtype)
        query = 1.0 - observed_mask

        if missing_mask is not None:
            missing_mask = t.cast(missing_mask, x.dtype)
            query *= 1.0 - missing_mask
            observed_mask *= 1.0 - missing_mask

        x_o = x * observed_mask
        x_u = x * query

        return x_o, x_u, observed_mask, query
//  [python]
def _process_inputs(self, x, observed_mask, missing_mask):
        observed_mask = tf.cast(observed_mask, x.dtype)
        query = 1.0 - observed_mask

        if missing_mask is not None:
            missing_mask = tf.cast(missing_mask, x.dtype)
            query *= 1.0 - missing_mask
            observed_mask *= 1.0 - missing_mask

        x_o = x * observed_mask
        x_u = x * query

        return x_o, x_u, observed_mask, query
//  [python]
def _dim_energy_sample(self, x_o, observed_mask, index, num_resampling_samples):
        outputs = self(
            [x_o, observed_mask],
            num_importance_samples=num_resampling_samples,
            selected_features=tf.broadcast_to(index, (tf.shape(x_o)[0], 1)),
        )

        resampled_inds = tf.to_int32(tf.random.categorical(outputs.log_ratios, 1))
        winners = tf.gather(
            outputs.proposal_samples[..., 0], resampled_inds, batch_dims=1
        )
        return winners[:, 0]
//  [python]
def _dim_energy_sample(self, x_o, observed_mask, index, num_resampling_samples):
        outputs = self(
            [x_o, observed_mask],
            num_importance_samples=num_resampling_samples,
            selected_features=tf.broadcast_to(index, (tf.shape(x_o)[0], 1)),
        )

        resampled_inds = tf.cast(tf.random.categorical(outputs.log_ratios, 1), tf.int32)
        winners = tf.gather(
            outputs.proposal_samples[..., 0], resampled_inds, batch_dims=1
        )
        return winners[:, 0]
//  [python]
def blend(image1, image2, factor):
  """Blend image1 and image2 using 'factor'.

  Factor can be above 0.0.  A value of 0.0 means only image1 is used.
  A value of 1.0 means only image2 is used.  A value between 0.0 and
  1.0 means we linearly interpolate the pixel values between the two
  images.  A value greater than 1.0 "extrapolates" the difference
  between the two pixel values, and we clip the results to values
  between 0 and 255.

  Args:
    image1: An image Tensor of type uint8.
    image2: An image Tensor of type uint8.
    factor: A floating point value above 0.0.

  Returns:
    A blended image Tensor of type uint8.
  """
  if factor == 0.0:
    return tf1.convert_to_tensor(image1)
  if factor == 1.0:
    return tf1.convert_to_tensor(image2)

  image1 = tf1.to_float(image1)
  image2 = tf1.to_float(image2)

  difference = image2 - image1
  scaled = factor * difference

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

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

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

  Args:
    image: A 3D uint8 tensor.
    seed: the random seed.

  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """

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

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

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

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = tf1.stack([s1, s2, s3], 2)
  return image
//  [python]
def autocontrast(image, seed):
  """Implements Autocontrast function from PIL using TF ops.

  Args:
    image: A 3D uint8 tensor.
    seed: the random seed.

  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """

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

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

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

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

    def replica_step(x):
      x = t.cast(x, tf.float32)
      return tf.reduce_sum(x)

    return self.strategy.experimental_local_results(
        self.strategy.run(replica_step, args=(next(iterator),)))
//  [python]
def eval_step(self, iterator):

    def replica_step(x):
      x = tf.cast(x, tf.float32)
      return tf.reduce_sum(x)

    return self.strategy.experimental_local_results(
        self.strategy.run(replica_step, args=(next(iterator),)))
//  [python]
def _compute_clip_window(self, image_shapes):
    """Computes clip window for non max suppression based on image shapes.

    This function assumes that the clip window's left top corner is at (0, 0).

    Args:
      image_shapes: A 2-D int32 tensor of shape [batch_size, 3] containing
      shapes of images in the batch. Each row represents [height, width,
      channels] of an image.

    Returns:
      A 2-D float32 tensor of shape [batch_size, 4] containing the clip window
      for each image in the form [ymin, xmin, ymax, xmax].
    """
    clip_heights = image_shapes[:, 0]
    clip_widths = image_shapes[:, 1]
    clip_window = t.cast(
        t.stack([
            t.zeros_like(clip_heights),
            t.zeros_like(clip_heights), clip_heights, clip_widths
        ],
                 axis=1),
        dtype=t.float32)
    return clip_window
//  [python]
def _compute_clip_window(self, image_shapes):
    """Computes clip window for non max suppression based on image shapes.

    This function assumes that the clip window's left top corner is at (0, 0).

    Args:
      image_shapes: A 2-D int32 tensor of shape [batch_size, 3] containing
      shapes of images in the batch. Each row represents [height, width,
      channels] of an image.

    Returns:
      A 2-D float32 tensor of shape [batch_size, 4] containing the clip window
      for each image in the form [ymin, xmin, ymax, xmax].
    """
    clip_heights = image_shapes[:, 0]
    clip_widths = image_shapes[:, 1]
    clip_window = tf.cast(
        tf.stack([
            tf.zeros_like(clip_heights),
            tf.zeros_like(clip_heights), clip_heights, clip_widths
        ],
                 axis=1),
        dtype=tf.float32)
    return clip_window
//  [python]
def _gather_instance_masks(self, instance_masks, classes):
    """Gathers the masks that correspond to classes.

    Args:
      instance_masks: A 4-D float32 tensor with shape
        [K, num_classes, mask_height, mask_width].
      classes: A 2-D int32 tensor with shape [batch_size, max_detection].

    Returns:
      masks: a 3-D float32 tensor with shape [K, mask_height, mask_width].
    """
    _, num_classes, height, width = instance_masks.get_shape().as_list()
    k = tf.shape(instance_masks)[0]
    instance_masks = tf.reshape(instance_masks, [-1, height, width])
    classes = tf.cast(tf.reshape(classes, [-1]), dtype=tf.int32)
    gather_idx = tf.range(k) * num_classes + classes
    return tf.gather(instance_masks, gather_idx)
//  [python]
def _single_permutation_ll(
        self,
        x,
        observed_mask,
        missing_mask=None,
        num_importance_samples=20,
        **map_fn_kwargs,
    ):
        if missing_mask is None:
            missing_mask = tf.broadcast_to(0.0, tf.shape(observed_mask))

        query = (1.0 - observed_mask) * (1.0 - missing_mask)
        observed_mask *= 1.0 - missing_mask

        expanded_x, expanded_observed_mask, selected_features = tf.map_fn(
            _get_autoregressive_ll_batch,
            [x, observed_mask, query],
            fn_output_signature=(
                tf.RaggedTensorSpec((None, x.shape[-1]), x.dtype),
                tf.RaggedTensorSpec((None, x.shape[-1]), observed_mask.dtype),
                tf.RaggedTensorSpec((None, 1), tf.int64),
            ),
            **map_fn_kwargs,
        )

        query_total = tf.math.count_nonzero(query)
        batch_x = expanded_x.merge_dims(0, 1).to_tensor(
            shape=(query_total, self._num_features)
        )
        batch_observed_mask = expanded_observed_mask.merge_dims(0, 1).to_tensor()
        batch_selected_features = tf.expand_dims(selected_features.merge_dims(0, 2), 1)
        batch_selected_features = tf.to_int32(batch_selected_features)

        outputs = self(
            [batch_x, batch_observed_mask],
            num_importance_samples=num_importance_samples,
            selected_features=batch_selected_features,
        )

        u_count = tf.math.count_nonzero(query, axis=-1)

        energy_ll = tf.RaggedTensor.from_row_lengths(
            tf.squeeze(outputs.energy_ll, -1), u_count
        )
        proposal_ll = tf.RaggedTensor.from_row_lengths(
            tf.squeeze(outputs.proposal_ll, 1), u_count
        )

        energy_ll = tf.reduce_sum(energy_ll, axis=-1)
        proposal_ll = tf.reduce_sum(proposal_ll, axis=-1)

        return energy_ll, proposal_ll