jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _create_predictions(self, decoder_output, features, labels, losses=None):
    """Creates the dictionary of predictions that is returned by the model.
    """
    predictions = {}

    # Add features and, if available, labels to predictions
    predictions.update(_flatten_dict({"features": features}))
    if labels is not None:
      predictions.update(_flatten_dict({"labels": labels}))

    if losses is not None:
      predictions["losses"] = _transpose_batch_time(losses)

    # Decoders returns output in time-major form [T, B, ...]
    # Here we transpose everything back to batch-major for the user
    output_dict = collections.OrderedDict(
        zip(decoder_output._fields, decoder_output))
    decoder_output_flat = _flatten_dict(output_dict)
    
    decoder_output_flat = {
        k: _transpose_batch_time(v)
        for k, v in decoder_output_flat.items()
    }
    predictions.update(decoder_output_flat)

    # If we predict the ids also map them back into the vocab and process them
    if "predicted_ids" in predictions.keys():
      vocab_tables = graph_utils.get_dict_from_collection("vocab_tables")
      target_id_to_vocab = vocab_tables["target_id_to_vocab"]
      predicted_tokens = target_id_to_vocab.lookup(
          tf.cast(predictions["predicted_ids"], tf.int64))
      # Raw predicted tokens
      predictions["predicted_tokens"] = predicted_tokens

    return predictions
//  [python]
def _using_motion_vector(depth, translation, rotation_angles, intrinsic_mat):
  """A helper for using_motion_vector. See docstring therein."""

  if translation.shape.ndims not in (2, 4):
    raise ValueError('\'translation\' should have rank 2 or 4, not %d' %
                     translation.shape.ndims)
  if translation.shape[-1] != 3:
    raise ValueError('translation\'s last dimension should be 3, not %d' %
                     translation.shape[1])
  if translation.shape.ndims == 2:
    translation = tf.expand_dims(tf.expand_dims(translation, 1), 1)

  _, height, width = tf.unstack(tf.shape(depth))
  grid = tf.squeeze(
      tf.stack(tf.meshgrid(tf.range(width), tf.range(height), (1,))), axis=3)
  grid = tf.to_float(grid)
  intrinsic_mat_inv = tf.linalg.inv(intrinsic_mat)

  rot_mat = transform_utils.matrix_from_angles(rotation_angles)
  # We have to treat separately the case of a per-image rotation vector and a
  # per-image rotation field, because the broadcasting capabilities of einsum
  # are limited.
  if rotation_angles.shape.ndims == 2:
    # The calculation here is identical to the one in inverse_warp above.
    # Howeverwe use einsum for better clarity. Under the hood, einsum performs
    # the reshaping and invocation of BatchMatMul, instead of doing it manually,
    # as in inverse_warp.
    projected_rotation = tf.einsum('bij,bjk,bkl->bil', intrinsic_mat, rot_mat,
                                   intrinsic_mat_inv)
    pcoords = tf.einsum('bij,jhw,bhw->bihw', projected_rotation, grid, depth)
  elif rotation_angles.shape.ndims == 4:
    # We push the H and W dimensions to the end, and transpose the rotation
    # matrix elements (as noted above).
    rot_mat = tf.transpose(rot_mat, [0, 3, 4, 1, 2])
    projected_rotation = tf.einsum('bij,bjkhw,bkl->bilhw', intrinsic_mat,
                                   rot_mat, intrinsic_mat_inv)
    pcoords = tf.einsum('bijhw,jhw,bhw->bihw', projected_rotation, grid, depth)

  projected_translation = tf.einsum('bij,bhwj->bihw', intrinsic_mat,
                                    translation)
  pcoords += projected_translation
  x, y, z = tf.unstack(pcoords, axis=1)
  return x / z, y / z, z
//  [python]
def _using_motion_vector(depth, translation, rotation_angles, intrinsic_mat):
  """A helper for using_motion_vector. See docstring therein."""

  if translation.shape.ndims not in (2, 4):
    raise ValueError('\'translation\' should have rank 2 or 4, not %d' %
                     translation.shape.ndims)
  if translation.shape[-1] != 3:
    raise ValueError('translation\'s last dimension should be 3, not %d' %
                     translation.shape[1])
  if translation.shape.ndims == 2:
    translation = tf.expand_dims(tf.expand_dims(translation, 1), 1)

  _, height, width = tf.unstack(tf.shape(depth))
  grid = tf.squeeze(
      tf.stack(tf.meshgrid(tf.range(width), tf.range(height), (1,))), axis=3)
  grid = tf.cast(grid, tf.float32)
  intrinsic_mat_inv = tf.linalg.inv(intrinsic_mat)

  rot_mat = transform_utils.matrix_from_angles(rotation_angles)
  # We have to treat separately the case of a per-image rotation vector and a
  # per-image rotation field, because the broadcasting capabilities of einsum
  # are limited.
  if rotation_angles.shape.ndims == 2:
    # The calculation here is identical to the one in inverse_warp above.
    # Howeverwe use einsum for better clarity. Under the hood, einsum performs
    # the reshaping and invocation of BatchMatMul, instead of doing it manually,
    # as in inverse_warp.
    projected_rotation = tf.einsum('bij,bjk,bkl->bil', intrinsic_mat, rot_mat,
                                   intrinsic_mat_inv)
    pcoords = tf.einsum('bij,jhw,bhw->bihw', projected_rotation, grid, depth)
  elif rotation_angles.shape.ndims == 4:
    # We push the H and W dimensions to the end, and transpose the rotation
    # matrix elements (as noted above).
    rot_mat = tf.transpose(rot_mat, [0, 3, 4, 1, 2])
    projected_rotation = tf.einsum('bij,bjkhw,bkl->bilhw', intrinsic_mat,
                                   rot_mat, intrinsic_mat_inv)
    pcoords = tf.einsum('bijhw,jhw,bhw->bihw', projected_rotation, grid, depth)

  projected_translation = tf.einsum('bij,bhwj->bihw', intrinsic_mat,
                                    translation)
  pcoords += projected_translation
  x, y, z = tf.unstack(pcoords, axis=1)
  return x / z, y / z, z
