62d2ea56ae4a090aa68baf133137982a836700bd,thumt/utils/search.py,,create_inference_graph,#Any#Any#Any#,430

Before Change

        else:
            return ids[:, :top_beams, 1:]
    else:
        if top_beams == 1:
            return ids[:, 0, 1:], scores[:, 0]
        else:
            return ids[:, :top_beams, 1:], scores[:, :top_beams]

After Change

    // shape: [batch, beam_size]
    normalized_scores = scores / output_length

    if params.decode_normalize:
        scores, indices = tf.nn.top_k(normalized_scores, k=top_beams)
        // shape of ids: [batch, beam_size, max_length]
        // shape of coordinates: [batch, beam_size, 2]
        batch_pos = compute_batch_indices(batch_size, beam_size)
        coordinates = tf.stack([batch_pos, indices], axis=2)
        ids = tf.gather_nd(ids, coordinates)

    // Return `top_beams` decoding
    // (also remove initial id from the beam search)
    return ids[:, :top_beams, 1:], scores[:, :top_beams]

Instances