import torch


def apply_4x4(RT, XYZ):
    """
    RT: B x 4 x 4
    XYZ: B x N x 3
    """
    RT = RT.to(XYZ.device)
    B, N, _ = list(XYZ.shape)
    ones = torch.ones([B, N, 1], device=XYZ.device)
    XYZ1 = torch.cat([XYZ, ones], 2)
    XYZ1_t = torch.transpose(XYZ1, 1, 2)
    # this is B x 4 x N
    XYZ2_t = torch.matmul(RT, XYZ1_t)
    XYZ2 = torch.transpose(XYZ2_t, 1, 2)
    XYZ2 = XYZ2[:,:,:3]
    return XYZ2