jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def testPerArmRewards(self, policy_class):
    tf.compat.v1.set_random_seed(3000)
    obs_spec = bandit_spec_utils.create_per_arm_observation_spec(2, 3, 4)
    time_step_spec = ts.time_step_spec(obs_spec)
    action_spec = tensor_spec.BoundedTensorSpec((), tf.int32, 0, 3)
    reward_network = (
        global_and_arm_feature_network.create_feed_forward_common_tower_network(
            obs_spec, (4, 3), (3, 4), (4, 2)
        )
    )

    policy = policy_class(
        time_step_spec,
        action_spec,
        reward_network=reward_network,
        accepts_per_arm_features=True,
        emit_policy_info=('predicted_rewards_mean',),
    )
    action_feature = tf.cast(
        tf.reshape(tf.random.shuffle(tf.range(24)), shape=[2, 4, 3]),
        dtype=tf.float32,
    )
    observations = {
        bandit_spec_utils.GLOBAL_FEATURE_KEY: tf.constant(
            [[1, 2], [3, 4]], dtype=tf.float32
        ),
        bandit_spec_utils.PER_ARM_FEATURE_KEY: action_feature,
    }
    time_step = ts.restart(observations, batch_size=2)
    action_step = policy.action(time_step, seed=1)
    self.assertEqual(action_step.action.shape.as_list(), [2])
    self.assertEqual(action_step.action.dtype, tf.int32)
    # Initialize all variables
    self.evaluate(tf.compat.v1.global_variables_initializer())
    action, p_info, first_arm_features = self.evaluate([
        action_step.action,
        action_step.info,
        observations[bandit_spec_utils.PER_ARM_FEATURE_KEY][0],
    ])
    self.assertAllEqual(action.shape, [2])
    self.assertAllEqual(p_info.predicted_rewards_mean.shape, [2, 4])
    self.assertAllEqual(p_info.chosen_arm_features.shape, [2, 3])
    first_action = action[0]
    self.assertAllEqual(
        p_info.chosen_arm_features[0], first_arm_features[first_action]
    )

    # Check that zeroing out some of the actions does not affect the predicted
    # rewards for unchanged actions. This is to make sure that action feature
    # padding does not influence the behavior.

    if not tf.executing_eagerly():
      # The below comparison will only work in tf2 because of the random per-arm
      # observations get re-drawn in tf1.
      return
    padded_action_feature = tf.concat(
        [
            action_feature[:, 0:1, :],
            tf.zeros(shape=[2, 3, 3], dtype=tf.float32),
        ],
        axis=1,
    )
    observations = {
        bandit_spec_utils.GLOBAL_FEATURE_KEY: tf.constant(
            [[1, 2], [3, 4]], dtype=tf.float32
        ),
        bandit_spec_utils.PER_ARM_FEATURE_KEY: padded_action_feature,
    }
    time_step = ts.restart(observations, batch_size=2)
    padded_action_step = policy.action(time_step, seed=1)
    padded_p_info = self.evaluate(padded_action_step.info)
    self.assertAllEqual(
        p_info.predicted_rewards_mean[:, 0],
        padded_p_info.predicted_rewards_mean[:, 0],
    )
//  [python]
def generate_sampled_gt_matrix(gt_matrix, samples, predicted_flags, train_flag,
                               config):

    if train_flag:
        samples_per_vertex = config.SAMPLES_PER_VERTEX
        max_vertices = config.TRAIN_ROIS_PER_IMAGE
        num_batch = config.BATCH_SIZE
    else:
        samples_per_vertex = config.DETECTION_MAX_INSTANCES
        max_vertices = config.DETECTION_MAX_INSTANCES
        num_batch = config.BATCH_SIZE

    mask = mask = tf.tile(tf.expand_dims(predicted_flags, -1),
                          [1, 1, samples_per_vertex])
    samples = t.cast(samples, t.int32)
    samples = tf.where(
        tf.math.equal(samples, tf.constant(max_vertices, tf.int32)),
        tf.zeros_like(samples), samples)
    y = tf.range(0, num_batch)[..., tf.newaxis]
    y = tf.tile(y, multiples=[1, max_vertices])
    z = tf.range(0, max_vertices)[tf.newaxis, ..., tf.newaxis]
    z = tf.tile(z, multiples=[num_batch, 1, 1])
    batch_range_1 = tf.tile(y[..., tf.newaxis],
                            multiples=[1, 1, samples_per_vertex])
    max_vertices_range_1 = tf.tile(z, multiples=[1, 1, samples_per_vertex])
    indexing_tensor_for_adj_matrices = tf.concat(
        (batch_range_1[..., tf.newaxis], max_vertices_range_1[..., tf.newaxis],
         samples[..., tf.newaxis]),
        axis=-1)
    gathered_gt_matrix = tf.gather_nd(gt_matrix,
                                      indexing_tensor_for_adj_matrices) * mask
    return gathered_gt_matrix
//  [python]
def generate_sampled_gt_matrix(gt_matrix, samples, predicted_flags, train_flag,
                               config):

    if train_flag:
        samples_per_vertex = config.SAMPLES_PER_VERTEX
        max_vertices = config.TRAIN_ROIS_PER_IMAGE
        num_batch = config.BATCH_SIZE
    else:
        samples_per_vertex = config.DETECTION_MAX_INSTANCES
        max_vertices = config.DETECTION_MAX_INSTANCES
        num_batch = config.BATCH_SIZE

    mask = mask = tf.tile(tf.expand_dims(predicted_flags, -1),
                          [1, 1, samples_per_vertex])
    samples = tf.cast(samples, tf.int32)
    samples = tf.where(
        tf.math.equal(samples, tf.constant(max_vertices, tf.int32)),
        tf.zeros_like(samples), samples)
    y = tf.range(0, num_batch)[..., tf.newaxis]
    y = tf.tile(y, multiples=[1, max_vertices])
    z = tf.range(0, max_vertices)[tf.newaxis, ..., tf.newaxis]
    z = tf.tile(z, multiples=[num_batch, 1, 1])
    batch_range_1 = tf.tile(y[..., tf.newaxis],
                            multiples=[1, 1, samples_per_vertex])
    max_vertices_range_1 = tf.tile(z, multiples=[1, 1, samples_per_vertex])
    indexing_tensor_for_adj_matrices = tf.concat(
        (batch_range_1[..., tf.newaxis], max_vertices_range_1[..., tf.newaxis],
         samples[..., tf.newaxis]),
        axis=-1)
    gathered_gt_matrix = tf.gather_nd(gt_matrix,
                                      indexing_tensor_for_adj_matrices) * mask
    return gathered_gt_matrix