//  [python]
def __init__(
      self,
      batch_size,
      optimizer_name,
      optimizer_kwargs,
      init_w,
  ):
    """A simple DQN agent.

    Args:
      batch_size: Size of update batch.
      optimizer_name: Name of an optimizer from tf.train
      optimizer_kwargs: Keyword arguments for the optimizer.
      init_w: The initial cumulant weight.
    """
    self._batch_size = batch_size
    self._init_w = np.array(init_w)
    self._replay = []

    # Regress w by gradient descent, could also use closed-form solution.
    self._n_cumulants = len(init_w)
    self._regressed_w = tf.get_variable(
        "regressed_w",
        dtype=tf.float32,
        initializer=lambda: tf.to_float(init_w))
    cumulants_ph = tf.placeholder(
        shape=(None, self._n_cumulants), dtype=tf.float32)
    rewards_ph = tf.placeholder(shape=(None,), dtype=tf.float32)
    predicted_rewards = tf.reduce_sum(
        tf.multiply(self._regressed_w, cumulants_ph), axis=-1)
    loss = tf.reduce_sum(tf.square(predicted_rewards - rewards_ph))

    with tf.variable_scope("optimizer"):
      self._optimizer = getattr(tf.train, optimizer_name)(**optimizer_kwargs)
      train_op = self._optimizer.minimize(loss)

    # Make session and callables.
    session = tf.Session()
    self._update_fn = session.make_callable(train_op,
                                            [cumulants_ph, rewards_ph])
    self._action = session.make_callable(self._regressed_w.read_value(), [])
    session.run(tf.global_variables_initializer())
//  [python]
def __init__(
      self,
      batch_size,
      optimizer_name,
      optimizer_kwargs,
      init_w,
  ):
    """A simple DQN agent.

    Args:
      batch_size: Size of update batch.
      optimizer_name: Name of an optimizer from tf.train
      optimizer_kwargs: Keyword arguments for the optimizer.
      init_w: The initial cumulant weight.
    """
    self._batch_size = batch_size
    self._init_w = np.array(init_w)
    self._replay = []

    # Regress w by gradient descent, could also use closed-form solution.
    self._n_cumulants = len(init_w)
    self._regressed_w = tf.get_variable(
        "regressed_w",
        dtype=tf.float32,
        initializer=lambda: tf.cast(init_w, tf.float32))
    cumulants_ph = tf.placeholder(
        shape=(None, self._n_cumulants), dtype=tf.float32)
    rewards_ph = tf.placeholder(shape=(None,), dtype=tf.float32)
    predicted_rewards = tf.reduce_sum(
        tf.multiply(self._regressed_w, cumulants_ph), axis=-1)
    loss = tf.reduce_sum(tf.square(predicted_rewards - rewards_ph))

    with tf.variable_scope("optimizer"):
      self._optimizer = getattr(tf.train, optimizer_name)(**optimizer_kwargs)
      train_op = self._optimizer.minimize(loss)

    # Make session and callables.
    session = tf.Session()
    self._update_fn = session.make_callable(train_op,
                                            [cumulants_ph, rewards_ph])
    self._action = session.make_callable(self._regressed_w.read_value(), [])
    session.run(tf.global_variables_initializer())
//  [python]
def input_fn_builder(input_file, is_training, drop_remainder,
                     names_to_features):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

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

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

    return example

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

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

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

    return d

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

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

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

    return example

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

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

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

    return d

  return input_fn
//  [python]
def padded_rmse(predictions, labels, weights_fn=common_layers.weights_all):
  predictions = tf.to_float(predictions)
  labels = tf.to_float(labels)
  predictions, labels = common_layers.pad_with_zeros(predictions, labels)
  weights = weights_fn(labels)
  error = tf.pow(predictions - labels, 2)
  error_sqrt = tf.sqrt(tf.reduce_mean(error * weights))
  return error_sqrt, tf.reduce_sum(weights)
//  [python]
def padded_rmse(predictions, labels, weights_fn=common_layers.weights_all):
  predictions = tf.cast(predictions, tf.float32)
  labels = tf.cast(labels, tf.float32)
  predictions, labels = common_layers.pad_with_zeros(predictions, labels)
  weights = weights_fn(labels)
  error = tf.pow(predictions - labels, 2)
  error_sqrt = tf.sqrt(tf.reduce_mean(error * weights))
  return error_sqrt, tf.reduce_sum(weights)
