Repository: Project-MONAI/MONAI

Issue:
bug: MetaTensor .shape() returns Tuple instead of torch.Size
MetaTensor .shape() returns tuple instead of torch.Size. If MetaTensor is subclass of torch.Tensor, it should not change it methods

```
import torch
from monai.data import MetaTensor

a = torch.rand(1,2,3)
print(a.shape, type(a.shape)) # prints torch.Size([1, 2, 3]) <class 'torch.Size'>

a_meta = MetaTensor(a)
print(a_meta.shape, type(a_meta.shape)) # prints (1, 2, 3) <class 'tuple'>

print('numel of last two dims', a.shape[1:].numel()) # prints 6
print('numel of last two dims', a_meta.shape[1:].numel()) # crashes

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: 'tuple' object has no attribute 'numel'
```


Discussion:

Resolving patch:
diff --git a/monai/data/meta_tensor.py b/monai/data/meta_tensor.py
--- a/monai/data/meta_tensor.py
+++ b/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:


Test patch:
diff --git a/tests/test_meta_tensor.py b/tests/test_meta_tensor.py
--- a/tests/test_meta_tensor.py
+++ b/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):