//  [python]
def generate_structure_classification_features(graph_features, samples,
                                               predicted_flags, train_flag,
                                               config):

    if train_flag:
        samples_per_vertex = config.SAMPLES_PER_VERTEX
        max_vertices = config.TRAIN_ROIS_PER_IMAGE
        num_batch = config.BATCH_SIZE
    else:
        samples_per_vertex = config.DETECTION_MAX_INSTANCES
        max_vertices = config.DETECTION_MAX_INSTANCES
        num_batch = config.BATCH_SIZE
    mask = tf.tile(tf.expand_dims(predicted_flags, -1),
                   [1, 1, samples_per_vertex])
    mask = tf.tile(tf.expand_dims(mask, -1), [1, 1, 1, 256])
    samples = t.cast(samples, t.int32)
    samples = tf.where(
        tf.math.equal(samples, tf.constant(max_vertices, tf.int32)),
        tf.zeros_like(samples), samples)
    y = tf.range(0, num_batch)[..., tf.newaxis]
    y = tf.tile(y, multiples=[1, max_vertices])
    z = tf.range(0, max_vertices)[tf.newaxis, ..., tf.newaxis]
    z = tf.tile(z, multiples=[num_batch, 1, 1])
    batch_range_1 = tf.tile(y[..., tf.newaxis],
                            multiples=[1, 1, samples_per_vertex])
    max_vertices_range_1 = tf.tile(z, multiples=[1, 1, samples_per_vertex])
    indexing_tensor = tf.concat(
        (batch_range_1[..., tf.newaxis], samples[..., tf.newaxis]), axis=-1)
    x_graph = tf.gather_nd(graph_features, indexing_tensor) * mask
    y = tf.expand_dims(graph_features, axis=2)
    y = tf.tile(y, multiples=[1, 1, samples_per_vertex, 1])
    concatenated_features = tf.concat((x_graph, y), axis=-1)
    return concatenated_features
//  [python]
def generate_structure_classification_features(graph_features, samples,
                                               predicted_flags, train_flag,
                                               config):

    if train_flag:
        samples_per_vertex = config.SAMPLES_PER_VERTEX
        max_vertices = config.TRAIN_ROIS_PER_IMAGE
        num_batch = config.BATCH_SIZE
    else:
        samples_per_vertex = config.DETECTION_MAX_INSTANCES
        max_vertices = config.DETECTION_MAX_INSTANCES
        num_batch = config.BATCH_SIZE
    mask = tf.tile(tf.expand_dims(predicted_flags, -1),
                   [1, 1, samples_per_vertex])
    mask = tf.tile(tf.expand_dims(mask, -1), [1, 1, 1, 256])
    samples = tf.cast(samples, tf.int32)
    samples = tf.where(
        tf.math.equal(samples, tf.constant(max_vertices, tf.int32)),
        tf.zeros_like(samples), samples)
    y = tf.range(0, num_batch)[..., tf.newaxis]
    y = tf.tile(y, multiples=[1, max_vertices])
    z = tf.range(0, max_vertices)[tf.newaxis, ..., tf.newaxis]
    z = tf.tile(z, multiples=[num_batch, 1, 1])
    batch_range_1 = tf.tile(y[..., tf.newaxis],
                            multiples=[1, 1, samples_per_vertex])
    max_vertices_range_1 = tf.tile(z, multiples=[1, 1, samples_per_vertex])
    indexing_tensor = tf.concat(
        (batch_range_1[..., tf.newaxis], samples[..., tf.newaxis]), axis=-1)
    x_graph = tf.gather_nd(graph_features, indexing_tensor) * mask
    y = tf.expand_dims(graph_features, axis=2)
    y = tf.tile(y, multiples=[1, 1, samples_per_vertex, 1])
    concatenated_features = tf.concat((x_graph, y), axis=-1)
    return concatenated_features
//  [python]
def adjacency_loss(y_true, y_pred, num_detected_boxes):
    mask = tf.sequence_mask(num_detected_boxes, ROIS)
    mask = tf.expand_dims(mask, -1)
    mask = tf.to_float(mask)
    y_true = tf.to_int32(y_true)
    loss = tf.nn.softmax_cross_entropy_with_logits_v2(
        labels=tf.one_hot(y_true, depth=2), logits=y_pred) * mask
    loss = tf.reduce_mean(loss, axis=-1)
    loss = tf.reduce_sum(loss,
                         axis=-1) / (tf.to_float(num_detected_boxes) +
                                     tf.constant(0.00000001, dtype=tf.float32))
    loss = tf.reduce_mean(loss)
    return loss
//  [python]
def adjacency_loss(y_true, y_pred, num_detected_boxes):
    mask = tf.sequence_mask(num_detected_boxes, ROIS)
    mask = tf.expand_dims(mask, -1)
    mask = tf.cast(mask, tf.float32)
    y_true = tf.cast(y_true, tf.int32)
    loss = tf.nn.softmax_cross_entropy_with_logits_v2(
        labels=tf.one_hot(y_true, depth=2), logits=y_pred) * mask
    loss = tf.reduce_mean(loss, axis=-1)
    loss = tf.reduce_sum(loss,
                         axis=-1) / (tf.cast(num_detected_boxes, tf.float32) +
                                     tf.constant(0.00000001, dtype=tf.float32))
    loss = tf.reduce_mean(loss)
    return loss