//  [python]
def _using_motion_vector_with_distortion(depth,
                                         translation,
                                         rotation_angles,
                                         intrinsic_mat,
                                         distortion_coeff=0.0):
  """A helper for using_motion_vector. See docstring therein."""

  if translation.shape.ndims not in (2, 4):
    raise ValueError('\'translation\' should have rank 2 or 4, not %d' %
                     translation.shape.ndims)
  if translation.shape[-1] != 3:
    raise ValueError('translation\'s last dimension should be 3, not %d' %
                     translation.shape[1])
  if translation.shape.ndims == 2:
    translation = tf.expand_dims(tf.expand_dims(translation, 1), 1)

  _, height, width = tf.unstack(tf.shape(depth))
  grid = tf.squeeze(
      tf.stack(tf.meshgrid(tf.range(width), tf.range(height), (1,))),
      axis=3)  # 3 x height x width
  grid = tf.to_float(grid)
  intrinsic_mat_inv = tf.linalg.inv(intrinsic_mat)

  normalized_grid = tf.einsum('bij,jhw->bihw', intrinsic_mat_inv, grid)
  radii_squared = tf.reduce_sum(tf.square(normalized_grid[:, :2, :, :]), axis=1)

  undistortion_factor = quadratic_inverse_distortion_scale(
      distortion_coeff, radii_squared)
  undistortion_factor = tf.stack([
      undistortion_factor, undistortion_factor,
      tf.ones_like(undistortion_factor)
  ],
                                 axis=1)
  normalized_grid *= undistortion_factor

  rot_mat = transform_utils.matrix_from_angles(rotation_angles)
  # We have to treat separately the case of a per-image rotation vector and a
  # per-image rotation field, because the broadcasting capabilities of einsum
  # are limited.
  if rotation_angles.shape.ndims == 2:
    # The calculation here is identical to the one in inverse_warp above.
    # Howeverwe use einsum for better clarity. Under the hood, einsum performs
    # the reshaping and invocation of BatchMatMul, instead of doing it manually,
    # as in inverse_warp.
    pcoords = tf.einsum('bij,bjhw,bhw->bihw', rot_mat, normalized_grid, depth)
  elif rotation_angles.shape.ndims == 4:
    # We push the H and W dimensions to the end, and transpose the rotation
    # matrix elements (as noted above).
    rot_mat = tf.transpose(rot_mat, [0, 3, 4, 1, 2])
    pcoords = tf.einsum('bijhw,bjhw,bhw->bihw', rot_mat, normalized_grid, depth)

  pcoords += tf.transpose(translation, [0, 3, 1, 2])

  x, y, z = tf.unstack(pcoords, axis=1)
  x /= z
  y /= z
  scale = quadraric_distortion_scale(distortion_coeff,
                                     tf.square(x) + tf.square(y))
  x *= scale
  y *= scale

  pcoords = tf.einsum('bij,bjhw->bihw', intrinsic_mat,
                      tf.stack([x, y, tf.ones_like(x)], axis=1))
  x, y, _ = tf.unstack(pcoords, axis=1)

  return x, y, z
//  [python]
def _using_motion_vector_with_distortion(depth,
                                         translation,
                                         rotation_angles,
                                         intrinsic_mat,
                                         distortion_coeff=0.0):
  """A helper for using_motion_vector. See docstring therein."""

  if translation.shape.ndims not in (2, 4):
    raise ValueError('\'translation\' should have rank 2 or 4, not %d' %
                     translation.shape.ndims)
  if translation.shape[-1] != 3:
    raise ValueError('translation\'s last dimension should be 3, not %d' %
                     translation.shape[1])
  if translation.shape.ndims == 2:
    translation = tf.expand_dims(tf.expand_dims(translation, 1), 1)

  _, height, width = tf.unstack(tf.shape(depth))
  grid = tf.squeeze(
      tf.stack(tf.meshgrid(tf.range(width), tf.range(height), (1,))),
      axis=3)  # 3 x height x width
  grid = tf.cast(grid, dtype=tf.float32)
  intrinsic_mat_inv = tf.linalg.inv(intrinsic_mat)

  normalized_grid = tf.einsum('bij,jhw->bihw', intrinsic_mat_inv, grid)
  radii_squared = tf.reduce_sum(tf.square(normalized_grid[:, :2, :, :]), axis=1)

  undistortion_factor = quadratic_inverse_distortion_scale(
      distortion_coeff, radii_squared)
  undistortion_factor = tf.stack([
      undistortion_factor, undistortion_factor,
      tf.ones_like(undistortion_factor)
  ],
                                 axis=1)
  normalized_grid *= undistortion_factor

  rot_mat = transform_utils.matrix_from_angles(rotation_angles)
  # We have to treat separately the case of a per-image rotation vector and a
  # per-image rotation field, because the broadcasting capabilities of einsum
  # are limited.
  if rotation_angles.shape.ndims == 2:
    # The calculation here is identical to the one in inverse_warp above.
    # Howeverwe use einsum for better clarity. Under the hood, einsum performs
    # the reshaping and invocation of BatchMatMul, instead of doing it manually,
    # as in inverse_warp.
    pcoords = tf.einsum('bij,bjhw,bhw->bihw', rot_mat, normalized_grid, depth)
  elif rotation_angles.shape.ndims == 4:
    # We push the H and W dimensions to the end, and transpose the rotation
    # matrix elements (as noted above).
    rot_mat = tf.transpose(rot_mat, [0, 3, 4, 1, 2])
    pcoords = tf.einsum('bijhw,bjhw,bhw->bihw', rot_mat, normalized_grid, depth)

  pcoords += tf.transpose(translation, [0, 3, 1, 2])

  x, y, z = tf.unstack(pcoords, axis=1)
  x /= z
  y /= z
  scale = quadraric_distortion_scale(distortion_coeff,
                                     tf.square(x) + tf.square(y))
  x *= scale
  y *= scale

  pcoords = tf.einsum('bij,bjhw->bihw', intrinsic_mat,
                      tf.stack([x, y, tf.ones_like(x)], axis=1))
  x, y, _ = tf.unstack(pcoords, axis=1)

  return x, y, z
//  [python]
def unpadded_mse(predictions, labels, weights_fn=common_layers.weights_all):
  predictions = tf.to_float(predictions)
  labels = tf.to_float(labels)
  weights = weights_fn(labels)
  error = tf.pow(predictions - labels, 2)
  mean_error = tf.reduce_mean(error * weights)
  return mean_error, tf.reduce_sum(weights)
