fix: add squeeze and detach to comp backend

Signed-off-by: anna-charlotte <[email protected]>
docarray · samsja · Jan 25, 2023 · Jan 17, 2023 · Jan 18, 2023 · Jan 18, 2023
commit c3b44bd4da85285cb4b15bef4431a79507fa2d72
diff --git a/docarray/base_document/mixins/plot.py b/docarray/base_document/mixins/plot.py
@@ -1,6 +1,5 @@
 from typing import TYPE_CHECKING, Any, Optional
 
-import numpy as np
 import rich
 from rich.highlighter import RegexHighlighter
 from rich.tree import Tree
@@ -84,7 +83,6 @@ def __rich_console__(self, console, options):
         id_abbrv = getattr(self, 'id')[:7]
         yield f":page_facing_up: [b]{kls}" f"[/b]: [cyan]{id_abbrv} ...[cyan]"
 
-        import torch
         from rich import box, text
         from rich.table import Table
 
@@ -105,15 +103,15 @@ def __rich_console__(self, console, options):
                 if len(v) > 50:
                     col_2 += f' ... (length: {len(v)})'
                 table.add_row(col_1, text.Text(col_2))
-            elif isinstance(v, (np.ndarray, torch.Tensor)):
-                if isinstance(v, torch.Tensor):
-                    v = v.detach().cpu().numpy()
-                if v.squeeze().ndim == 1 and len(v) < 200:
-                    table.add_row(col_1, ColorBoxArray(v.squeeze()))
+            elif isinstance(v, AbstractTensor):
+                comp = v.get_comp_backend()
+                v_squeezed = comp.squeeze(comp.detach(v))
+                if comp.n_dim(v_squeezed) == 1 and comp.shape(v_squeezed)[0] < 200:
+                    table.add_row(col_1, ColorBoxArray(v_squeezed))
                 else:
                     table.add_row(
                         col_1,
-                        text.Text(f'{type(v)} of shape {v.shape}, dtype: {v.dtype}'),
+                        text.Text(f'{type(v)} of shape {comp.shape(v)}'),
                     )
             elif isinstance(v, (tuple, list)):
                 col_2 = ''
@@ -177,7 +175,8 @@ class ColorBoxArray:
     """
 
     def __init__(self, array: AbstractTensor):
-        self._array = array.get_comp_backend().minmax_normalize(array, (0, 5))
+        comp_be = array.get_comp_backend()
+        self._array = comp_be.minmax_normalize(comp_be.detach(array), (0, 5))
 
     def __rich_console__(
         self, console: 'Console', options: 'ConsoleOptions'

diff --git a/docarray/computation/abstract_comp_backend.py b/docarray/computation/abstract_comp_backend.py
@@ -37,6 +37,14 @@ def n_dim(array: 'TTensor') -> int:
         """
         ...
 
+    @staticmethod
+    @abstractmethod
+    def squeeze(tensor: 'TTensor') -> 'TTensor':
+        """
+        Returns a tensor with all the dimensions of tensor of size 1 removed.
+        """
+        ...
+
     @staticmethod
     @abstractmethod
     def to_numpy(array: 'TTensor') -> 'np.ndarray':
@@ -85,6 +93,17 @@ def reshape(tensor: 'TTensor', shape: Tuple[int, ...]) -> 'TTensor':
         """
         ...
 
+    @staticmethod
+    @abstractmethod
+    def detach(tensor: 'TTensor') -> 'TTensor':
+        """
+        Returns the tensor detached from its current graph.
+
+        :param tensor: tensor to be detached
+        :return: a detached tensor with the same data.
+        """
+        ...
+
     @staticmethod
     @abstractmethod
     def minmax_normalize(

diff --git a/docarray/computation/numpy_backend.py b/docarray/computation/numpy_backend.py
@@ -49,6 +49,13 @@ def to_device(tensor: 'np.ndarray', device: str) -> 'np.ndarray':
     def n_dim(array: 'np.ndarray') -> int:
         return array.ndim
 
+    @staticmethod
+    def squeeze(tensor: 'np.ndarray') -> 'np.ndarray':
+        """
+        Returns a tensor with all the dimensions of tensor of size 1 removed.
+        """
+        return tensor.squeeze()
+
     @staticmethod
     def to_numpy(array: 'np.ndarray') -> 'np.ndarray':
         return array
@@ -85,6 +92,16 @@ def reshape(array: 'np.ndarray', shape: Tuple[int, ...]) -> 'np.ndarray':
         """
         return array.reshape(shape)
 
+    @staticmethod
+    def detach(tensor: 'np.ndarray') -> 'np.ndarray':
+        """
+        Returns the tensor detached from its current graph.
+
+        :param tensor: tensor to be detached
+        :return: a detached tensor with the same data.
+        """
+        return tensor
+
     @staticmethod
     def minmax_normalize(
         tensor: 'np.ndarray',

diff --git a/docarray/computation/torch_backend.py b/docarray/computation/torch_backend.py
@@ -63,6 +63,13 @@ def empty(
     def n_dim(array: 'torch.Tensor') -> int:
         return array.ndim
 
+    @staticmethod
+    def squeeze(tensor: 'torch.Tensor') -> 'torch.Tensor':
+        """
+        Returns a tensor with all the dimensions of tensor of size 1 removed.
+        """
+        return torch.squeeze(tensor)
+
     @staticmethod
     def to_numpy(array: 'torch.Tensor') -> 'np.ndarray':
         return array.cpu().detach().numpy()
@@ -89,6 +96,16 @@ def reshape(tensor: 'torch.Tensor', shape: Tuple[int, ...]) -> 'torch.Tensor':
         """
         return tensor.reshape(shape)
 
+    @staticmethod
+    def detach(tensor: 'torch.Tensor') -> 'torch.Tensor':
+        """
+        Returns the tensor detached from its current graph.
+
+        :param tensor: tensor to be detached
+        :return: a detached tensor with the same data.
+        """
+        return tensor.detach()
+
     @staticmethod
     def minmax_normalize(
         tensor: 'torch.Tensor',

diff --git a/tests/units/computation_backends/numpy_backend/test_basics.py b/tests/units/computation_backends/numpy_backend/test_basics.py
@@ -52,6 +52,12 @@ def test_empty_device():
         NumpyCompBackend.empty((10, 3), device='meta')
 
 
+def test_squeeze():
+    tensor = np.zeros(shape=(1, 1, 3, 1))
+    squeezed = NumpyCompBackend.squeeze(tensor)
+    assert squeezed.shape == (3,)
+
+
 @pytest.mark.parametrize(
     'array,t_range,x_range,result',
     [

diff --git a/tests/units/computation_backends/torch_backend/test_basics.py b/tests/units/computation_backends/torch_backend/test_basics.py
@@ -55,6 +55,12 @@ def test_empty_device():
     assert tensor.device == torch.device('meta')
 
 
+def test_squeeze():
+    tensor = torch.zeros(size=(1, 1, 3, 1))
+    squeezed = TorchCompBackend.squeeze(tensor)
+    assert squeezed.shape == (3,)
+
+
 @pytest.mark.parametrize(
     'array,t_range,x_range,result',
     [