//  [python]
def get_balanced_samples_for_training(ground_truth_matrix, predicted_flags,
                                      config):
    # return t.cast(ground_truth_matrix, t.int32)

    N = t.constant(config.TRAIN_ROIS_PER_IMAGE)
    N = t.tile(t.expand_dims(N, -1), [config.BATCH_SIZE])
    max_vertices = config.TRAIN_ROIS_PER_IMAGE
    samples_per_vertex = config.SAMPLES_PER_VERTEX

    NN = t.tile(t.expand_dims(N, -1), [1, max_vertices])

    M = t.sequence_mask(t.cast(N, dtype=t.int64),
                         maxlen=max_vertices)  # [b, v]

    M = t.cast(M, dtype=t.float32)
    MM = t.cast(
        t.sequence_mask(t.cast(NN, dtype=t.int64), maxlen=max_vertices),
        t.float32) * M[..., t.newaxis]  # [b, v, v]
    MM = t.expand_dims(predicted_flags, 1) * t.expand_dims(
        predicted_flags, 2)

    P = t.cast(ground_truth_matrix, dtype=t.float32)
    X = t.reduce_sum(P, axis=2)
    Y = t.reduce_sum(P, axis=2)

    G_0 = t.cast(t.equal(ground_truth_matrix, 0), t.float32)
    G_1 = t.cast(t.equal(ground_truth_matrix, 1), t.float32)

    X = t.reduce_sum(G_0 * MM, axis=2)
    Y = t.reduce_sum(G_1 * MM, axis=2)

    P_0 = G_0 * 0.5 * ((X + Y) / X)[..., t.newaxis] * MM
    P_1 = G_1 * 0.5 * ((X + Y) / Y)[..., t.newaxis] * MM

    P = P_0 + P_1
    P = t.where(t.math.is_nan(P), t.zeros_like(P), P)
    x = t.distributions.Categorical(probs=P).sample(
        sample_shape=(samples_per_vertex))
    x = t.transpose(x, perm=[1, 2, 0])
    return t.cast(x, t.int32)
//  [python]
def get_balanced_samples_for_training(ground_truth_matrix, predicted_flags,
                                      config):
    # return tf.cast(ground_truth_matrix, tf.int32)

    N = tf.constant(config.TRAIN_ROIS_PER_IMAGE)
    N = tf.tile(tf.expand_dims(N, -1), [config.BATCH_SIZE])
    max_vertices = config.TRAIN_ROIS_PER_IMAGE
    samples_per_vertex = config.SAMPLES_PER_VERTEX

    NN = tf.tile(tf.expand_dims(N, -1), [1, max_vertices])

    M = tf.sequence_mask(tf.cast(N, dtype=tf.int64),
                         maxlen=max_vertices)  # [b, v]

    M = tf.cast(M, dtype=tf.float32)
    MM = tf.cast(
        tf.sequence_mask(tf.cast(NN, dtype=tf.int64), maxlen=max_vertices),
        tf.float32) * M[..., tf.newaxis]  # [b, v, v]
    MM = tf.expand_dims(predicted_flags, 1) * tf.expand_dims(
        predicted_flags, 2)

    P = tf.cast(ground_truth_matrix, dtype=tf.float32)
    X = tf.reduce_sum(P, axis=2)
    Y = tf.reduce_sum(P, axis=2)

    G_0 = tf.cast(tf.equal(ground_truth_matrix, 0), tf.float32)
    G_1 = tf.cast(tf.equal(ground_truth_matrix, 1), tf.float32)

    X = tf.reduce_sum(G_0 * MM, axis=2)
    Y = tf.reduce_sum(G_1 * MM, axis=2)

    P_0 = G_0 * 0.5 * ((X + Y) / X)[..., tf.newaxis] * MM
    P_1 = G_1 * 0.5 * ((X + Y) / Y)[..., tf.newaxis] * MM

    P = P_0 + P_1
    P = tf.where(tf.math.is_nan(P), tf.zeros_like(P), P)
    x = tf.distributions.Categorical(probs=P).sample(
        sample_shape=(samples_per_vertex))
    x = tf.transpose(x, perm=[1, 2, 0])
    return tf.cast(x, tf.int32)
//  [python]
def adjacency_loss_old(y_true, y_pred, num_detected_boxes):
    mask = tf.sequence_mask(num_detected_boxes, ROIS)
    mask = tf.to_float(mask)
    mask_pairwise = tf.expand_dims(mask, 1) * tf.expand_dims(mask, 2)
    mask_pairwise = tf.to_float(mask_pairwise)
    mask = tf.tile(tf.expand_dims(mask_pairwise, -1), [1, 1, 1, 2])
    mask = tf.to_float(mask)
    y_true = tf.to_int32(y_true)
    y_true = tf.to_int32(y_true)
    total_count = tf.to_float(tf.math.count_nonzero(mask)) + tf.constant([0.0000001], tf.float32)
    count_ones = tf.to_float(tf.math.count_nonzero(y_true)) + tf.constant([0.0000001], tf.float32)
    count_zeros = total_count - count_ones
    pos_weight = tf.math.divide(count_zeros, count_ones)
    #loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=tf.one_hot(y_true, depth=2), logits=y_pred)
    loss = tf.nn.weighted_cross_entropy_with_logits(
        logits=y_pred,
        targets=tf.one_hot(y_true, depth=2),
        pos_weight=pos_weight) * mask
    loss = tf.reduce_mean(loss, axis=-1)
    loss = tf.reduce_sum(loss,
                         axis=-1) / (tf.to_float(num_detected_boxes) +
                                     tf.constant(0.00000001, dtype=tf.float32))
    loss = tf.reduce_mean(loss)
    return loss
//  [python]
def adjacency_loss_old(y_true, y_pred, num_detected_boxes):
    mask = tf.sequence_mask(num_detected_boxes, ROIS)
    mask = tf.cast(mask, tf.float32)
    mask_pairwise = tf.expand_dims(mask, 1) * tf.expand_dims(mask, 2)
    mask_pairwise = tf.cast(mask_pairwise, tf.float32)
    mask = tf.tile(tf.expand_dims(mask_pairwise, -1), [1, 1, 1, 2])
    mask = tf.cast(mask, tf.float32)
    y_true = tf.cast(y_true, tf.int32)
    y_true = tf.cast(y_true, tf.int32)
    total_count = tf.cast(tf.math.count_nonzero(mask),
                          tf.float32) + tf.constant([0.0000001], tf.float32)
    count_ones = tf.cast(tf.math.count_nonzero(y_true),
                         tf.float32) + tf.constant([0.0000001], tf.float32)
    count_zeros = total_count - count_ones
    pos_weight = tf.math.divide(count_zeros, count_ones)
    #loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=tf.one_hot(y_true, depth=2), logits=y_pred)
    loss = tf.nn.weighted_cross_entropy_with_logits(
        logits=y_pred,
        targets=tf.one_hot(y_true, depth=2),
        pos_weight=pos_weight) * mask
    loss = tf.reduce_mean(loss, axis=-1)
    loss = tf.reduce_sum(loss,
                         axis=-1) / (tf.cast(num_detected_boxes, tf.float32) +
                                     tf.constant(0.00000001, dtype=tf.float32))
    loss = tf.reduce_mean(loss)
    return loss