//  [python]
def unpadded_mse(predictions, labels, weights_fn=common_layers.weights_all):
  predictions = tf.cast(predictions, tf.float32)
  labels = tf.cast(labels, tf.float32)
  weights = weights_fn(labels)
  error = tf.pow(predictions - labels, 2)
  mean_error = tf.reduce_mean(error * weights)
  return mean_error, tf.reduce_sum(weights)
//  [python]
def triplet_loss(scope, imgs, is_train, n, margin, batch_all=True):
    anchor_img, pos_img, neg_img = tf.split(imgs, [1, 1, 1], 3, name="split")

    anchor, _ = inception.inception_v3(anchor_img, is_training=is_train, scope=scope, num_classes=n)
    tf.get_variable_scope().reuse_variables()
    positive, _ = inception.inception_v3(pos_img, is_training=is_train, scope=scope, num_classes=n)
    tf.get_variable_scope().reuse_variables()
    negative, _ = inception.inception_v3(neg_img, is_training=is_train, scope=scope, num_classes=n)

    # distance between the anchor and the positive
    pos_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, positive)), 1)
    # distance between the anchor and the negative
    neg_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, negative)), 1)

    if batch_all:
        # batch all
        t_loss = pos_dist - neg_dist + margin
        t_loss = tf.maximum(t_loss, 0.0)
        valid_triplets = tf.to_float(tf.greater(t_loss, 1e-16))
        num_positive_triplets = tf.reduce_sum(valid_triplets)
        # Get final mean triplet loss over the positive valid triplets
        t_loss = tf.reduce_sum(t_loss) / (num_positive_triplets + 1e-16)
    else:
        # batch hard
        hardest_positive_dist = tf.reduce_max(pos_dist, axis=0, keepdims=True)
        # For each anchor, get the hardest negative
        hardest_negative_dist = tf.reduce_min(neg_dist, axis=0, keepdims=True)
        t_loss = hardest_positive_dist - hardest_negative_dist + margin
        t_loss = tf.maximum(t_loss, 0.0)

    return pos_dist, neg_dist, t_loss
//  [python]
def triplet_loss(scope, imgs, is_train, n, margin, batch_all=True):
    anchor_img, pos_img, neg_img = tf.split(imgs, [1, 1, 1], 3, name="split")

    anchor, _ = inception.inception_v3(anchor_img, is_training=is_train, scope=scope, num_classes=n)
    tf.get_variable_scope().reuse_variables()
    positive, _ = inception.inception_v3(pos_img, is_training=is_train, scope=scope, num_classes=n)
    tf.get_variable_scope().reuse_variables()
    negative, _ = inception.inception_v3(neg_img, is_training=is_train, scope=scope, num_classes=n)

    # distance between the anchor and the positive
    pos_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, positive)), 1)
    # distance between the anchor and the negative
    neg_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, negative)), 1)

    if batch_all:
        # batch all
        t_loss = pos_dist - neg_dist + margin
        t_loss = tf.maximum(t_loss, 0.0)
        valid_triplets = tf.cast(tf.greater(t_loss, 1e-16), tf.float32)
        num_positive_triplets = tf.reduce_sum(valid_triplets)
        # Get final mean triplet loss over the positive valid triplets
        t_loss = tf.reduce_sum(t_loss) / (num_positive_triplets + 1e-16)
    else:
        # batch hard
        hardest_positive_dist = tf.reduce_max(pos_dist, axis=0, keepdims=True)
        # For each anchor, get the hardest negative
        hardest_negative_dist = tf.reduce_min(neg_dist, axis=0, keepdims=True)
        t_loss = hardest_positive_dist - hardest_negative_dist + margin
        t_loss = tf.maximum(t_loss, 0.0)

    return pos_dist, neg_dist, t_loss
//  [python]
def ctc_decode(y_pred, input_length, greedy=True, beam_width=100,
               top_paths=1):
    """Decodes the output of a softmax.

    Can use either greedy search (also known as best path)
    or a constrained dictionary search.

    # Arguments
        y_pred: tensor `(samples, time_steps, num_categories)`
            containing the prediction, or output of the softmax.
        input_length: tensor `(samples, )` containing the sequence length for
            each batch item in `y_pred`.
        greedy: perform much faster best-path search if `true`.
            This does not use a dictionary.
        beam_width: if `greedy` is `false`: a beam search decoder will be used
            with a beam of this width.
        top_paths: if `greedy` is `false`,
            how many of the most probable paths will be returned.

    # Returns
        Tuple:
            List: if `greedy` is `true`, returns a list of one element that
                contains the decoded sequence.
                If `false`, returns the `top_paths` most probable
                decoded sequences.
                Important: blank labels are returned as `-1`.
            Tensor `(top_paths, )` that contains
                the log probability of each decoded sequence.
    """
    y_pred = tf.log(tf.transpose(y_pred, perm=[1, 0, 2]) + epsilon())
    input_length = tf.to_int32(input_length)

    if greedy:
        (decoded, log_prob) = ctc.ctc_greedy_decoder(
            inputs=y_pred,
            sequence_length=input_length)
    else:
        (decoded, log_prob) = ctc.ctc_beam_search_decoder(
            inputs=y_pred,
            sequence_length=input_length, beam_width=beam_width,
            top_paths=top_paths)

    decoded_dense = [tf.sparse_to_dense(st.indices, st.dense_shape, st.values, default_value=-1)
                     for st in decoded]
    return (decoded_dense, log_prob)
