    def gumbel_with_max(self, phi, T) -> torch.Tensor:
        """
        Sample `Gumbel(phi)` conditioned on the maximum value being equal to `T`.

        `phi` should have shape `(batch_size, num_classes)` and `T` should have
        shape `(batch_size, 1)`.
        """
        # Shape: (batch_size, num_classes)
        G_phi = self.gumbel(phi)

        # Now we find the maximum from these samples.
        # Shape: (batch_size, )
        Z, _ = G_phi.max(dim=-1)

        # Shape: (batch_size, num_classes)
        v = T - G_phi + torch.log1p(-torch.exp(G_phi - Z.unsqueeze(-1)))

        # Shape: (batch_size, num_classes)
        return T - torch.nn.functional.relu(v) - torch.log1p(torch.exp(-v.abs()))
Sample `Gumbel(phi)` conditioned on the maximum value being equal to `T`.

`phi` should have shape `(batch_size, num_classes)` and `T` should have
shape `(batch_size, 1)`.
allennlp/nn/beam_search.py