import torch


def get_affine_grid(matrix, grid_shape):
    """Given a matrix and a grid shape, returns the grid transformed by the
    matrix (typically a rotation matrix).

    Args:
        matrix (torch.Tensor): Batch of matrices of size (batch_size, 3, 3).
        grid_shape (torch.size): Shape of returned affine grid. Should be of the
            form (batch_size, channels, depth, height, width).

    Notes:
        We use align_corners=False in affine_grid. See
        https://discuss.pytorch.org/t/what-we-should-use-align-corners-false/22663/9
        for a nice illustration of why this is.
    """
    batch_size = matrix.shape[0]
    # Last column of affine matrix corresponds to translation which is 0 in our
    # case. Therefore pad original matrix with zeros, so shape changes from
    # (batch_size, 3, 3) to (batch_size, 3, 4)
    translations = torch.zeros(batch_size, 3, 1, device=matrix.device)
    affine_matrix = torch.cat([matrix, translations], dim=2)
    return torch.nn.functional.affine_grid(affine_matrix, grid_shape,
                                           align_corners=False)