import torch


def ssm(vec):
    """
    Get the skew-symmetric matrix equivalent to cross-prodict for a (batch of) 3-element vector(s).
    """
    if len(vec.shape) == 1:
        return ssm(vec.unsuqeeze(0)).squeeze()
    
    out = torch.zeros(vec.shape[0], 3, 3, device=vec.device)
    out[:, 0, 1] = -vec[:, 2]
    out[:, 0, 2] = vec[:, 1]
    out[:, 1, 0] = vec[:, 2]
    out[:, 1, 2] = -vec[:, 0]
    out[:, 2, 0] = -vec[:, 1]
    out[:, 2, 1] = vec[:, 0]

    return out