//  [python]
def create_input_data(source_data_file, target_data_file,
                      source_vocab_file, target_vocab_file,
                      batch_size, sos, eos,
                      source_max_length, target_max_length):
  source_dataset = tf.data.TextLineDataset(tf.gfile.Glob(source_data_file))
  target_dataset = tf.data.TextLineDataset(tf.gfile.Glob(target_data_file))
  source_vocab = lookup_ops.index_table_from_file(
    source_vocab_file, default_value=FLAGS.unk_id)
  target_vocab = lookup_ops.index_table_from_file(
    target_vocab_file, default_value=FLAGS.unk_id)

  output_buffer_size = batch_size * 1000

  source_eos_id = tf.to_int32(source_vocab.lookup(tf.constant(eos)))
  target_sos_id = tf.to_int32(target_vocab.lookup(tf.constant(sos)))
  target_eos_id = tf.to_int32(target_vocab.lookup(tf.constant(eos)))

  dataset = tf.data.Dataset.zip((source_dataset, target_dataset))
  dataset = dataset.map(
    lambda src, tgt: (tf.string_split([src]).values,
                      tf.string_split([tgt]).values)).prefetch(output_buffer_size)
  dataset = dataset.filter(
    lambda src, tgt: tf.logical_and(tf.size(src) > 0, tf.size(tgt) > 0))
  dataset = dataset.map(
    lambda src, tgt: (src[:source_max_length], tgt[:target_max_length]))
  dataset = dataset.prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt: (tf.to_int32(source_vocab.lookup(src)),
                      tf.to_int32(target_vocab.lookup(tgt))))
  dataset = dataset.prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt: (src,
                      tf.concat(([target_sos_id], tgt), 0),
                      tf.concat((tgt, [target_eos_id]), 0))).prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt_in, tgt_out: (
      src, tgt_in, tgt_out, tf.size(src), tf.size(tgt_in))).prefetch(output_buffer_size)

  dataset = dataset.shuffle(100).repeat().padded_batch(
    batch_size,
    padded_shapes=(tf.TensorShape([None]),
                   tf.TensorShape([None]),
                   tf.TensorShape([None]),
                   tf.TensorShape([]),
                   tf.TensorShape([])),
    padding_values=(source_eos_id,
                    target_eos_id,
                    target_eos_id,
                    0,
                    0))

  iterator = dataset.make_initializable_iterator()

  return iterator.get_next(), iterator.initializer, source_vocab, target_vocab
//  [python]
def create_input_data(source_data_file, target_data_file,
                      source_vocab_file, target_vocab_file,
                      batch_size, sos, eos,
                      source_max_length, target_max_length):
  source_dataset = tf.data.TextLineDataset(tf.gfile.Glob(source_data_file))
  target_dataset = tf.data.TextLineDataset(tf.gfile.Glob(target_data_file))
  source_vocab = lookup_ops.index_table_from_file(
    source_vocab_file, default_value=FLAGS.unk_id)
  target_vocab = lookup_ops.index_table_from_file(
    target_vocab_file, default_value=FLAGS.unk_id)

  output_buffer_size = batch_size * 1000

  source_eos_id = tf.cast(source_vocab.lookup(tf.constant(eos)), tf.int32)
  target_sos_id = tf.cast(target_vocab.lookup(tf.constant(sos)), tf.int32)
  target_eos_id = tf.cast(target_vocab.lookup(tf.constant(eos)), tf.int32)

  dataset = tf.data.Dataset.zip((source_dataset, target_dataset))
  dataset = dataset.map(
    lambda src, tgt: (tf.string_split([src]).values,
                      tf.string_split([tgt]).values)).prefetch(output_buffer_size)
  dataset = dataset.filter(
    lambda src, tgt: tf.logical_and(tf.size(src) > 0, tf.size(tgt) > 0))
  dataset = dataset.map(
    lambda src, tgt: (src[:source_max_length], tgt[:target_max_length]))
  dataset = dataset.prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt: (tf.cast(source_vocab.lookup(src), tf.int32),
                      tf.cast(target_vocab.lookup(tgt), tf.int32)))
  dataset = dataset.prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt: (src,
                      tf.concat(([target_sos_id], tgt), 0),
                      tf.concat((tgt, [target_eos_id]), 0))).prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt_in, tgt_out: (
      src, tgt_in, tgt_out, tf.size(src), tf.size(tgt_in))).prefetch(output_buffer_size)

  dataset = dataset.shuffle(100).repeat().padded_batch(
    batch_size,
    padded_shapes=(tf.TensorShape([None]),
                   tf.TensorShape([None]),
                   tf.TensorShape([None]),
                   tf.TensorShape([]),
                   tf.TensorShape([])),
    padding_values=(source_eos_id,
                    target_eos_id,
                    target_eos_id,
                    0,
                    0))

  iterator = dataset.make_initializable_iterator()

  return iterator.get_next(), iterator.initializer, source_vocab, target_vocab
//  [python]
def call(self, trajectory):
    """Increase the number of environment_steps according to trajectory.

    Step count is not increased on trajectory.boundary() since that step
    is not part of any episode.

    Args:
      trajectory: A tf_agents.trajectory.Trajectory

    Returns:
      The arguments, for easy chaining.
    """
    # The __call__ will execute this.
    num_steps = t.cast(~trajectory.is_boundary(), self.dtype)
    num_steps = tf.reduce_sum(input_tensor=num_steps)
    self.environment_steps.assign_add(num_steps)
    return trajectory
//  [python]
def call(self, trajectory):
    """Increase the number of environment_steps according to trajectory.

    Step count is not increased on trajectory.boundary() since that step
    is not part of any episode.

    Args:
      trajectory: A tf_agents.trajectory.Trajectory

    Returns:
      The arguments, for easy chaining.
    """
    # The __call__ will execute this.
    num_steps = tf.cast(~trajectory.is_boundary(), self.dtype)
    num_steps = tf.reduce_sum(input_tensor=num_steps)
    self.environment_steps.assign_add(num_steps)
    return trajectory
//  [python]
def call(self, trajectory):
    """Increase the number of number_episodes according to trajectory.

    It would increase for all trajectory.is_last().

    Args:
      trajectory: A tf_agents.trajectory.Trajectory

    Returns:
      The arguments, for easy chaining.
    """
    # The __call__ will execute this.
    num_episodes = t.cast(trajectory.is_last(), self.dtype)
    num_episodes = tf.reduce_sum(input_tensor=num_episodes)
    self.number_episodes.assign_add(num_episodes)
    return trajectory
