import torch


def identity_grid(shape, dtype=None, device=None, jitter=False):
    """Returns an identity deformation field.

    Parameters
    ----------
    shape : (dim,) sequence of int
        Spatial dimension of the field.
    dtype : torch.dtype, default=`get_default_dtype()`
        Data type.
    device torch.device, optional
        Device.
    jitter : bool or 'reproducible', default=False
        Jitter identity grid.

    Returns
    -------
    grid : (*shape, dim) tensor
        Transformation field

    """
    mesh1d = [torch.arange(float(s), dtype=dtype, device=device)
              for s in shape]
    grid = torch.meshgrid(*mesh1d)
    grid = torch.stack(grid, dim=-1)
    if jitter:
        reproducible = jitter == 'reproducible'
        device_ids = [grid.device.index] if grid.device.type == 'cuda' else None
        with torch.random.fork_rng(device_ids, enabled=reproducible):
            if reproducible:
                torch.manual_seed(0)
            jitter = torch.rand_like(grid).sub_(0.5).mul_(0.1)
            grid += jitter
    return grid