//  [python]
def ctc_decode(y_pred, input_length, greedy=True, beam_width=100,
               top_paths=1):
    """Decodes the output of a softmax.

    Can use either greedy search (also known as best path)
    or a constrained dictionary search.

    # Arguments
        y_pred: tensor `(samples, time_steps, num_categories)`
            containing the prediction, or output of the softmax.
        input_length: tensor `(samples, )` containing the sequence length for
            each batch item in `y_pred`.
        greedy: perform much faster best-path search if `true`.
            This does not use a dictionary.
        beam_width: if `greedy` is `false`: a beam search decoder will be used
            with a beam of this width.
        top_paths: if `greedy` is `false`,
            how many of the most probable paths will be returned.

    # Returns
        Tuple:
            List: if `greedy` is `true`, returns a list of one element that
                contains the decoded sequence.
                If `false`, returns the `top_paths` most probable
                decoded sequences.
                Important: blank labels are returned as `-1`.
            Tensor `(top_paths, )` that contains
                the log probability of each decoded sequence.
    """
    y_pred = tf.log(tf.transpose(y_pred, perm=[1, 0, 2]) + epsilon())
    input_length = tf.cast(input_length, tf.int32)

    if greedy:
        (decoded, log_prob) = ctc.ctc_greedy_decoder(
            inputs=y_pred,
            sequence_length=input_length)
    else:
        (decoded, log_prob) = ctc.ctc_beam_search_decoder(
            inputs=y_pred,
            sequence_length=input_length, beam_width=beam_width,
            top_paths=top_paths)

    decoded_dense = [tf.sparse_to_dense(st.indices, st.dense_shape, st.values, default_value=-1)
                     for st in decoded]
    return (decoded_dense, log_prob)
//  [python]
def padded_accuracy_topk(predictions,
                         labels,
                         k,
                         weights_fn=common_layers.weights_nonzero):
  """Percentage of times that top-k predictions matches labels on non-0s."""
  with tf.variable_scope("padded_accuracy_topk", values=[predictions, labels]):
    padded_predictions, padded_labels = common_layers.pad_with_zeros(
        predictions, labels)
    weights = weights_fn(padded_labels)
    effective_k = tf.minimum(k,
                             common_layers.shape_list(padded_predictions)[-1])
    _, outputs = tf.nn.top_k(padded_predictions, k=effective_k)
    outputs = tf.to_int32(outputs)
    padded_labels = tf.to_int32(padded_labels)
    padded_labels = tf.expand_dims(padded_labels, axis=-1)
    padded_labels += tf.zeros_like(outputs)  # Pad to same shape.
    same = tf.to_float(tf.equal(outputs, padded_labels))
    same_topk = tf.reduce_sum(same, axis=-1)
    return same_topk, weights
//  [python]
def padded_accuracy_topk(predictions,
                         labels,
                         k,
                         weights_fn=common_layers.weights_nonzero):
  """Percentage of times that top-k predictions matches labels on non-0s."""
  with tf.variable_scope("padded_accuracy_topk", values=[predictions, labels]):
    padded_predictions, padded_labels = common_layers.pad_with_zeros(
        predictions, labels)
    weights = weights_fn(padded_labels)
    effective_k = tf.minimum(k,
                             common_layers.shape_list(padded_predictions)[-1])
    _, outputs = tf.nn.top_k(padded_predictions, k=effective_k)
    outputs = tf.cast(outputs, tf.int32)
    padded_labels = tf.cast(padded_labels, tf.int32)
    padded_labels = tf.expand_dims(padded_labels, axis=-1)
    padded_labels += tf.zeros_like(outputs)  # Pad to same shape.
    same = tf.to_float(tf.equal(outputs, padded_labels))
    same_topk = tf.reduce_sum(same, axis=-1)
    return same_topk, weights
//  [python]
def add_intrinsics_head(bottleneck, image_height, image_width):
    """Adds a head the preficts camera intrinsics.

  Args:
    bottleneck: A tf.Tensor of shape [B, 1, 1, C], typically the bottlenech
      features of a netrowk.
    image_height: A scalar tf.Tensor or an python scalar, the image height in
      pixels.
    image_width: A scalar tf.Tensor or an python scalar, the image width in
      pixels.

  image_height and image_width are used to provide the right scale for the focal
  length and the offest parameters.

  Returns:
    a tf.Tensor of shape [B, 3, 3], and type float32, where the 3x3 part is the
    intrinsic matrix: (fx, 0, x0), (0, fy, y0), (0, 0, 1).
  """
    with tf.variable_scope('CameraIntrinsics'):
        # Since the focal lengths in pixels tend to be in the order of magnitude of
        # the image width and height, we multiply the network prediction by them.
        focal_lengths = tf.squeeze(
            layers.conv2d(bottleneck, 2, [1, 1], stride=1, activation_fn=tf.nn.softplus, weights_regularizer=None, scope='foci'),
            axis=(1, 2)
        ) * tf.to_float(tf.convert_to_tensor([[image_width, image_height]]))

        # The pixel offsets tend to be around the center of the image, and they
        # are typically a fraction the image width and height in pixels. We thus
        # multiply the network prediction by the width and height, and the
        # additional 0.5 them by default at the center of the image.
        offsets = (tf.squeeze(
            layers.conv2d(bottleneck, 2, [1, 1], stride=1, activation_fn=None, weights_regularizer=None, biases_initializer=None, scope='offsets'),
            axis=(1, 2)) + 0.5
        ) * tf.to_float(tf.convert_to_tensor([[image_width, image_height]]))

        foci = tf.linalg.diag(focal_lengths)

        intrinsic_mat = tf.concat([foci, tf.expand_dims(offsets, -1)], axis=2)
        batch_size = tf.shape(bottleneck)[0]
        last_row = tf.tile([[[0.0, 0.0, 1.0]]], [batch_size, 1, 1])
        intrinsic_mat = tf.concat([intrinsic_mat, last_row], axis=1)
        return intrinsic_mat