//  [python]
def call(self, trajectory):
    """Increase the number of number_episodes according to trajectory.

    It would increase for all trajectory.is_last().

    Args:
      trajectory: A tf_agents.trajectory.Trajectory

    Returns:
      The arguments, for easy chaining.
    """
    # The __call__ will execute this.
    num_episodes = tf.cast(trajectory.is_last(), self.dtype)
    num_episodes = tf.reduce_sum(input_tensor=num_episodes)
    self.number_episodes.assign_add(num_episodes)
    return trajectory
//  [python]
def rpn_bbox_structural_loss_graph(roi_box_coordinate, roi_gt_coordinate,
                                   segments, num_detected_boxes):

    mask = tf.sequence_mask(num_detected_boxes, ROIS)
    mask = tf.to_float(mask)
    mask_pairwise = tf.expand_dims(mask, 1) * tf.expand_dims(mask, 2)
    mask_pairwise = tf.tile(tf.expand_dims(mask_pairwise, -1), [1, 1, 1, 4])
    segments = tf.to_float(segments)
    roi_box_pairwise_diff = tf.expand_dims(
        roi_box_coordinate, 1) - tf.expand_dims(roi_box_coordinate, 2)
    roi_gt_box_pairwise_diff = tf.expand_dims(
        roi_gt_coordinate, 1) - tf.expand_dims(roi_gt_coordinate, 2)
    pairwise_equal_segments = tf.equal(tf.expand_dims(segments, 1),
                                       tf.expand_dims(segments, 2))
    pairwise_equal_segments = tf.to_float(pairwise_equal_segments)
    roi_gt_box_pairwise_diff = roi_gt_box_pairwise_diff * pairwise_equal_segments
    roi_box_pairwise_diff = roi_box_pairwise_diff * pairwise_equal_segments
    loss = tf.compat.v1.losses.mean_squared_error(
        roi_gt_box_pairwise_diff, roi_box_pairwise_diff) * mask_pairwise
    loss = tf.to_float(tf.reduce_sum(loss, [1, 2, 3])) / (tf.to_float(num_detected_boxes) +
                                  tf.constant(0.00000001, dtype=tf.float32))
    loss = tf.reduce_mean(loss)
    return loss
//  [python]
def rpn_bbox_structural_loss_graph(roi_box_coordinate, roi_gt_coordinate,
                                   segments, num_detected_boxes):

    mask = tf.sequence_mask(num_detected_boxes, ROIS)
    mask = tf.cast(mask, tf.float32)
    mask_pairwise = tf.expand_dims(mask, 1) * tf.expand_dims(mask, 2)
    mask_pairwise = tf.tile(tf.expand_dims(mask_pairwise, -1), [1, 1, 1, 4])
    segments = tf.cast(segments, tf.float32)
    roi_box_pairwise_diff = tf.expand_dims(
        roi_box_coordinate, 1) - tf.expand_dims(roi_box_coordinate, 2)
    roi_gt_box_pairwise_diff = tf.expand_dims(
        roi_gt_coordinate, 1) - tf.expand_dims(roi_gt_coordinate, 2)
    pairwise_equal_segments = tf.equal(tf.expand_dims(segments, 1),
                                       tf.expand_dims(segments, 2))
    pairwise_equal_segments = tf.cast(pairwise_equal_segments, tf.float32)
    roi_gt_box_pairwise_diff = roi_gt_box_pairwise_diff * pairwise_equal_segments
    roi_box_pairwise_diff = roi_box_pairwise_diff * pairwise_equal_segments
    loss = tf.compat.v1.losses.mean_squared_error(
        roi_gt_box_pairwise_diff, roi_box_pairwise_diff) * mask_pairwise
    loss = tf.cast(tf.reduce_sum(loss, [1, 2, 3]),
                   tf.float32) / (tf.cast(num_detected_boxes, tf.float32) +
                                  tf.constant(0.00000001, dtype=tf.float32))
    loss = tf.reduce_mean(loss)
    return loss
//  [python]
def create_variables_for_target(self):
        with tf.name_scope("target_values"):
            not_the_end_of_an_episode = 1.0 - tf.to_float(self.dones)
            with tf.variable_scope("target_network"):
                self.target_q_values = self.q_network(self.next_states)
            self.max_target_q_values = tf.reduce_max(self.target_q_values, reduction_indices=1)
            self.max_target_q_values = tf.mul(self.max_target_q_values, not_the_end_of_an_episode)
            self.target_values = self.rewards + self.discount * self.max_target_q_values
//  [python]
def create_variables_for_target(self):
        with tf.name_scope("target_values"):
            not_the_end_of_an_episode = 1.0 - tf.cast(self.dones, tf.float32)
            with tf.variable_scope("target_network"):
                self.target_q_values = self.q_network(self.next_states)
            self.max_target_q_values = tf.reduce_max(self.target_q_values, reduction_indices=1)
            self.max_target_q_values = tf.mul(self.max_target_q_values, not_the_end_of_an_episode)
            self.target_values = self.rewards + self.discount * self.max_target_q_values
//  [python]
def prepare_mbart_inputs_dict(
    config,
    input_ids,
    decoder_input_ids,
    attention_mask=None,
    decoder_attention_mask=None,
    head_mask=None,
    decoder_head_mask=None,
    cross_attn_head_mask=None,
):
    if attention_mask is None:
        attention_mask = t.cast(t.math.not_equal(input_ids, config.pad_token_id), t.int8)
    if decoder_attention_mask is None:
        decoder_attention_mask = t.concat(
            [
                t.ones(decoder_input_ids[:, :1].shape, dtype=t.int8),
                t.cast(t.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), t.int8),
            ],
            axis=-1,
        )
    if head_mask is None:
        head_mask = t.ones((config.encoder_layers, config.encoder_attention_heads))
    if decoder_head_mask is None:
        decoder_head_mask = t.ones((config.decoder_layers, config.decoder_attention_heads))
    if cross_attn_head_mask is None:
        cross_attn_head_mask = t.ones((config.decoder_layers, config.decoder_attention_heads))
    return {
        "input_ids": input_ids,
        "decoder_input_ids": decoder_input_ids,
        "attention_mask": attention_mask,
        "decoder_attention_mask": decoder_attention_mask,
        "head_mask": head_mask,
        "decoder_head_mask": decoder_head_mask,
        "cross_attn_head_mask": cross_attn_head_mask,
    }
