import torch


def apply_homogeneous_affine_transform(transform: torch.Tensor, position: torch.Tensor):
    """
    Apply an homogeneous affine transform (4x4 for 3D or 3x3 for 2D) to a position

    Args:
        transform: an homogeneous affine transformation
        position: XY(Z) position

    Returns:
        a transformed position XY(Z)
    """
    assert len(transform.shape) == 2
    assert len(position.shape) == 1
    dim = position.shape[0]
    assert transform.shape[0] == transform.shape[1]
    assert transform.shape[0] == dim + 1
    # decompose the transform as a (3x3 transform, translation) components
    position = position.unsqueeze(1).type(torch.float32)
    return transform[:dim, :dim].mm(position).squeeze(1) + transform[:dim, dim]