import torch


def transform_coords(Ms, coords, wrap_around=False):
    """
    Ms - affine transform matrices, (B,2,3), assumes x-y coordinate system in math
    coords - (B,dimx,dimy,2)
    """
    B, dimx, dimy = coords.size()[0], coords.size()[1], coords.size()[2]
    augmented_coords = torch.cat((coords, torch.ones(B,dimx,dimy,1).cuda()),dim=3) # (B,dimx,dimy,3)
    new_coords = torch.einsum('bij,bxyj->bxyi',Ms,augmented_coords) # (B,dimx,dimy,2)
    if wrap_around:
        new_coords[:,:,:,0] = new_coords[:,:,:,0] % dimx
        new_coords[:,:,:,1] = new_coords[:,:,:,1] % dimy
    return new_coords