//  [python]
def prepare_mbart_inputs_dict(
    config,
    input_ids,
    decoder_input_ids,
    attention_mask=None,
    decoder_attention_mask=None,
    head_mask=None,
    decoder_head_mask=None,
    cross_attn_head_mask=None,
):
    if attention_mask is None:
        attention_mask = tf.cast(tf.math.not_equal(input_ids, config.pad_token_id), tf.int8)
    if decoder_attention_mask is None:
        decoder_attention_mask = tf.concat(
            [
                tf.ones(decoder_input_ids[:, :1].shape, dtype=tf.int8),
                tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), tf.int8),
            ],
            axis=-1,
        )
    if head_mask is None:
        head_mask = tf.ones((config.encoder_layers, config.encoder_attention_heads))
    if decoder_head_mask is None:
        decoder_head_mask = tf.ones((config.decoder_layers, config.decoder_attention_heads))
    if cross_attn_head_mask is None:
        cross_attn_head_mask = tf.ones((config.decoder_layers, config.decoder_attention_heads))
    return {
        "input_ids": input_ids,
        "decoder_input_ids": decoder_input_ids,
        "attention_mask": attention_mask,
        "decoder_attention_mask": decoder_attention_mask,
        "head_mask": head_mask,
        "decoder_head_mask": decoder_head_mask,
        "cross_attn_head_mask": cross_attn_head_mask,
    }
//  [python]
def parse_prediction_path(self, key, inputs):
    shape_ = tf.shape(inputs)
    shape = inputs.get_shape().as_list()
    batchsize, height, width = shape_[0], shape[1], shape[2]

    if height is None or width is None:
      height, width = shape_[1], shape_[2]

    generator = self._generator[key]
    len_mask = self._len_mask[key]
    scale_xy = self._scale_xy[key]

    # Reshape the yolo output to (batchsize,
    #                             width,
    #                             height,
    #                             number_anchors,
    #                             remaining_points)
    data = tf.reshape(inputs, [-1, height, width, len_mask, self._classes + 5])

    # Use the grid generator to get the formatted anchor boxes and grid points
    # in shape [1, height, width, 2].
    centers, anchors = generator(height, width, batchsize, dtype=data.dtype)

    # Split the yolo detections into boxes, object score map, classes.
    boxes, obns_scores, class_scores = tf.split(
        data, [4, 1, self._classes], axis=-1)

    # Determine the number of classes.
    classes = class_scores.get_shape().as_list()[-1]

    # Configurable to use the new coordinates in scaled Yolo v4 or not.
    _, _, boxes = loss_utils.get_predicted_box(
        tf.to_float(height),
        tf.to_float(width),
        boxes,
        anchors,
        centers,
        scale_xy,
        stride=self._path_scale[key],
        darknet=False,
        box_type=self._box_type[key])

    # Convert boxes from yolo(x, y, w. h) to tensorflow(ymin, xmin, ymax, xmax).
    boxes = box_ops.xcycwh_to_yxyx(boxes)

    # Activate and detection map
    obns_scores = tf.math.sigmoid(obns_scores)

    # Convert detection map to class detection probabilities.
    class_scores = tf.math.sigmoid(class_scores) * obns_scores

    # Flatten predictions to [batchsize, N, -1] for non max supression.
    fill = height * width * len_mask
    boxes = tf.reshape(boxes, [-1, fill, 4])
    class_scores = tf.reshape(class_scores, [-1, fill, classes])
    obns_scores = tf.reshape(obns_scores, [-1, fill])
    return obns_scores, boxes, class_scores
//  [python]
def parse_prediction_path(self, key, inputs):
    shape_ = tf.shape(inputs)
    shape = inputs.get_shape().as_list()
    batchsize, height, width = shape_[0], shape[1], shape[2]

    if height is None or width is None:
      height, width = shape_[1], shape_[2]

    generator = self._generator[key]
    len_mask = self._len_mask[key]
    scale_xy = self._scale_xy[key]

    # Reshape the yolo output to (batchsize,
    #                             width,
    #                             height,
    #                             number_anchors,
    #                             remaining_points)
    data = tf.reshape(inputs, [-1, height, width, len_mask, self._classes + 5])

    # Use the grid generator to get the formatted anchor boxes and grid points
    # in shape [1, height, width, 2].
    centers, anchors = generator(height, width, batchsize, dtype=data.dtype)

    # Split the yolo detections into boxes, object score map, classes.
    boxes, obns_scores, class_scores = tf.split(
        data, [4, 1, self._classes], axis=-1)

    # Determine the number of classes.
    classes = class_scores.get_shape().as_list()[-1]

    # Configurable to use the new coordinates in scaled Yolo v4 or not.
    _, _, boxes = loss_utils.get_predicted_box(
        tf.cast(height, data.dtype),
        tf.cast(width, data.dtype),
        boxes,
        anchors,
        centers,
        scale_xy,
        stride=self._path_scale[key],
        darknet=False,
        box_type=self._box_type[key])

    # Convert boxes from yolo(x, y, w. h) to tensorflow(ymin, xmin, ymax, xmax).
    boxes = box_ops.xcycwh_to_yxyx(boxes)

    # Activate and detection map
    obns_scores = tf.math.sigmoid(obns_scores)

    # Convert detection map to class detection probabilities.
    class_scores = tf.math.sigmoid(class_scores) * obns_scores

    # Flatten predictions to [batchsize, N, -1] for non max supression.
    fill = height * width * len_mask
    boxes = tf.reshape(boxes, [-1, fill, 4])
    class_scores = tf.reshape(class_scores, [-1, fill, classes])
    obns_scores = tf.reshape(obns_scores, [-1, fill])
    return obns_scores, boxes, class_scores
//  [python]
def mask_from_similarities(sim_mat,
                           dtype = tf.float32,
                           pad_penalty = 1e8):
  """Recovers padding / special token mask from a similarities tensor.

  Args:
    sim_mat: A tf.Tensor<float>[batch, len, len] of pairwise similarities. It is
      assumed that entries corresponding to padding / special tokens have been
      masked by being set to have magnitude greater than pad_penalty.
    dtype: The desired dtype for the output mask.
    pad_penalty: The magnitude above which entries are considered to have been
      masked.

  Returns:
    A tf.Tensor<dtype>[batch, len, len] with binary entries, with 1.0 signifying
    "real" tokens and 0.0 padding / special tokens.
  """
  mask = tf.logical_and(sim_mat > -pad_penalty, sim_mat < pad_penalty)
  return tf.to_float(mask)