//  [python]
def add_intrinsics_head(bottleneck, image_height, image_width):
    """Adds a head the preficts camera intrinsics.

  Args:
    bottleneck: A tf.Tensor of shape [B, 1, 1, C], typically the bottlenech
      features of a netrowk.
    image_height: A scalar tf.Tensor or an python scalar, the image height in
      pixels.
    image_width: A scalar tf.Tensor or an python scalar, the image width in
      pixels.

  image_height and image_width are used to provide the right scale for the focal
  length and the offest parameters.

  Returns:
    a tf.Tensor of shape [B, 3, 3], and type float32, where the 3x3 part is the
    intrinsic matrix: (fx, 0, x0), (0, fy, y0), (0, 0, 1).
  """
    with tf.variable_scope('CameraIntrinsics'):
        # Since the focal lengths in pixels tend to be in the order of magnitude of
        # the image width and height, we multiply the network prediction by them.
        focal_lengths = tf.squeeze(
            layers.conv2d(bottleneck, 2, [1, 1], stride=1, activation_fn=tf.nn.softplus, weights_regularizer=None, scope='foci'),
            axis=(1, 2)
        ) * tf.cast(tf.convert_to_tensor([[image_width, image_height]]), tf.float32)

        # The pixel offsets tend to be around the center of the image, and they
        # are typically a fraction the image width and height in pixels. We thus
        # multiply the network prediction by the width and height, and the
        # additional 0.5 them by default at the center of the image.
        offsets = (tf.squeeze(
            layers.conv2d(bottleneck, 2, [1, 1], stride=1, activation_fn=None, weights_regularizer=None, biases_initializer=None, scope='offsets'),
            axis=(1, 2)) + 0.5
        ) * tf.cast(tf.convert_to_tensor([[image_width, image_height]]), tf.float32)

        foci = tf.linalg.diag(focal_lengths)

        intrinsic_mat = tf.concat([foci, tf.expand_dims(offsets, -1)], axis=2)
        batch_size = tf.shape(bottleneck)[0]
        last_row = tf.tile([[[0.0, 0.0, 1.0]]], [batch_size, 1, 1])
        intrinsic_mat = tf.concat([intrinsic_mat, last_row], axis=1)
        return intrinsic_mat
//  [python]
def two_class_accuracy(predictions, labels, weights_fn=None):
  """Accuracy for two class classification with 0/1 labels."""
  with tf.variable_scope("two_class_accuracy", values=[predictions, labels]):
    del weights_fn
    hard_predictions = tf.to_int32(tf.math.round(tf.squeeze(predictions)))
    int_labels = tf.to_int32(labels)
    _, accuracy = tf.metrics.accuracy(labels=int_labels,
                                      predictions=hard_predictions)
    return accuracy, tf.constant(1.0)
//  [python]
def two_class_accuracy(predictions, labels, weights_fn=None):
  """Accuracy for two class classification with 0/1 labels."""
  with tf.variable_scope("two_class_accuracy", values=[predictions, labels]):
    del weights_fn
    hard_predictions = tf.cast(tf.math.round(tf.squeeze(predictions)), tf.int32)
    int_labels = tf.cast(labels, tf.int32)
    _, accuracy = tf.metrics.accuracy(labels=int_labels,
                                      predictions=hard_predictions)
    return accuracy, tf.constant(1.0)
//  [python]
def rounding_sequence_accuracy(predictions,
                               labels,
                               weights_fn=common_layers.weights_nonzero):
  """Sequence accuracy for L1/L2 losses: round down the predictions to ints."""
  outputs = tf.squeeze(tf.to_int32(predictions), axis=-1)
  weights = weights_fn(labels)
  labels = tf.to_int32(labels)
  not_correct = tf.to_float(tf.not_equal(outputs, labels)) * weights
  axis = list(range(1, len(outputs.get_shape())))
  correct_seq = 1.0 - tf.minimum(1.0, tf.reduce_sum(not_correct, axis=axis))
  return correct_seq, tf.constant(1.0)
//  [python]
def rounding_sequence_accuracy(predictions,
                               labels,
                               weights_fn=common_layers.weights_nonzero):
  """Sequence accuracy for L1/L2 losses: round down the predictions to ints."""
  outputs = tf.squeeze(tf.cast(predictions, tf.int32), axis=-1)
  weights = weights_fn(labels)
  labels = tf.cast(labels, tf.int32)
  not_correct = tf.cast(tf.not_equal(outputs, labels), tf.float32) * weights
  axis = list(range(1, len(outputs.get_shape())))
  correct_seq = 1.0 - tf.minimum(1.0, tf.reduce_sum(not_correct, axis=axis))
  return correct_seq, tf.constant(1.0)
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder, task_name, use_tpu, bsz,
                                multiple=1,slist=None):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""
  labeltype = tf.float32 if task_name == "sts-b" else tf.int64
  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "label_ids": tf.FixedLenFeature([], labeltype),
      "is_real_example": tf.FixedLenFeature([], tf.int64),
  }

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

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

    return example

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

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if not slist==None:
      sublist = tf.data.Dataset.from_tensor_slices(tf.cast(slist, tf.int64))
      d=tf.data.experimental.choose_from_datasets(d, sublist)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        contrib_data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder, task_name, use_tpu, bsz,
                                multiple=1,slist=None):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""
  labeltype = tf.float32 if task_name == "sts-b" else tf.int64
  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "label_ids": tf.FixedLenFeature([], labeltype),
      "is_real_example": tf.FixedLenFeature([], tf.int64),
  }

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

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

    return example

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

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if not slist==None:
      sublist = tf.data.Dataset.from_tensor_slices(tf.cast(slist, tf.int64))
      d=tf.data.experimental.choose_from_datasets(d, sublist)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        contrib_data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def rounding_accuracy(predictions,
                      labels,
                      weights_fn=common_layers.weights_nonzero):
  """Rounding accuracy for L1/L2 losses: round down the predictions to ints."""
  outputs = tf.squeeze(tf.to_int32(predictions))
  labels = tf.squeeze(labels)
  weights = weights_fn(labels)
  labels = tf.to_int32(labels)
  return tf.to_float(tf.equal(outputs, labels)), weights
