import torch


def softmax_2d(inp):
    """Compute the softmax with the last two tensor dimensions combined."""
    size = inp.size()
    flat = inp.view(-1, size[-1] * size[-2])
    flat = torch.nn.functional.softmax(flat)
    return flat.view(*size)