5532 tensor.shape in tensor.Size (#5533)

Signed-off-by: Wenqi Li <[email protected]>

Fixes #5532 

### Description
returns early if it's torch.Size

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [x] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [x] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

Signed-off-by: Wenqi Li <[email protected]>
Co-authored-by: Nic Ma <[email protected]>

Author: wyli

monai/data/meta_tensor.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import functools
 import warnings
 from copy import deepcopy
 from typing import Any, Sequence
@@ -29,6 +30,24 @@
 __all__ = ["MetaTensor"]
 
 
+@functools.lru_cache(None)
+def _get_named_tuple_like_type(func):
+    if (
+        hasattr(torch, "return_types")
+        and hasattr(func, "__name__")
+        and hasattr(torch.return_types, func.__name__)
+        and isinstance(getattr(torch.return_types, func.__name__), type)
+    ):
+        return getattr(torch.return_types, func.__name__)
+    return None
+
+
+def _not_requiring_metadata(ret):
+    return isinstance(ret, (int, str, bytes, torch.Size, torch.dtype, torch.device, np.ndarray)) or not (
+        isinstance(ret, MetaTensor) or (isinstance(ret, Sequence) and any(isinstance(x, MetaTensor) for x in ret))
+    )
+
+
 class MetaTensor(MetaObj, torch.Tensor):
     """
     Class that inherits from both `torch.Tensor` and `MetaObj`, adding support for metadata.
@@ -253,20 +272,16 @@ def __torch_function__(cls, func, types, args=(), kwargs=None) -> Any:
         # we might have 1 or multiple outputs. Might be MetaTensor, might be something
         # else (e.g., `__repr__` returns a string).
         # Convert to list (if necessary), process, and at end remove list if one was added.
-        if (
-            hasattr(torch, "return_types")
-            and hasattr(func, "__name__")
-            and hasattr(torch.return_types, func.__name__)
-            and isinstance(getattr(torch.return_types, func.__name__), type)
-            and isinstance(ret, getattr(torch.return_types, func.__name__))
-        ):
+        if _not_requiring_metadata(ret):
+            return ret
+        if _get_named_tuple_like_type(func) is not None and isinstance(ret, _get_named_tuple_like_type(func)):
             # for torch.max(torch.tensor(1.0), dim=0), the return type is named-tuple like
             out_items = MetaTensor.update_meta(ret, func, args, kwargs)
             for idx in range(ret.n_fields):
                 ret[idx].meta = out_items[idx].meta
                 ret[idx].applied_operations = out_items[idx].applied_operations
             return ret
-        if isinstance(ret, (str, bytes)) or not isinstance(ret, Sequence):
+        if not isinstance(ret, Sequence):
             ret = [ret]
             unpack = True
         else:

tests/test_meta_tensor.py

@@ -424,6 +424,12 @@ def test_str(self):
         t = MetaTensor([1.0], affine=torch.tensor(1), meta={"fname": "filename"})
         self.assertEqual(str(t), "tensor([1.])")
 
+    def test_shape(self):
+        s = MetaTensor([1])
+        self.assertEqual(s.shape, torch.Size([1]))
+        self.assertEqual(s.size(), torch.Size([1]))
+        self.assertEqual(s.size(0), 1)
+
     def test_astype(self):
         t = MetaTensor([1.0], affine=torch.tensor(1), meta={"fname": "filename"})
         for np_types in ("float32", "np.float32", "numpy.float32", np.float32, float, "int", np.compat.long, np.uint16):