import torch


def _compute_translation_matrix(translation):
    """
        Computes affine matrix for translation.
    """
    matrix = torch.eye(3, device=translation.device, dtype=translation.dtype)
    matrix = matrix.repeat(translation.shape[0], 1, 1)

    dx, dy = torch.chunk(translation, chunks=2, dim=-1)

    matrix[..., 0, 2:3] += dx
    matrix[..., 1, 2:3] += dy

    return matrix