import torch


def matrix_to_homogeneous(batch: torch.Tensor) -> torch.Tensor:
    """
    Transforms a given transformation matrix to a homogeneous
    transformation matrix.

    Args:
        batch: the batch of matrices to convert [N, dim, dim]

    Returns:
        torch.Tensor: the converted batch of matrices

    """
    if batch.size(-1) == batch.size(-2):
        missing = batch.new_zeros(size=(*batch.shape[:-1], 1))
        batch = torch.cat([batch, missing], dim=-1)

    missing = torch.zeros(
        (batch.size(0), *[1 for tmp in batch.shape[1:-1]], batch.size(-1)), device=batch.device, dtype=batch.dtype
    )

    missing[..., -1] = 1

    return torch.cat([batch, missing], dim=-2)