Source code for torchvision.tv_tensors

import torch

from ._bounding_boxes import BoundingBoxes, BoundingBoxFormat, is_rotated_bounding_format
from ._image import Image
from ._keypoints import KeyPoints
from ._mask import Mask
from ._torch_function_helpers import set_return_type
from ._tv_tensor import TVTensor
from ._video import Video


# TODO: Fix this. We skip this method as it leads to
# RecursionError: maximum recursion depth exceeded while calling a Python object
# Until `disable` is removed, there will be graph breaks after all calls to functional transforms
@torch.compiler.disable
def wrap(wrappee, *, like, **kwargs):
    """Convert a :class:`torch.Tensor` (``wrappee``) into the same :class:`~torchvision.tv_tensors.TVTensor` subclass as ``like``.

    If ``like`` carries metadata (e.g. ``format``, ``canvas_size``), that
    metadata is copied to the output. Individual metadata fields can be
    overridden via ``kwargs``.

    Subclass authors can override :meth:`~torchvision.tv_tensors.TVTensor.wrap`
    to define how their metadata is propagated.

    Args:
        wrappee (Tensor): The tensor to convert.
        like (:class:`~torchvision.tv_tensors.TVTensor`): The reference.
            ``wrappee`` will be converted into the same subclass as ``like``.
        kwargs: Metadata overrides passed to the subclass's
            :meth:`~torchvision.tv_tensors.TVTensor.wrap` method.
    """
    if (wrap_method := getattr(type(like), "wrap", None)) is not None:
        return wrap_method(wrappee, like, **kwargs)
    return wrappee.as_subclass(type(like))