Class Manipulable

flowmap/misc/manipulable.py:19–123

Give a dataclass that bundles tensors, arrays, and lists functionality that mimics regular tensor manipulation. For example, this allows you to slice every element of a dataclass at once.

class Manipulable:
    """Give a dataclass that bundles tensors, arrays, and lists functionality that
    mimics regular tensor manipulation. For example, this allows you to slice every
    element of a dataclass at once.
    """

    def to(self: T, device: torch.device) -> T:
        """Return a shallow copy of this instance in which all tensors have been moved
        to the specified device.
        """

        replacements = {}

        for field in fields(self):
            # Move fields that are torch.Tensor to the specified device.
            value = getattr(self, field.name)
            if isinstance(value, Tensor):
                replacements[field.name] = value.to(device)

        return replace(self, **replacements)

    def __getitem__(self: T, slices: slice | tuple[SliceLike, ...]) -> T:
        """Return a shallow copy of this instance in which all tensors, arrays, and
        lists have been sliced according to the specified slices. If the number of
        slices exceeds the number of dimensions of a particular dataclass element,
        ignore the trailing slices. For now, ellipses are not supported.
        """

        # To simplify the implementation below, ensure that slices is a tuple of slices.
        if not isinstance(slices, tuple):
            slices = (slices,)

        replacements = {}

        # Recursively slice any sliceable fields.
        for field in fields(self):
            value = getattr(self, field.name)
            value = Manipulable.recursive_slice(value, slices)
            replacements[field.name] = value

        return replace(self, **replacements)

    @staticmethod
    def recursive_slice(value: Any, slices: tuple[SliceLike, ...]) -> Any:
        if not isinstance(value, Sliceable):
            return value

        # Torch tensors and NumPy arrays don't need to be recursively sliced, since they
        # already support tuples of slices.
        if isinstance(value, Tensor) or isinstance(value, np.ndarray):
            # If the number of slices exceeds the number of dimensions, drop the
            # trailing slices.
            slices = slices[: value.ndim]

            # Apply the remaining slices.
            return value[slices]

        # Lists and tuples need to be recursively sliced. To check if we should recurse