    def output_layer(
        self,
        features: Tensor,
        attn: Tensor,
        src_tokens: Tensor,
        p_gens: Tensor
    ) -> Tensor:
        """
        Project features to the vocabulary size and mix with the attention
        distributions.
        """
        if self.force_p_gen is not None:
            p_gens = self.force_p_gen

        # project back to size of vocabulary
        if self.adaptive_softmax is None:
            logits = self.output_projection(features)
        else:
            logits = features

        batch_size = logits.shape[0]
        output_length = logits.shape[1]
        assert logits.shape[2] == self.num_embeddings
        assert src_tokens.shape[0] == batch_size
        src_length = src_tokens.shape[1]

        # The final output distribution will be a mixture of the normal output
        # distribution (softmax of logits) and attention weights.
        gen_dists = self.get_normalized_probs_scriptable(
            (logits, None), log_probs=False, sample=None
        )
        gen_dists = torch.mul(gen_dists, p_gens)
        padding_size = (batch_size, output_length, self.num_oov_types)
        padding = gen_dists.new_zeros(padding_size)
        gen_dists = torch.cat((gen_dists, padding), 2)
        assert gen_dists.shape[2] == self.num_types

        # Scatter attention distributions to distributions over the extended
        # vocabulary in a tensor of shape [batch_size, output_length,
        # vocab_size]. Each attention weight will be written into a location
        # that is for other dimensions the same as in the index tensor, but for
        # the third dimension it's the value of the index tensor (the token ID).
        attn = torch.mul(attn.float(), 1 - p_gens)
        index = src_tokens[:, None, :]
        index = index.expand(batch_size, output_length, src_length)
        attn_dists_size = (batch_size, output_length, self.num_types)
        attn_dists = attn.new_zeros(attn_dists_size)
        attn_dists.scatter_add_(2, index, attn.float())

        # Final distributions, [batch_size, output_length, num_types].
        return gen_dists + attn_dists
Project features to the vocabulary size and mix with the attention
distributions.
examples/pointer_generator/pointer_generator_src/transformer_pg.py