From 1a6c3c80cd4165120b8378e1c15a51ca92dcfc06 Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Mon, 11 Jul 2022 22:43:42 +0100 Subject: [PATCH 1/7] adds __array_function__ impl. Signed-off-by: Wenqi Li --- monai/data/meta_tensor.py | 21 +++++++++++++++++++++ tests/test_meta_tensor.py | 11 +++++++++++ 2 files changed, 32 insertions(+) diff --git a/monai/data/meta_tensor.py b/monai/data/meta_tensor.py index d0818bd25e..e1ad84cb45 100644 --- a/monai/data/meta_tensor.py +++ b/monai/data/meta_tensor.py @@ -280,6 +280,27 @@ def __torch_function__(cls, func, types, args=(), kwargs=None) -> Any: ret = MetaTensor.update_meta(ret, func, args, kwargs) return ret[0] if unpack else ret + def __array_function__(self, func, types, args, kwargs): + """for numpy Interoperability""" + if not func.__module__.startswith("numpy"): + return NotImplemented + if not any(issubclass(t, MetaTensor) for t in types): + # Defer to any other subclasses that implement __array_function__ + return NotImplemented + _args = list( + convert_data_type(x, output_type=np.ndarray, wrap_sequence=False)[0] + if isinstance(x, (MetaTensor, tuple, list)) + else x + for x in args + ) + _kwargs = { + k: convert_data_type(v, output_type=np.ndarray, wrap_sequence=False)[0] + if isinstance(v, (MetaTensor, tuple, list)) + else v + for k, v in kwargs.items() + } + return func(*_args, **_kwargs) + def get_default_affine(self, dtype=torch.float64) -> torch.Tensor: return torch.eye(4, device=self.device, dtype=dtype) diff --git a/tests/test_meta_tensor.py b/tests/test_meta_tensor.py index f3a6205b94..ebd0e127c6 100644 --- a/tests/test_meta_tensor.py +++ b/tests/test_meta_tensor.py @@ -520,6 +520,17 @@ def test_multiprocessing(self, device=None, dtype=None): self.assertIsInstance(obj, MetaTensor) assert_allclose(obj.as_tensor(), t) + @parameterized.expand(TESTS) + def test_array_function(self, device, _dtype): + a = np.random.RandomState().randn(100, 100) + b = MetaTensor(a, device=device) + assert_allclose(np.sum(a), np.sum(b)) + assert_allclose(np.sum(a, axis=1), np.sum(b, axis=1)) + assert_allclose(np.linalg.qr(a), np.linalg.qr(b)) + c = MetaTensor([1.0, 2.0, 3.0]) + assert_allclose(np.argwhere(c == 1.0).astype(int).tolist(), [[0]]) + assert_allclose(np.concatenate([c, c]), np.asarray([1.0, 2.0, 3.0, 1.0, 2.0, 3.0])) + if __name__ == "__main__": unittest.main() From 03666d9fb4420dfd3a8a92b9d0de62a2f570544f Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Tue, 12 Jul 2022 08:23:56 +0100 Subject: [PATCH 2/7] update based on comments Signed-off-by: Wenqi Li --- monai/data/meta_tensor.py | 28 ++++++++++++---------------- 1 file changed, 12 insertions(+), 16 deletions(-) diff --git a/monai/data/meta_tensor.py b/monai/data/meta_tensor.py index e1ad84cb45..982d45778c 100644 --- a/monai/data/meta_tensor.py +++ b/monai/data/meta_tensor.py @@ -282,23 +282,19 @@ def __torch_function__(cls, func, types, args=(), kwargs=None) -> Any: def __array_function__(self, func, types, args, kwargs): """for numpy Interoperability""" - if not func.__module__.startswith("numpy"): + try: + if not func.__module__.startswith("numpy"): + return NotImplemented + except AttributeError: return NotImplemented - if not any(issubclass(t, MetaTensor) for t in types): - # Defer to any other subclasses that implement __array_function__ - return NotImplemented - _args = list( - convert_data_type(x, output_type=np.ndarray, wrap_sequence=False)[0] - if isinstance(x, (MetaTensor, tuple, list)) - else x - for x in args - ) - _kwargs = { - k: convert_data_type(v, output_type=np.ndarray, wrap_sequence=False)[0] - if isinstance(v, (MetaTensor, tuple, list)) - else v - for k, v in kwargs.items() - } + + def _convert(x): + if isinstance(x, (MetaTensor, torch.Tensor, tuple, list)): + return convert_data_type(x, output_type=np.ndarray, wrap_sequence=False)[0] + return x + + _args = list(map(_convert, args)) + _kwargs = {k: _convert(v) for k, v in kwargs.items()} return func(*_args, **_kwargs) def get_default_affine(self, dtype=torch.float64) -> torch.Tensor: From feca937ea65afe36a45ec38e0dc59da23ad6d4b2 Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Tue, 12 Jul 2022 10:16:25 +0100 Subject: [PATCH 3/7] update ufunc Signed-off-by: Wenqi Li --- monai/data/meta_tensor.py | 29 ++++++++++++++++++++--------- tests/test_meta_tensor.py | 1 + 2 files changed, 21 insertions(+), 9 deletions(-) diff --git a/monai/data/meta_tensor.py b/monai/data/meta_tensor.py index 982d45778c..b323a3f9a7 100644 --- a/monai/data/meta_tensor.py +++ b/monai/data/meta_tensor.py @@ -280,23 +280,34 @@ def __torch_function__(cls, func, types, args=(), kwargs=None) -> Any: ret = MetaTensor.update_meta(ret, func, args, kwargs) return ret[0] if unpack else ret + @staticmethod + def _convert(x): + if isinstance(x, (MetaTensor, torch.Tensor, tuple, list)): + return convert_data_type(x, output_type=np.ndarray, wrap_sequence=False)[0] + return x + def __array_function__(self, func, types, args, kwargs): - """for numpy Interoperability""" + """for numpy Interoperability, so that we can compute ``np.sum(MetaTensor([1.0]))``.""" try: if not func.__module__.startswith("numpy"): return NotImplemented except AttributeError: return NotImplemented - - def _convert(x): - if isinstance(x, (MetaTensor, torch.Tensor, tuple, list)): - return convert_data_type(x, output_type=np.ndarray, wrap_sequence=False)[0] - return x - - _args = list(map(_convert, args)) - _kwargs = {k: _convert(v) for k, v in kwargs.items()} + _args = list(map(MetaTensor._convert, args)) + _kwargs = {k: MetaTensor._convert(v) for k, v in kwargs.items()} return func(*_args, **_kwargs) + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + """for numpy interoperability, so that we can compute ``MetaTensor([1.0]) >= np.asarray([1.0])``.""" + try: + if not type(ufunc).__module__.startswith("numpy"): + return NotImplemented + except AttributeError: + return NotImplemented + _inputs = map(MetaTensor._convert, inputs) + _kwargs = {k: MetaTensor._convert(v) for k, v in kwargs.items()} + return getattr(ufunc, method)(*_inputs, **_kwargs) + def get_default_affine(self, dtype=torch.float64) -> torch.Tensor: return torch.eye(4, device=self.device, dtype=dtype) diff --git a/tests/test_meta_tensor.py b/tests/test_meta_tensor.py index ebd0e127c6..4fc2366cdd 100644 --- a/tests/test_meta_tensor.py +++ b/tests/test_meta_tensor.py @@ -530,6 +530,7 @@ def test_array_function(self, device, _dtype): c = MetaTensor([1.0, 2.0, 3.0]) assert_allclose(np.argwhere(c == 1.0).astype(int).tolist(), [[0]]) assert_allclose(np.concatenate([c, c]), np.asarray([1.0, 2.0, 3.0, 1.0, 2.0, 3.0])) + assert_allclose(c > np.asarray([1.0, 1.0, 1.0]), np.asarray([False, True, True])) if __name__ == "__main__": From c51df8b59818a7e35aafe5587abdab84490647bb Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Tue, 12 Jul 2022 11:06:17 +0100 Subject: [PATCH 4/7] pt1.7 Signed-off-by: Wenqi Li --- monai/data/meta_tensor.py | 5 ++++- tests/test_meta_tensor.py | 8 ++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/monai/data/meta_tensor.py b/monai/data/meta_tensor.py index b323a3f9a7..30afc847c6 100644 --- a/monai/data/meta_tensor.py +++ b/monai/data/meta_tensor.py @@ -298,7 +298,10 @@ def __array_function__(self, func, types, args, kwargs): return func(*_args, **_kwargs) def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): - """for numpy interoperability, so that we can compute ``MetaTensor([1.0]) >= np.asarray([1.0])``.""" + """ + For numpy interoperability, so that we can compute ``MetaTensor([1.0]) >= np.asarray([1.0])``. + This is for pytorch > 1.8. + """ try: if not type(ufunc).__module__.startswith("numpy"): return NotImplemented diff --git a/tests/test_meta_tensor.py b/tests/test_meta_tensor.py index 4fc2366cdd..02fb07d9ec 100644 --- a/tests/test_meta_tensor.py +++ b/tests/test_meta_tensor.py @@ -521,7 +521,7 @@ def test_multiprocessing(self, device=None, dtype=None): assert_allclose(obj.as_tensor(), t) @parameterized.expand(TESTS) - def test_array_function(self, device, _dtype): + def test_array_function(self, device="cpu", _dtype=float): a = np.random.RandomState().randn(100, 100) b = MetaTensor(a, device=device) assert_allclose(np.sum(a), np.sum(b)) @@ -530,7 +530,11 @@ def test_array_function(self, device, _dtype): c = MetaTensor([1.0, 2.0, 3.0]) assert_allclose(np.argwhere(c == 1.0).astype(int).tolist(), [[0]]) assert_allclose(np.concatenate([c, c]), np.asarray([1.0, 2.0, 3.0, 1.0, 2.0, 3.0])) - assert_allclose(c > np.asarray([1.0, 1.0, 1.0]), np.asarray([False, True, True])) + if pytorch_after(1, 8, 1): + assert_allclose(c > np.asarray([1.0, 1.0, 1.0]), np.asarray([False, True, True])) + assert_allclose( + c > torch.as_tensor([1.0, 1.0, 1.0], device=device), torch.as_tensor([False, True, True], device=device) + ) if __name__ == "__main__": From 3a1e35814199f1d2f4bef872cdc372eca0ce2368 Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Tue, 12 Jul 2022 11:50:51 +0100 Subject: [PATCH 5/7] fixes typo Signed-off-by: Wenqi Li --- tests/test_meta_tensor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_meta_tensor.py b/tests/test_meta_tensor.py index 02fb07d9ec..ad6374d994 100644 --- a/tests/test_meta_tensor.py +++ b/tests/test_meta_tensor.py @@ -527,7 +527,7 @@ def test_array_function(self, device="cpu", _dtype=float): assert_allclose(np.sum(a), np.sum(b)) assert_allclose(np.sum(a, axis=1), np.sum(b, axis=1)) assert_allclose(np.linalg.qr(a), np.linalg.qr(b)) - c = MetaTensor([1.0, 2.0, 3.0]) + c = MetaTensor([1.0, 2.0, 3.0], device=device) assert_allclose(np.argwhere(c == 1.0).astype(int).tolist(), [[0]]) assert_allclose(np.concatenate([c, c]), np.asarray([1.0, 2.0, 3.0, 1.0, 2.0, 3.0])) if pytorch_after(1, 8, 1): From 72f672b1827e95e2080b64f9510faa953f8ca305 Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Tue, 12 Jul 2022 11:53:54 +0100 Subject: [PATCH 6/7] fixes tests Signed-off-by: Wenqi Li --- tests/test_meta_tensor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_meta_tensor.py b/tests/test_meta_tensor.py index ad6374d994..838f48b34e 100644 --- a/tests/test_meta_tensor.py +++ b/tests/test_meta_tensor.py @@ -521,13 +521,13 @@ def test_multiprocessing(self, device=None, dtype=None): assert_allclose(obj.as_tensor(), t) @parameterized.expand(TESTS) - def test_array_function(self, device="cpu", _dtype=float): + def test_array_function(self, device="cpu", dtype=float): a = np.random.RandomState().randn(100, 100) b = MetaTensor(a, device=device) assert_allclose(np.sum(a), np.sum(b)) assert_allclose(np.sum(a, axis=1), np.sum(b, axis=1)) assert_allclose(np.linalg.qr(a), np.linalg.qr(b)) - c = MetaTensor([1.0, 2.0, 3.0], device=device) + c = MetaTensor([1.0, 2.0, 3.0], device=device, dtype=dtype) assert_allclose(np.argwhere(c == 1.0).astype(int).tolist(), [[0]]) assert_allclose(np.concatenate([c, c]), np.asarray([1.0, 2.0, 3.0, 1.0, 2.0, 3.0])) if pytorch_after(1, 8, 1): From 54b0155ddd90b491c56142619985dc97c8055dfa Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Tue, 12 Jul 2022 22:18:18 +0100 Subject: [PATCH 7/7] more checking Signed-off-by: Wenqi Li --- monai/data/meta_tensor.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/monai/data/meta_tensor.py b/monai/data/meta_tensor.py index 30afc847c6..58240968f0 100644 --- a/monai/data/meta_tensor.py +++ b/monai/data/meta_tensor.py @@ -307,9 +307,16 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): return NotImplemented except AttributeError: return NotImplemented + if method != "__call__": + return NotImplemented _inputs = map(MetaTensor._convert, inputs) _kwargs = {k: MetaTensor._convert(v) for k, v in kwargs.items()} - return getattr(ufunc, method)(*_inputs, **_kwargs) + if "out" in _kwargs: + return NotImplemented # not supported + try: + return getattr(ufunc, method)(*_inputs, **_kwargs) + except AttributeError: + return NotImplemented def get_default_affine(self, dtype=torch.float64) -> torch.Tensor: return torch.eye(4, device=self.device, dtype=dtype)