Support custom TVTensor types in wrap() function (registration or like-API)

[gabrielfruet]

🚀 The feature

Problem

The wrap() function currently only supports the built-in TVTensor types from torchvision.
If I make a custom TVTensor, I cannot use wrap() to restore metadata in the same way.

Proposal

Allow users to register custom wrapping logic, for example using functools.singledispatch.
Alternatively, add a method or API like CustomTensor.wrap(tensor, like=obj) so each custom type can define its own wrapping.

Example

class CustomTensor(tv_tensors.TVTensor):
    extra_info: str

    @classmethod
    def _wrap(cls, tensor, *, extra_info):
        out = tensor.as_subclass(cls)
        out.extra_info = extra_info
        return out

# Current wrap does not handle extra_info
custom = CustomTensor(torch.ones(2,2), extra_info="abc")
result = custom + 1
wrapped = tv_tensors.wrap(result, like=custom)  # lost metadata

# 1st Proposal (dispatch)
@wrap.register(CustomTensor)
def _(tensor, *, like, **kwargs):
    return CustomTensor._wrap(tensor, extra_info=like.extra_info)

# 2nd proposal (wrap method)
def wrap(wrappee, *, like, **kwargs):
     return like.wrap(wrapee, like=like) # like.__wrap__(...)

Motivation, pitch

• I want to create my own TVTensor types for domain-specific tasks (like depth maps, segmentation, classification, 3d point clouds etc).
• Ideally, I want to use the same wrap/unwrap workflow as built-in TVTensors so transforms and kernels remain easy and clean.
• Supporting custom TVTensors in wrap() helps keep the code extensible but simple. No breaking changes.
• Some users are already requesting functionality that would be possible by this feature (e.g tv_tensors.wrap doesn't work with subclasses of BoundingBoxes or KeyPoints #9328 )

Alternatives

• Manually re-implementing wrapping logic in every custom transform (hard to maintain)
• Monkeypatching wrap() or patching transforms
• Not supporting custom TVTensor use-cases at all

Additional context

This proposal enables advanced users to extend TVTensors for new applications, following existing patterns like register_kernel.
The implementation can be very minimal at first, using the existing Python singledispatch pattern.
If maintainers accept, I'm happy to draft a small PR or demo.

[NicolasHug]

Hi @gabrielfruet , thanks for the request.

I think what you want is technically feasible, but I'd like to better understand your use-case. Your suggested solution would lead to you being able to rely on tv_tensors.wrap to wrap custom subclasses:

wrapped = tv_tensors.wrap(result, like=custom)  # lost metadata

My question is: could you just use a custom wrap() function to do that? Or do you really need support for tv_tensors.wrap because it is called by other torchvision utils that your code relies on?

[gabrielfruet]

Hi Nicholas! Definitely, i could implement my own wrap.

But for example, when custom subclassing a tv_tensor (e.g BoundingBox) and i want to add another field like the confidence of the annotator, the v2 transforms wouldn't pass through this field, and would instantiate a bounding box instead.

AFAIK, it wouldn't directly dispatch the kernel to the subclass, but it would be easier, given this new wrap implementation, reuse the kernels for custom types.

[NicolasHug]

Can you share a minimal reproducing example of the behavior you're observing that you'd like to change? Sorry, I a little low on time these days, so that would greatly help. My current understanding is that you are hitting something like this

class CustomSubType(TVTensor):
    # some meta-data here




@F.register_kernel(functional="resize", tv_tensor_cls=CustomSubType)
def my_resize(custom, *args, **kwargs):
   # pass

custom = CustomSubType()
F.resize(custom)  # metadata is lost here??

But I'd like to confirm to make sure a) we're addressing your problem adequately, and b) we're not implementing unnecessary complexity within TV.

Thanks!

[nkaenzig]

I have the same issue in a similar usecase, often in my custom transform classes I'd use wrap to make sure that the output type of the transform is as expected (as depending on the tv_tensor type and operations applied (e.g. clone()) the input tv_tensors type can get lost):

class RandRotate90(v2.Transform):
    def __init__(
        self,
        prob: float = 0.1,
        max_k: int = 3,
        spatial_axes: tuple[int, int] = (1, 2),
    ) -> None:
        super().__init__()

        self._rotate = monai_spatial_transforms.RandRotate90(
            prob=prob, max_k=max_k, spatial_axes=spatial_axes
        )

    @override
    def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
        self._rotate.randomize()
        return {}

    @functools.singledispatchmethod
    @override
    def transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
        return inpt

    @transform.register(tv_tensors.Image)
    @transform.register(custom_tv_tensors.Volume)
    @transform.register(tv_tensors.Mask)
    def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
        inpt_rotated = self._rotate(img=inpt, randomize=False)
        return tv_tensors.wrap(inpt_rotated, like=inpt)

In this example, I'm using a custom custom_tv_tensors.Volume type, which subclasses from tv_tensors.Video but adds some custom attributes. Those custom attributes of course are being lost when calling .wrap. Right now tv_tensors.wrap just hardcodes the attribute preservation for BoundingBoxes and KeyPoints, but it would be very handy to have a system here that supports custom tv_tensor classes (e.g. if I implement a _wrap method in my tv_tensors method, couldn't that be called?).

https://github.com/pytorch/vision/blob/main/torchvision/tv_tensors/__init__.py#L29-L37

[gabrielfruet]

If i were to make any example, would most likely be like @nkaenzig 's one.

@NicolasHug WDYT about this?

Supporting custom TVTensor types in wrap() would unlock extensibility for domain-specific tasks. Would also reduce maintenance overhead for people that extensively use the library features.

If you want i can make a POC for that. Do you prefer the wrap method (interface-like API) or singledispatch (registry-like API). Or any other that you may find more suitable for the repo.

[zy1git]

@gabrielfruet Hi, thanks a lot for posting this issue and the discussion. Feel free to make a POC for that. Could you please implement the wrap method?

Thank you!

[gabrielfruet]

Thank you! I'll draft a PR in the next few days.