//  [python]
def rounding_accuracy(predictions,
                      labels,
                      weights_fn=common_layers.weights_nonzero):
  """Rounding accuracy for L1/L2 losses: round down the predictions to ints."""
  outputs = tf.squeeze(tf.cast(predictions, tf.int32))
  labels = tf.squeeze(labels)
  weights = weights_fn(labels)
  labels = tf.cast(labels, tf.int32)
  return tf.cast(tf.equal(outputs, labels), tf.float32), weights
//  [python]
def padded_accuracy(predictions,
                    labels,
                    weights_fn=common_layers.weights_nonzero):
  """Percentage of times that predictions matches labels on non-0s."""
  # If the last dimension is 1 then we're using L1/L2 loss.
  if common_layers.shape_list(predictions)[-1] == 1:
    return rounding_accuracy(predictions, labels, weights_fn=weights_fn)
  with tf.variable_scope("padded_accuracy", values=[predictions, labels]):
    padded_predictions, padded_labels = common_layers.pad_with_zeros(
        predictions, labels)
    weights = weights_fn(padded_labels)
    outputs = tf.to_int32(tf.argmax(padded_predictions, axis=-1))
    padded_labels = tf.to_int32(padded_labels)
    return tf.to_float(tf.equal(outputs, padded_labels)), weights
//  [python]
def padded_accuracy(predictions,
                    labels,
                    weights_fn=common_layers.weights_nonzero):
  """Percentage of times that predictions matches labels on non-0s."""
  # If the last dimension is 1 then we're using L1/L2 loss.
  if common_layers.shape_list(predictions)[-1] == 1:
    return rounding_accuracy(predictions, labels, weights_fn=weights_fn)
  with tf.variable_scope("padded_accuracy", values=[predictions, labels]):
    padded_predictions, padded_labels = common_layers.pad_with_zeros(
        predictions, labels)
    weights = weights_fn(padded_labels)
    outputs = tf.cast(tf.argmax(padded_predictions, axis=-1), tf.int32)
    padded_labels = tf.cast(padded_labels, tf.int32)
    return tf.cast(tf.equal(outputs, padded_labels), tf.float32), weights
//  [python]
def padded_sequence_accuracy(predictions,
                             labels,
                             weights_fn=common_layers.weights_nonzero):
  """Percentage of times that predictions matches labels everywhere (non-0)."""
  # If the last dimension is 1 then we're using L1/L2 loss.
  if common_layers.shape_list(predictions)[-1] == 1:
    return rounding_sequence_accuracy(
        predictions, labels, weights_fn=weights_fn)
  with tf.variable_scope(
      "padded_sequence_accuracy", values=[predictions, labels]):
    padded_predictions, padded_labels = common_layers.pad_with_zeros(
        predictions, labels)
    weights = weights_fn(padded_labels)

    # Flatten, keeping batch dim (and num_classes dim for predictions)
    # TPU argmax can only deal with a limited number of dimensions
    predictions_shape = common_layers.shape_list(padded_predictions)
    batch_size = predictions_shape[0]
    num_classes = predictions_shape[-1]
    flat_size = common_layers.list_product(
        common_layers.shape_list(padded_labels)[1:])
    padded_predictions = tf.reshape(
        padded_predictions,
        [batch_size, common_layers.list_product(predictions_shape[1:-1]),
         num_classes])
    padded_labels = tf.reshape(padded_labels, [batch_size, flat_size])
    weights = tf.reshape(weights, [batch_size, flat_size])

    outputs = tf.to_int32(tf.argmax(padded_predictions, axis=-1))
    padded_labels = tf.to_int32(padded_labels)
    not_correct = tf.to_float(tf.not_equal(outputs, padded_labels)) * weights
    axis = list(range(1, len(outputs.get_shape())))
    correct_seq = 1.0 - tf.minimum(1.0, tf.reduce_sum(not_correct, axis=axis))
    return correct_seq, tf.constant(1.0)
//  [python]
def padded_sequence_accuracy(predictions,
                             labels,
                             weights_fn=common_layers.weights_nonzero):
  """Percentage of times that predictions matches labels everywhere (non-0)."""
  # If the last dimension is 1 then we're using L1/L2 loss.
  if common_layers.shape_list(predictions)[-1] == 1:
    return rounding_sequence_accuracy(
        predictions, labels, weights_fn=weights_fn)
  with tf.variable_scope(
      "padded_sequence_accuracy", values=[predictions, labels]):
    padded_predictions, padded_labels = common_layers.pad_with_zeros(
        predictions, labels)
    weights = weights_fn(padded_labels)

    # Flatten, keeping batch dim (and num_classes dim for predictions)
    # TPU argmax can only deal with a limited number of dimensions
    predictions_shape = common_layers.shape_list(padded_predictions)
    batch_size = predictions_shape[0]
    num_classes = predictions_shape[-1]
    flat_size = common_layers.list_product(
        common_layers.shape_list(padded_labels)[1:])
    padded_predictions = tf.reshape(
        padded_predictions,
        [batch_size, common_layers.list_product(predictions_shape[1:-1]),
         num_classes])
    padded_labels = tf.reshape(padded_labels, [batch_size, flat_size])
    weights = tf.reshape(weights, [batch_size, flat_size])

    outputs = tf.cast(tf.argmax(padded_predictions, axis=-1), tf.int32)
    padded_labels = tf.cast(padded_labels, tf.int32)
    not_correct = tf.to_float(tf.not_equal(outputs, padded_labels)) * weights
    axis = list(range(1, len(outputs.get_shape())))
    correct_seq = 1.0 - tf.minimum(1.0, tf.reduce_sum(not_correct, axis=axis))
    return correct_seq, tf.constant(1.0)
