import torch


def zero_pad(x, target_shape):
    """
    when x is of shape (a,b,c,d) and target shape is (e,f), x will be centered padded to (a,b,e,f)
    :param x: Tensor, input data that need to pad
    :param target_shape: target shape of last few dims
    :return: a padded tensor
    """
    pads = []
    for xs, ts in zip(x.shape[::-1], target_shape[::-1]):
        assert ts >= xs
        assert (ts - xs) % 2 == 0
        pads.append((ts - xs) // 2)
        pads.append((ts - xs) // 2)
    return torch.nn.functional.pad(x, pads)