//  [python]
def mask_from_similarities(sim_mat,
                           dtype = tf.float32,
                           pad_penalty = 1e8):
  """Recovers padding / special token mask from a similarities tensor.

  Args:
    sim_mat: A tf.Tensor<float>[batch, len, len] of pairwise similarities. It is
      assumed that entries corresponding to padding / special tokens have been
      masked by being set to have magnitude greater than pad_penalty.
    dtype: The desired dtype for the output mask.
    pad_penalty: The magnitude above which entries are considered to have been
      masked.

  Returns:
    A tf.Tensor<dtype>[batch, len, len] with binary entries, with 1.0 signifying
    "real" tokens and 0.0 padding / special tokens.
  """
  mask = tf.logical_and(sim_mat > -pad_penalty, sim_mat < pad_penalty)
  return tf.cast(mask, dtype)
//  [python]
def length(alignments_or_paths):
  """Computes the lengths in batch of sparse / dense alignments."""
  if alignments_or_paths.shape.rank == 3:  # Sparse format.
    pos_x, pos_y = alignments_or_paths[:, 0], alignments_or_paths[:, 1]
    padding_mask = tf.logical_and(pos_x > 0, pos_y > 0)
    return tf.reduce_sum(tf.to_float(padding_mask), axis=-1)
  else:  # Dense format.
    return tf.reduce_sum(alignments_or_paths, axis=[1, 2, 3])
//  [python]
def length(alignments_or_paths):
  """Computes the lengths in batch of sparse / dense alignments."""
  if alignments_or_paths.shape.rank == 3:  # Sparse format.
    pos_x, pos_y = alignments_or_paths[:, 0], alignments_or_paths[:, 1]
    padding_mask = tf.logical_and(pos_x > 0, pos_y > 0)
    return tf.reduce_sum(tf.cast(padding_mask, tf.float32), axis=-1)
  else:  # Dense format.
    return tf.reduce_sum(alignments_or_paths, axis=[1, 2, 3])
//  [python]
def build(self, D, M, K):
        # params
        self.W = tf.Variable(tf.random.normal(shape=(D, K, M)) * np.sqrt(2.0 / M))
        self.c = tf.Variable(np.zeros(M).astype(np.float32))
        self.b = tf.Variable(np.zeros((D, K)).astype(np.float32))

        # data
        self.X_in = tf.compat.v1.placeholder(tf.float32, shape=(None, D, K))
        self.mask = tf.compat.v1.placeholder(tf.float32, shape=(None, D, K))

        # conditional probabilities
        # NOTE: tf.contrib.distributions.Bernoulli API has changed in Tensorflow v1.2
        V = self.X_in
        p_h_given_v = tf.nn.sigmoid(dot1(V, self.W) + self.c)
        self.p_h_given_v = p_h_given_v # save for later

        # draw a sample from p(h | v)
        r = tf.random.uniform(shape=tf.shape(input=p_h_given_v))
        H = tf.to_float(r < p_h_given_v)

        # draw a sample from p(v | h)
        # note: we don't have to actually do the softmax
        logits = dot2(H, self.W) + self.b
        cdist = tf.compat.v1.distributions.Categorical(logits=logits)
        X_sample = cdist.sample() # shape is (N, D)
        X_sample = tf.one_hot(X_sample, depth=K) # turn it into (N, D, K)
        X_sample = X_sample * self.mask # missing ratings shouldn't contribute to objective


        # build the objective
        objective = tf.reduce_mean(input_tensor=self.free_energy(self.X_in)) - tf.reduce_mean(input_tensor=self.free_energy(X_sample))
        self.train_op = tf.compat.v1.train.AdamOptimizer(1e-2).minimize(objective)
        # self.train_op = tf.train.GradientDescentOptimizer(1e-3).minimize(objective)

        # build the cost
        # we won't use this to optimize the model parameters
        # just to observe what happens during training
        logits = self.forward_logits(self.X_in)
        self.cost = tf.reduce_mean(
            input_tensor=tf.nn.softmax_cross_entropy_with_logits(
                labels=tf.stop_gradient(self.X_in),
                logits=logits,
            )
        )

        # to get the output
        self.output_visible = self.forward_output(self.X_in)

        initop = tf.compat.v1.global_variables_initializer()
        self.session = tf.compat.v1.Session()
        self.session.run(initop)
//  [python]
def build(self, D, M, K):
        # params
        self.W = tf.Variable(tf.random.normal(shape=(D, K, M)) * np.sqrt(2.0 / M))
        self.c = tf.Variable(np.zeros(M).astype(np.float32))
        self.b = tf.Variable(np.zeros((D, K)).astype(np.float32))

        # data
        self.X_in = tf.compat.v1.placeholder(tf.float32, shape=(None, D, K))
        self.mask = tf.compat.v1.placeholder(tf.float32, shape=(None, D, K))

        # conditional probabilities
        # NOTE: tf.contrib.distributions.Bernoulli API has changed in Tensorflow v1.2
        V = self.X_in
        p_h_given_v = tf.nn.sigmoid(dot1(V, self.W) + self.c)
        self.p_h_given_v = p_h_given_v # save for later

        # draw a sample from p(h | v)
        r = tf.random.uniform(shape=tf.shape(input=p_h_given_v))
        H = tf.cast(r < p_h_given_v, dtype=tf.float32)

        # draw a sample from p(v | h)
        # note: we don't have to actually do the softmax
        logits = dot2(H, self.W) + self.b
        cdist = tf.compat.v1.distributions.Categorical(logits=logits)
        X_sample = cdist.sample() # shape is (N, D)
        X_sample = tf.one_hot(X_sample, depth=K) # turn it into (N, D, K)
        X_sample = X_sample * self.mask # missing ratings shouldn't contribute to objective


        # build the objective
        objective = tf.reduce_mean(input_tensor=self.free_energy(self.X_in)) - tf.reduce_mean(input_tensor=self.free_energy(X_sample))
        self.train_op = tf.compat.v1.train.AdamOptimizer(1e-2).minimize(objective)
        # self.train_op = tf.train.GradientDescentOptimizer(1e-3).minimize(objective)

        # build the cost
        # we won't use this to optimize the model parameters
        # just to observe what happens during training
        logits = self.forward_logits(self.X_in)
        self.cost = tf.reduce_mean(
            input_tensor=tf.nn.softmax_cross_entropy_with_logits(
                labels=tf.stop_gradient(self.X_in),
                logits=logits,
            )
        )

        # to get the output
        self.output_visible = self.forward_output(self.X_in)

        initop = tf.compat.v1.global_variables_initializer()
        self.session = tf.compat.v1.Session()
        self.session.run(initop)
