def reshape_to_images(x, im_h=None):
    """Reshape set of vectors to images, assuming square or given height."""
    if im_h is None:
        im_h = int(x.shape[1]**0.5)
    im_w = x.shape[1] // im_h
    return x.reshape(x.shape[0], im_h, im_w)