import torch


def dot(a, b, keepdim=False, out=None):
    """(Batched) dot product

    Parameters
    ----------
    a : (..., N) tensor
    b : (..., N) tensor
    keepdim : bool, default=False
    out : tensor, optional

    Returns
    -------
    ab : (..., [1]) tensor

    """
    a = a[..., None, :]
    b = b[..., :, None]
    ab = torch.matmul(a, b, out=out)
    if keepdim:
        ab = ab[..., 0, 0]
    else:
        ab = ab[..., 0]
    return ab