import torch


def batch_log_maxvecmul(A, b):
    """Similar to batch_log_matvecmul, but takes a maximum instead of
    logsumexp. Returns both the max and the argmax.

    Arguments
    ---------
    A : torch.Tensor (batch, dim1, dim2)
        Tensor.
    b : torch.Tensor (batch, dim1)
        Tensor

    Outputs
    -------
    x : torch.Tensor (batch, dim1)
        Tensor.
    argmax : torch.Tensor (batch, dim1)
        Tensor.

    Example
    -------
    >>> A = torch.tensor([[[   0., -1.],
    ...                    [ -1e5,  0.]]])
    >>> b = torch.tensor([[0., 0.,]])
    >>> x, argmax = batch_log_maxvecmul(A, b)
    >>> x
    tensor([[0., 0.]])
    >>> argmax
    tensor([[0, 1]])
    """
    b = b.unsqueeze(1)
    x, argmax = torch.max(A + b, dim=2)

    return x, argmax