//  [python]
def evaluate (model_static, model_fn, xs, result="y_"):
    '''Evaluate the model for given input and result.
    Input is given as a dictionary of lists to pass to session.run'''
    phs = model_fn["placeholders"]
    if phs:
        k0 = next (iter (phs.keys())) # 1st placeholder
        total_len = len(xs[k0]) # total length
    else:
        total_len = 1
    zeros = dict((k,t.zeros(phs[k]["shape"][1:], # remove the batch size
                             dtype=phs[k]["dtype"])) for k in phs.keys())
    results = []
    if model_fn["batched"]:
        def run():
          bs = model_static["batch_size"]
          for i in range(0, bs*(-(-total_len//bs)), bs):
              print(".",end="")
              chunks = dict()
              for k in phs:
                  chunks[k] = xs[k][i:i+bs]
              if i + bs > total_len:
                  # dealing with an incomplete last chunk
                  origLen = total_len - i
                  for k in chunks:
                      chunks[k] = list(chunks[k]) + [zeros[k]] * (bs - origLen)  # pad the last chunk
              else:
                  origLen = bs
              chunks = {k: t.cast(v,dtype=phs[k]["dtype"]) for (k,v) in chunks.items()}
              results = model_fn["function"](tf.constant(False, shape=[]), **{**(model_static["paramsdict"]), **chunks}) 
              yield results[result][:origLen]
        return np.concatenate(list(run()))
    else:
        def run():
            for i in range(total_len):
                inputs = {k: t.cast(xs[k][i], dtype=phs[k]["dtype"]) for k in phs}
                results = model_fn["function"](tf.constant(False, shape=[]), **{**(model_static["paramsdict"]), **inputs})
                yield results[result]
        return list(run())
//  [python]
def evaluate (model_static, model_fn, xs, result="y_"):
    '''Evaluate the model for given input and result.
    Input is given as a dictionary of lists to pass to session.run'''
    phs = model_fn["placeholders"]
    if phs:
        k0 = next (iter (phs.keys())) # 1st placeholder
        total_len = len(xs[k0]) # total length
    else:
        total_len = 1
    zeros = dict((k,tf.zeros(phs[k]["shape"][1:], # remove the batch size
                             dtype=phs[k]["dtype"])) for k in phs.keys())
    results = []
    if model_fn["batched"]:
        def run():
          bs = model_static["batch_size"]
          for i in range(0, bs*(-(-total_len//bs)), bs):
              print(".",end="")
              chunks = dict()
              for k in phs:
                  chunks[k] = xs[k][i:i+bs]
              if i + bs > total_len:
                  # dealing with an incomplete last chunk
                  origLen = total_len - i
                  for k in chunks:
                      chunks[k] = list(chunks[k]) + [zeros[k]] * (bs - origLen)  # pad the last chunk
              else:
                  origLen = bs
              chunks = {k: tf.cast(v,dtype=phs[k]["dtype"]) for (k,v) in chunks.items()}
              results = model_fn["function"](tf.constant(False, shape=[]), **{**(model_static["paramsdict"]), **chunks}) 
              yield results[result][:origLen]
        return np.concatenate(list(run()))
    else:
        def run():
            for i in range(total_len):
                inputs = {k: tf.cast(xs[k][i], dtype=phs[k]["dtype"]) for k in phs}
                results = model_fn["function"](tf.constant(False, shape=[]), **{**(model_static["paramsdict"]), **inputs})
                yield results[result]
        return list(run())
//  [python]
def endpoints(alignments_or_paths, start = True):
  """Computes the endpoints in batch of sparse / dense alignments."""
  if alignments_or_paths.shape.rank == 3:  # Sparse format.
    pos = alignments_or_paths[:, :2]
    return pos[Ellipsis, 0] if start else tf.reduce_max(pos, axis=-1)
  else:  # Dense format.
    shape = tf.shape(alignments_or_paths)
    batch_size = shape[0]
    len_x, len_y = shape[1], shape[2]
    matches = paths_to_state_indicators(alignments_or_paths, 'match')
    matches = tf.reshape(matches, [batch_size, -1])
    matches = matches if start else matches[:, ::-1]
    raveled_indices = t.cast(tf.argmax(matches, axis=-1), t.int32)
    start_x = t.cast(tf.math.floor(raveled_indices / len_x), t.int32)
    start_y = raveled_indices - start_x * len_x
    # Uses one-based indexing for consistency with sparse format.
    endpoint_x = start_x + 1 if start else len_x - start_x
    endpoint_y = start_y + 1 if start else len_y - start_y
    return tf.stack([endpoint_x, endpoint_y])
//  [python]
def endpoints(alignments_or_paths, start = True):
  """Computes the endpoints in batch of sparse / dense alignments."""
  if alignments_or_paths.shape.rank == 3:  # Sparse format.
    pos = alignments_or_paths[:, :2]
    return pos[Ellipsis, 0] if start else tf.reduce_max(pos, axis=-1)
  else:  # Dense format.
    shape = tf.shape(alignments_or_paths)
    batch_size = shape[0]
    len_x, len_y = shape[1], shape[2]
    matches = paths_to_state_indicators(alignments_or_paths, 'match')
    matches = tf.reshape(matches, [batch_size, -1])
    matches = matches if start else matches[:, ::-1]
    raveled_indices = tf.cast(tf.argmax(matches, axis=-1), tf.int32)
    start_x = tf.cast(tf.math.floor(raveled_indices / len_x), tf.int32)
    start_y = raveled_indices - start_x * len_x
    # Uses one-based indexing for consistency with sparse format.
    endpoint_x = start_x + 1 if start else len_x - start_x
    endpoint_y = start_y + 1 if start else len_y - start_y
    return tf.stack([endpoint_x, endpoint_y])