//  [python]
def prefix_accuracy(predictions,
                    labels,
                    weights_fn=common_layers.weights_nonzero):
  """Average # of correct tokens at start of sequences, ignoring padding 0s.

  See section 4.3 of Learning to Transduce with Unbounded Memory,
  Grefenstette et al., 2015.

  Args:
    predictions: Tensor of shape [`batch_size`, `length`, 1, `num_classes`] and
        type tf.float32 representing the logits, 0-padded.
    labels: Tensor of shape [`batch_size`, `length`, 1, 1] and type tf.int32
        representing the labels of same length as logits and 0-padded.
    weights_fn: ignored. The weights returned are the total length of the ground
        truth labels, excluding 0-paddings.

  Returns:
    (prefix accuracy, 1.0)

  Raises:
    ValueError: if weights_fn is not common_layers.weights_nonzero.
  """
  if weights_fn is not common_layers.weights_nonzero:
    raise ValueError("Only weights_nonzero can be used for this metric.")

  predictions = tf.to_int32(tf.squeeze(tf.argmax(predictions, axis=-1), axis=2))
  labels = tf.squeeze(labels, axis=(2, 3))
  seq_len = tf.reduce_sum(
      tf.cast(tf.not_equal(labels, tf.constant(0)), dtype=tf.float32), axis=1)
  matching_elements = tf.equal(labels, predictions)
  prefix_len = tf.reduce_sum(
      tf.cumprod(tf.cast(matching_elements, tf.float32), axis=1), axis=1)
  return tf.reduce_mean(prefix_len / seq_len), tf.constant(1.0)
//  [python]
def prefix_accuracy(predictions,
                    labels,
                    weights_fn=common_layers.weights_nonzero):
  """Average # of correct tokens at start of sequences, ignoring padding 0s.

  See section 4.3 of Learning to Transduce with Unbounded Memory,
  Grefenstette et al., 2015.

  Args:
    predictions: Tensor of shape [`batch_size`, `length`, 1, `num_classes`] and
        type tf.float32 representing the logits, 0-padded.
    labels: Tensor of shape [`batch_size`, `length`, 1, 1] and type tf.int32
        representing the labels of same length as logits and 0-padded.
    weights_fn: ignored. The weights returned are the total length of the ground
        truth labels, excluding 0-paddings.

  Returns:
    (prefix accuracy, 1.0)

  Raises:
    ValueError: if weights_fn is not common_layers.weights_nonzero.
  """
  if weights_fn is not common_layers.weights_nonzero:
    raise ValueError("Only weights_nonzero can be used for this metric.")

  predictions = tf.cast(tf.squeeze(tf.argmax(predictions, axis=-1), axis=2), tf.int32)
  labels = tf.squeeze(labels, axis=(2, 3))
  seq_len = tf.reduce_sum(
      tf.cast(tf.not_equal(labels, tf.constant(0)), dtype=tf.float32), axis=1)
  matching_elements = tf.equal(labels, predictions)
  prefix_len = tf.reduce_sum(
      tf.cumprod(tf.cast(matching_elements, tf.float32), axis=1), axis=1)
  return tf.reduce_mean(prefix_len / seq_len), tf.constant(1.0)
//  [python]
def set_precision(predictions, labels,
                  weights_fn=common_layers.weights_nonzero):
  """Precision of set predictions.

  Args:
    predictions : A Tensor of scores of shape [batch, nlabels].
    labels: A Tensor of int32s giving true set elements,
      of shape [batch, seq_length].
    weights_fn: A function to weight the elements.

  Returns:
    hits: A Tensor of shape [batch, nlabels].
    weights: A Tensor of shape [batch, nlabels].
  """
  with tf.variable_scope("set_precision", values=[predictions, labels]):
    labels = tf.squeeze(labels, [2, 3])
    weights = weights_fn(labels)
    labels = tf.one_hot(labels, predictions.shape[-1])
    labels = tf.reduce_max(labels, axis=1)
    labels = tf.cast(labels, tf.bool)
    return tf.to_float(tf.equal(labels, predictions)), weights
//  [python]
def set_precision(predictions, labels,
                  weights_fn=common_layers.weights_nonzero):
  """Precision of set predictions.

  Args:
    predictions : A Tensor of scores of shape [batch, nlabels].
    labels: A Tensor of int32s giving true set elements,
      of shape [batch, seq_length].
    weights_fn: A function to weight the elements.

  Returns:
    hits: A Tensor of shape [batch, nlabels].
    weights: A Tensor of shape [batch, nlabels].
  """
  with tf.variable_scope("set_precision", values=[predictions, labels]):
    labels = tf.squeeze(labels, [2, 3])
    weights = weights_fn(labels)
    labels = tf.one_hot(labels, predictions.shape[-1])
    labels = tf.reduce_max(labels, axis=1)
    labels = tf.cast(labels, tf.bool)
    return tf.cast(tf.equal(labels, predictions), tf.float32), weights
//  [python]
def oov_metric(self, word_ids):
    """Metric for checking the OOV rate.

    Args:
      word_ids: int32 Tensor of arbitrary shape with word indices.

    Returns:
      metric: A metric computing the OOV rate. For use with tf.Estimator.
    """
    return tf.metrics.mean(tf.to_float(tf.greater(
        word_ids, len(self._idx2str) - 1)))