import torch


def batch_log_matvecmul(A, b):
    """For each 'matrix' and 'vector' pair in the batch, do matrix-vector
    multiplication in the log domain, i.e., logsumexp instead of add,
    add instead of multiply.

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

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

    Example
    -------
    >>> A = torch.tensor([[[   0., 0.],
    ...                    [ -1e5, 0.]]])
    >>> b = torch.tensor([[0., 0.,]])
    >>> x = batch_log_matvecmul(A, b)
    >>> x
    tensor([[0.6931, 0.0000]])
    >>>
    >>> # non-log domain equivalent without batching functionality
    >>> A_ = torch.tensor([[1., 1.],
    ...                    [0., 1.]])
    >>> b_ = torch.tensor([1., 1.,])
    >>> x_ = torch.matmul(A_, b_)
    >>> x_
    tensor([2., 1.])
    """
    b = b.unsqueeze(1)
    x = torch.logsumexp(A + b, dim=2)

    return x