From a574d84b3eeba8cb0a8589e8a237249c5c08d700 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 29 Nov 2022 10:50:54 +0100
Subject: [PATCH 01/53] fix: inherit from cls in da

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index ba8f750463f..5ea106cc452 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -29,7 +29,7 @@ def __class_getitem__(cls, item: Type[BaseDocument]):
                 f'DocumentArray[item] item should be a Document not a {item} '
             )
 
-        class _DocumenArrayTyped(DocumentArray):
+        class _DocumenArrayTyped(cls):  # type: ignore
             document_type: Type[BaseDocument] = item
 
         for field in _DocumenArrayTyped.document_type.__fields__.keys():

From 05b90e9a3e934b8b265017bc571319679aaf46dc Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 29 Nov 2022 11:22:43 +0100
Subject: [PATCH 02/53] fix: move class atribute to top torch tensor

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/typing/tensor/torch_tensor.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py
index b564908be75..54253b01999 100644
--- a/docarray/typing/tensor/torch_tensor.py
+++ b/docarray/typing/tensor/torch_tensor.py
@@ -73,6 +73,8 @@ class MyDoc(Document):
 
     """
 
+    __parametrized_meta__ = metaTorchAndNode
+
     @classmethod
     def __get_validators__(cls):
         # one or more validators may be yielded which will be called in the
@@ -80,8 +82,6 @@ def __get_validators__(cls):
         # the value returned from the previous validator
         yield cls.validate
 
-    __parametrized_meta__ = metaTorchAndNode
-
     @classmethod
     def __validate_shape__(cls, t: T, shape: Tuple[int]) -> T:  # type: ignore
         if t.shape == shape:

From df5fe2806685ce35e0e836d774ceeb87897190e3 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Thu, 1 Dec 2022 15:42:11 +0100
Subject: [PATCH 03/53] feat: add stack and unstack mod

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py             | 41 ++++++++++++++++++++++++++++-
 docarray/array/mixins/attribute.py  |  6 +++--
 tests/units/array/test_stack_mod.py | 38 ++++++++++++++++++++++++++
 3 files changed, 82 insertions(+), 3 deletions(-)
 create mode 100644 tests/units/array/test_stack_mod.py

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 5ea106cc452..f0c9d5e1acf 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -1,8 +1,12 @@
-from typing import Iterable, Type
+from collections import defaultdict
+from typing import Dict, Iterable, List, Optional, Type
+
+import torch
 
 from docarray.array.abstract_array import AbstractDocumentArray
 from docarray.array.mixins import GetAttributeArrayMixin, ProtoArrayMixin
 from docarray.document import AnyDocument, BaseDocument, BaseNode
+from docarray.typing import TorchTensor
 
 
 class DocumentArray(
@@ -23,6 +27,14 @@ class DocumentArray(
     def __init__(self, docs: Iterable[BaseDocument]):
         super().__init__(doc_ for doc_ in docs)
 
+        self._tensor_columns: Optional[Dict[str, Optional[TorchTensor]]] = dict()
+
+        for field_name, field in self.document_type.__fields__.items():
+            if issubclass(field.type_, TorchTensor):
+                self._tensor_columns[field_name] = None
+        else:
+            self._tensor_columns = None
+
     def __class_getitem__(cls, item: Type[BaseDocument]):
         if not issubclass(item, BaseDocument):
             raise ValueError(
@@ -44,3 +56,30 @@ def _property_generator(val: str):
         _DocumenArrayTyped.__qualname__ = f'DocumentArray[{item.__name__}]'
 
         return _DocumenArrayTyped
+
+    def stacked(self):
+
+        if self.is_stack():
+            tensors_to_stack: Dict[str, List[TorchTensor]] = defaultdict(list)
+            for doc in self:
+                for tensor_field in self._tensor_columns.keys():
+                    tensors_to_stack[tensor_field].append(getattr(doc, tensor_field))
+                    setattr(doc, tensor_field, None)
+
+            for tensor_field, to_stack in tensors_to_stack.items():
+                self._tensor_columns[tensor_field] = torch.stack(to_stack)
+
+            for i, doc in enumerate(self):
+                for tensor_field in self._tensor_columns.keys():
+                    setattr(doc, tensor_field, self._tensor_columns[tensor_field][i])
+
+    def unstacked(self):
+        if self.is_stack():
+
+            for field in self._tensor_columns.keys():
+                del self._tensor_columns[field]
+
+            self._tensor_columns = None
+
+    def is_stack(self) -> bool:
+        return self._tensor_columns is not None
diff --git a/docarray/array/mixins/attribute.py b/docarray/array/mixins/attribute.py
index 712b5b6524e..796274fbf40 100644
--- a/docarray/array/mixins/attribute.py
+++ b/docarray/array/mixins/attribute.py
@@ -8,7 +8,8 @@ class GetAttributeArrayMixin(AbstractDocumentArray):
     """Helpers that provide attributes getter in bulk"""
 
     def _get_documents_attribute(
-        self, field: str
+        self,
+        field: str,
     ) -> Union[List, AbstractDocumentArray]:
         """Return all values of the fields from all docs this array contains
 
@@ -16,7 +17,6 @@ def _get_documents_attribute(
         :return: Returns a list of the field value for each document
         in the array like container
         """
-
         field_type = self.__class__.document_type._get_nested_document_class(field)
 
         if issubclass(field_type, BaseDocument):
@@ -26,5 +26,7 @@ def _get_documents_attribute(
             return self.__class__.__class_getitem__(field_type)(
                 (getattr(doc, field) for doc in self)
             )
+        elif field in self._tensor_columns.keys():
+            return self._tensor_columns[field]
         else:
             return [getattr(doc, field) for doc in self]
diff --git a/tests/units/array/test_stack_mod.py b/tests/units/array/test_stack_mod.py
new file mode 100644
index 00000000000..aa0fe4b2529
--- /dev/null
+++ b/tests/units/array/test_stack_mod.py
@@ -0,0 +1,38 @@
+import torch
+
+from docarray import Document, DocumentArray
+from docarray.typing import TorchTensor
+
+
+def test_stack():
+    class Image(Document):
+        tensor: TorchTensor[3, 224, 224]
+
+    batch = DocumentArray[Image](
+        [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)]
+    )
+
+    batch.stacked()
+
+    assert (batch._tensor_columns['tensor'] == torch.zeros(10, 3, 224, 224)).all()
+    assert (batch.tensor == torch.zeros(10, 3, 224, 224)).all()
+    assert batch._tensor_columns['tensor'].data_ptr() == batch.tensor.data_ptr()
+
+    for doc, tensor in zip(batch, batch.tensor):
+        assert doc.tensor.data_ptr() == tensor.data_ptr()
+
+
+def test_unstack():
+    class Image(Document):
+        tensor: TorchTensor[3, 224, 224]
+
+    batch = DocumentArray[Image](
+        [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)]
+    )
+
+    batch.stacked()
+
+    batch.unstacked()
+
+    for doc in batch:
+        assert (doc.tensor == torch.zeros(3, 224, 224)).all()

From 915cbf58d6db9f7a9083be6d8ed19ab17da0cde3 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Thu, 1 Dec 2022 16:25:52 +0100
Subject: [PATCH 04/53] feat: add blocker for stack mod

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py             | 33 ++++++++++++++++++++++++-----
 tests/units/array/test_stack_mod.py | 15 +++++++++++++
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index f0c9d5e1acf..a208a696b27 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -1,4 +1,5 @@
 from collections import defaultdict
+from functools import wraps
 from typing import Dict, Iterable, List, Optional, Type
 
 import torch
@@ -9,6 +10,18 @@
 from docarray.typing import TorchTensor
 
 
+def _stack_mode_blocker(func):
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if self.is_stack():
+            raise RuntimeError(
+                f'Cannot call {func.__name__} when the document array is in stack mode'
+            )
+        return func(self, *args, **kwargs)
+
+    return wrapper
+
+
 class DocumentArray(
     list,
     ProtoArrayMixin,
@@ -32,8 +45,6 @@ def __init__(self, docs: Iterable[BaseDocument]):
         for field_name, field in self.document_type.__fields__.items():
             if issubclass(field.type_, TorchTensor):
                 self._tensor_columns[field_name] = None
-        else:
-            self._tensor_columns = None
 
     def __class_getitem__(cls, item: Type[BaseDocument]):
         if not issubclass(item, BaseDocument):
@@ -59,7 +70,7 @@ def _property_generator(val: str):
 
     def stacked(self):
 
-        if self.is_stack():
+        if not (self.is_stack()):
             tensors_to_stack: Dict[str, List[TorchTensor]] = defaultdict(list)
             for doc in self:
                 for tensor_field in self._tensor_columns.keys():
@@ -79,7 +90,19 @@ def unstacked(self):
             for field in self._tensor_columns.keys():
                 del self._tensor_columns[field]
 
-            self._tensor_columns = None
+            self._tensor_columns = dict()
 
     def is_stack(self) -> bool:
-        return self._tensor_columns is not None
+        return self._tensor_columns != dict()
+
+    append = _stack_mode_blocker(list.append)
+    extend = _stack_mode_blocker(list.extend)
+    clear = _stack_mode_blocker(list.clear)
+    insert = _stack_mode_blocker(list.insert)
+    pop = _stack_mode_blocker(list.pop)
+    remove = _stack_mode_blocker(list.remove)
+    reverse = _stack_mode_blocker(list.reverse)
+    sort = _stack_mode_blocker(list.sort)
+
+    def copy(self) -> 'DocumentArray':
+        return super().copy()
diff --git a/tests/units/array/test_stack_mod.py b/tests/units/array/test_stack_mod.py
index aa0fe4b2529..4daab3cce61 100644
--- a/tests/units/array/test_stack_mod.py
+++ b/tests/units/array/test_stack_mod.py
@@ -1,3 +1,4 @@
+import pytest
 import torch
 
 from docarray import Document, DocumentArray
@@ -36,3 +37,17 @@ class Image(Document):
 
     for doc in batch:
         assert (doc.tensor == torch.zeros(3, 224, 224)).all()
+
+
+def test_stack_runtime_error():
+    class Image(Document):
+        tensor: TorchTensor[3, 224, 224]
+
+    batch = DocumentArray[Image](
+        [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)]
+    )
+
+    batch.stacked()
+
+    with pytest.raises(RuntimeError):
+        batch.append([])

From fd7d8a73249e9ef7c8bebd8cc03a9b93f4acafb9 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Thu, 1 Dec 2022 16:36:09 +0100
Subject: [PATCH 05/53] feat: add blocker for stack mod

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index a208a696b27..f4eca4fc6a9 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -103,6 +103,3 @@ def is_stack(self) -> bool:
     remove = _stack_mode_blocker(list.remove)
     reverse = _stack_mode_blocker(list.reverse)
     sort = _stack_mode_blocker(list.sort)
-
-    def copy(self) -> 'DocumentArray':
-        return super().copy()

From 95316af07614d1054b08523d99b567498f3c3db4 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Thu, 1 Dec 2022 17:31:29 +0100
Subject: [PATCH 06/53] fix: fix issue with dict iterationg

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index f4eca4fc6a9..e9b590672ef 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -87,13 +87,21 @@ def stacked(self):
     def unstacked(self):
         if self.is_stack():
 
-            for field in self._tensor_columns.keys():
+            for field in list(self._tensor_columns.keys()):
+                # list needed here otherwise we are modifying the dict while iterating
                 del self._tensor_columns[field]
 
             self._tensor_columns = dict()
 
     def is_stack(self) -> bool:
-        return self._tensor_columns != dict()
+        if self._tensor_columns != dict():
+            return all(
+                self._tensor_columns[field] is not None
+                for field in self._tensor_columns.keys()
+            )
+
+        else:
+            return False
 
     append = _stack_mode_blocker(list.append)
     extend = _stack_mode_blocker(list.extend)

From 70d6430e20a21c69f362a65b53f791dd1ea3b133 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Fri, 2 Dec 2022 08:54:37 +0100
Subject: [PATCH 07/53] refactor: init column in stacked mode entry and not
 init

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 22 +++++++++-------------
 1 file changed, 9 insertions(+), 13 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index e9b590672ef..eca9d342e85 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -40,11 +40,7 @@ class DocumentArray(
     def __init__(self, docs: Iterable[BaseDocument]):
         super().__init__(doc_ for doc_ in docs)
 
-        self._tensor_columns: Optional[Dict[str, Optional[TorchTensor]]] = dict()
-
-        for field_name, field in self.document_type.__fields__.items():
-            if issubclass(field.type_, TorchTensor):
-                self._tensor_columns[field_name] = None
+        self._tensor_columns: Optional[Dict[str, Optional[TorchTensor]]] = None
 
     def __class_getitem__(cls, item: Type[BaseDocument]):
         if not issubclass(item, BaseDocument):
@@ -71,6 +67,13 @@ def _property_generator(val: str):
     def stacked(self):
 
         if not (self.is_stack()):
+
+            self._tensor_columns: Optional[Dict[str, Optional[TorchTensor]]] = dict()
+
+            for field_name, field in self.document_type.__fields__.items():
+                if issubclass(field.type_, TorchTensor):
+                    self._tensor_columns[field_name] = None
+
             tensors_to_stack: Dict[str, List[TorchTensor]] = defaultdict(list)
             for doc in self:
                 for tensor_field in self._tensor_columns.keys():
@@ -94,14 +97,7 @@ def unstacked(self):
             self._tensor_columns = dict()
 
     def is_stack(self) -> bool:
-        if self._tensor_columns != dict():
-            return all(
-                self._tensor_columns[field] is not None
-                for field in self._tensor_columns.keys()
-            )
-
-        else:
-            return False
+        return self._tensor_columns is not None
 
     append = _stack_mode_blocker(list.append)
     extend = _stack_mode_blocker(list.extend)

From b2900400c5c77367195c818c310d28882903b582 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Fri, 2 Dec 2022 08:56:25 +0100
Subject: [PATCH 08/53] fix: fix typo

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py            | 12 ++++++------
 docarray/array/mixins/attribute.py |  2 +-
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index eca9d342e85..f602caf54d1 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -48,21 +48,21 @@ def __class_getitem__(cls, item: Type[BaseDocument]):
                 f'DocumentArray[item] item should be a Document not a {item} '
             )
 
-        class _DocumenArrayTyped(cls):  # type: ignore
+        class _DocumentArrayTyped(cls):  # type: ignore
             document_type: Type[BaseDocument] = item
 
-        for field in _DocumenArrayTyped.document_type.__fields__.keys():
+        for field in _DocumentArrayTyped.document_type.__fields__.keys():
 
             def _property_generator(val: str):
                 return property(lambda self: self._get_documents_attribute(val))
 
-            setattr(_DocumenArrayTyped, field, _property_generator(field))
+            setattr(_DocumentArrayTyped, field, _property_generator(field))
             # this generates property on the fly based on the schema of the item
 
-        _DocumenArrayTyped.__name__ = f'DocumentArray[{item.__name__}]'
-        _DocumenArrayTyped.__qualname__ = f'DocumentArray[{item.__name__}]'
+        _DocumentArrayTyped.__name__ = f'DocumentArray[{item.__name__}]'
+        _DocumentArrayTyped.__qualname__ = f'DocumentArray[{item.__name__}]'
 
-        return _DocumenArrayTyped
+        return _DocumentArrayTyped
 
     def stacked(self):
 
diff --git a/docarray/array/mixins/attribute.py b/docarray/array/mixins/attribute.py
index 796274fbf40..cb03e1b8aff 100644
--- a/docarray/array/mixins/attribute.py
+++ b/docarray/array/mixins/attribute.py
@@ -26,7 +26,7 @@ def _get_documents_attribute(
             return self.__class__.__class_getitem__(field_type)(
                 (getattr(doc, field) for doc in self)
             )
-        elif field in self._tensor_columns.keys():
+        elif self._tensor_columns is not None and field in self._tensor_columns.keys():
             return self._tensor_columns[field]
         else:
             return [getattr(doc, field) for doc in self]

From ac26cda1a3c02aad8b62933c03c93534c9226baf Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Fri, 2 Dec 2022 08:57:52 +0100
Subject: [PATCH 09/53] fix: fix typo

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index f602caf54d1..7e87a0be4f1 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -10,7 +10,7 @@
 from docarray.typing import TorchTensor
 
 
-def _stack_mode_blocker(func):
+def _stacked_mode_blocker(func):
     @wraps(func)
     def wrapper(self, *args, **kwargs):
         if self.is_stack():
@@ -99,11 +99,11 @@ def unstacked(self):
     def is_stack(self) -> bool:
         return self._tensor_columns is not None
 
-    append = _stack_mode_blocker(list.append)
-    extend = _stack_mode_blocker(list.extend)
-    clear = _stack_mode_blocker(list.clear)
-    insert = _stack_mode_blocker(list.insert)
-    pop = _stack_mode_blocker(list.pop)
-    remove = _stack_mode_blocker(list.remove)
-    reverse = _stack_mode_blocker(list.reverse)
-    sort = _stack_mode_blocker(list.sort)
+    append = _stacked_mode_blocker(list.append)
+    extend = _stacked_mode_blocker(list.extend)
+    clear = _stacked_mode_blocker(list.clear)
+    insert = _stacked_mode_blocker(list.insert)
+    pop = _stacked_mode_blocker(list.pop)
+    remove = _stacked_mode_blocker(list.remove)
+    reverse = _stacked_mode_blocker(list.reverse)
+    sort = _stacked_mode_blocker(list.sort)

From f50c27b319ea02c2a2a67105852b7a6723c435a9 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Fri, 2 Dec 2022 09:57:57 +0100
Subject: [PATCH 10/53] feat: add column mode for tensor

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py             | 61 ++++++++++++++++++++---------
 docarray/array/mixins/attribute.py  |  6 +--
 tests/units/array/test_stack_mod.py | 45 ++++++++++++++++++++-
 3 files changed, 88 insertions(+), 24 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 7e87a0be4f1..78f0859e279 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -1,6 +1,6 @@
 from collections import defaultdict
 from functools import wraps
-from typing import Dict, Iterable, List, Optional, Type
+from typing import Dict, Iterable, List, Optional, Type, Union
 
 import torch
 
@@ -40,7 +40,7 @@ class DocumentArray(
     def __init__(self, docs: Iterable[BaseDocument]):
         super().__init__(doc_ for doc_ in docs)
 
-        self._tensor_columns: Optional[Dict[str, Optional[TorchTensor]]] = None
+        self._columns: Optional[Dict[str, Optional[TorchTensor]]] = None
 
     def __class_getitem__(cls, item: Type[BaseDocument]):
         if not issubclass(item, BaseDocument):
@@ -68,36 +68,59 @@ def stacked(self):
 
         if not (self.is_stack()):
 
-            self._tensor_columns: Optional[Dict[str, Optional[TorchTensor]]] = dict()
+            self._columns: Optional[Dict[str, Optional[TorchTensor]]] = dict()
 
             for field_name, field in self.document_type.__fields__.items():
-                if issubclass(field.type_, TorchTensor):
-                    self._tensor_columns[field_name] = None
-
-            tensors_to_stack: Dict[str, List[TorchTensor]] = defaultdict(list)
+                if issubclass(field.type_, TorchTensor) or issubclass(
+                    field.type_, BaseDocument
+                ):
+                    self._columns[field_name] = None
+
+            columns_to_stack: Dict[
+                str, Union[List[TorchTensor], List[BaseDocument]]
+            ] = defaultdict(list)
             for doc in self:
-                for tensor_field in self._tensor_columns.keys():
-                    tensors_to_stack[tensor_field].append(getattr(doc, tensor_field))
-                    setattr(doc, tensor_field, None)
-
-            for tensor_field, to_stack in tensors_to_stack.items():
-                self._tensor_columns[tensor_field] = torch.stack(to_stack)
+                for field_to_stack in self._columns.keys():
+                    columns_to_stack[field_to_stack].append(
+                        getattr(doc, field_to_stack)
+                    )
+                    if issubclass(
+                        self.document_type.__fields__[field_to_stack].type_, TorchTensor
+                    ):
+                        self._columns[field_to_stack] = None
+                    setattr(doc, field_to_stack, None)
+
+            for field_to_stack, to_stack in columns_to_stack.items():
+                type_ = self.document_type.__fields__[field_to_stack].type_
+                if issubclass(type_, TorchTensor):
+                    self._columns[field_to_stack] = torch.stack(to_stack)
+                elif issubclass(type_, BaseDocument):
+                    self._columns[field_to_stack] = DocumentArray[type_](
+                        to_stack
+                    ).stacked()
 
             for i, doc in enumerate(self):
-                for tensor_field in self._tensor_columns.keys():
-                    setattr(doc, tensor_field, self._tensor_columns[tensor_field][i])
+                for field_to_stack in self._columns.keys():
+                    if issubclass(
+                        self.document_type.__fields__[field_to_stack].type_, TorchTensor
+                    ):
+                        setattr(doc, field_to_stack, self._columns[field_to_stack][i])
+
+        return self
 
     def unstacked(self):
         if self.is_stack():
 
-            for field in list(self._tensor_columns.keys()):
+            for field in list(self._columns.keys()):
                 # list needed here otherwise we are modifying the dict while iterating
-                del self._tensor_columns[field]
+                del self._columns[field]
+
+            self._columns = dict()
 
-            self._tensor_columns = dict()
+        return self
 
     def is_stack(self) -> bool:
-        return self._tensor_columns is not None
+        return self._columns is not None
 
     append = _stacked_mode_blocker(list.append)
     extend = _stacked_mode_blocker(list.extend)
diff --git a/docarray/array/mixins/attribute.py b/docarray/array/mixins/attribute.py
index cb03e1b8aff..ad1fb594f5a 100644
--- a/docarray/array/mixins/attribute.py
+++ b/docarray/array/mixins/attribute.py
@@ -19,14 +19,14 @@ def _get_documents_attribute(
         """
         field_type = self.__class__.document_type._get_nested_document_class(field)
 
-        if issubclass(field_type, BaseDocument):
+        if self._columns is not None and field in self._columns.keys():
+            return self._columns[field]
+        elif issubclass(field_type, BaseDocument):
             # calling __class_getitem__ ourselves is a hack otherwise mypy complain
             # most likely a bug in mypy though
             # bug reported here https://github.com/python/mypy/issues/14111
             return self.__class__.__class_getitem__(field_type)(
                 (getattr(doc, field) for doc in self)
             )
-        elif self._tensor_columns is not None and field in self._tensor_columns.keys():
-            return self._tensor_columns[field]
         else:
             return [getattr(doc, field) for doc in self]
diff --git a/tests/units/array/test_stack_mod.py b/tests/units/array/test_stack_mod.py
index 4daab3cce61..7b3e09d6c00 100644
--- a/tests/units/array/test_stack_mod.py
+++ b/tests/units/array/test_stack_mod.py
@@ -15,14 +15,39 @@ class Image(Document):
 
     batch.stacked()
 
-    assert (batch._tensor_columns['tensor'] == torch.zeros(10, 3, 224, 224)).all()
+    assert (batch._columns['tensor'] == torch.zeros(10, 3, 224, 224)).all()
     assert (batch.tensor == torch.zeros(10, 3, 224, 224)).all()
-    assert batch._tensor_columns['tensor'].data_ptr() == batch.tensor.data_ptr()
+    assert batch._columns['tensor'].data_ptr() == batch.tensor.data_ptr()
 
     for doc, tensor in zip(batch, batch.tensor):
         assert doc.tensor.data_ptr() == tensor.data_ptr()
 
 
+def test_stack_mod_nested_document():
+    class Image(Document):
+        tensor: TorchTensor[3, 224, 224]
+
+    class MMdoc(Document):
+        img: Image
+
+    batch = DocumentArray[MMdoc](
+        [MMdoc(img=Image(tensor=torch.zeros(3, 224, 224))) for _ in range(10)]
+    )
+
+    batch.stacked()
+
+    assert (
+        batch._columns['img']._columns['tensor'] == torch.zeros(10, 3, 224, 224)
+    ).all()
+
+    assert (batch.img.tensor == torch.zeros(10, 3, 224, 224)).all()
+
+    assert (
+        batch._columns['img']._columns['tensor'].data_ptr()
+        == batch.img.tensor.data_ptr()
+    )
+
+
 def test_unstack():
     class Image(Document):
         tensor: TorchTensor[3, 224, 224]
@@ -39,6 +64,22 @@ class Image(Document):
         assert (doc.tensor == torch.zeros(3, 224, 224)).all()
 
 
+def test_unstack_nested_document():
+    class Image(Document):
+        tensor: TorchTensor[3, 224, 224]
+
+    class MMdoc(Document):
+        img: Image
+
+    batch = DocumentArray[MMdoc](
+        [MMdoc(img=Image(tensor=torch.zeros(3, 224, 224))) for _ in range(10)]
+    )
+
+    batch.stacked()
+
+    batch.unstacked()
+
+
 def test_stack_runtime_error():
     class Image(Document):
         tensor: TorchTensor[3, 224, 224]

From 869c8ee0def0ee2a3851b8e84ac104cb6f5f2e4c Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Fri, 2 Dec 2022 12:01:48 +0100
Subject: [PATCH 11/53] feat: add decorator

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py             | 31 ++++++++++++++-----
 tests/units/array/test_stack_mod.py | 46 ++++++++++++++++++++++++-----
 2 files changed, 62 insertions(+), 15 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 78f0859e279..f1a844235bf 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -1,4 +1,5 @@
 from collections import defaultdict
+from contextlib import contextmanager
 from functools import wraps
 from typing import Dict, Iterable, List, Optional, Type, Union
 
@@ -13,7 +14,7 @@
 def _stacked_mode_blocker(func):
     @wraps(func)
     def wrapper(self, *args, **kwargs):
-        if self.is_stack():
+        if self.is_stacked():
             raise RuntimeError(
                 f'Cannot call {func.__name__} when the document array is in stack mode'
             )
@@ -64,9 +65,9 @@ def _property_generator(val: str):
 
         return _DocumentArrayTyped
 
-    def stacked(self):
+    def stack(self):
 
-        if not (self.is_stack()):
+        if not (self.is_stacked()):
 
             self._columns: Optional[Dict[str, Optional[TorchTensor]]] = dict()
 
@@ -97,7 +98,7 @@ def stacked(self):
                 elif issubclass(type_, BaseDocument):
                     self._columns[field_to_stack] = DocumentArray[type_](
                         to_stack
-                    ).stacked()
+                    ).stack()
 
             for i, doc in enumerate(self):
                 for field_to_stack in self._columns.keys():
@@ -108,18 +109,32 @@ def stacked(self):
 
         return self
 
-    def unstacked(self):
-        if self.is_stack():
+    @contextmanager
+    def stacked_mode(self):
+        try:
+            yield self.stack()
+        finally:
+            self.unstack()
+
+    @contextmanager
+    def unstacked_mode(self):
+        try:
+            yield self.unstack()
+        finally:
+            self.stack()
+
+    def unstack(self):
+        if self.is_stacked():
 
             for field in list(self._columns.keys()):
                 # list needed here otherwise we are modifying the dict while iterating
                 del self._columns[field]
 
-            self._columns = dict()
+            self._columns = None
 
         return self
 
-    def is_stack(self) -> bool:
+    def is_stacked(self) -> bool:
         return self._columns is not None
 
     append = _stacked_mode_blocker(list.append)
diff --git a/tests/units/array/test_stack_mod.py b/tests/units/array/test_stack_mod.py
index 7b3e09d6c00..43fc0ed00c1 100644
--- a/tests/units/array/test_stack_mod.py
+++ b/tests/units/array/test_stack_mod.py
@@ -13,7 +13,7 @@ class Image(Document):
         [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)]
     )
 
-    batch.stacked()
+    batch.stack()
 
     assert (batch._columns['tensor'] == torch.zeros(10, 3, 224, 224)).all()
     assert (batch.tensor == torch.zeros(10, 3, 224, 224)).all()
@@ -34,7 +34,7 @@ class MMdoc(Document):
         [MMdoc(img=Image(tensor=torch.zeros(3, 224, 224))) for _ in range(10)]
     )
 
-    batch.stacked()
+    batch.stack()
 
     assert (
         batch._columns['img']._columns['tensor'] == torch.zeros(10, 3, 224, 224)
@@ -56,9 +56,9 @@ class Image(Document):
         [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)]
     )
 
-    batch.stacked()
+    batch.stack()
 
-    batch.unstacked()
+    batch.unstack()
 
     for doc in batch:
         assert (doc.tensor == torch.zeros(3, 224, 224)).all()
@@ -75,9 +75,9 @@ class MMdoc(Document):
         [MMdoc(img=Image(tensor=torch.zeros(3, 224, 224))) for _ in range(10)]
     )
 
-    batch.stacked()
+    batch.stack()
 
-    batch.unstacked()
+    batch.unstack()
 
 
 def test_stack_runtime_error():
@@ -88,7 +88,39 @@ class Image(Document):
         [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)]
     )
 
-    batch.stacked()
+    batch.stack()
 
     with pytest.raises(RuntimeError):
         batch.append([])
+
+
+def test_context_stack():
+    class Image(Document):
+        tensor: TorchTensor[3, 224, 224]
+
+    batch = DocumentArray[Image](
+        [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)]
+    )
+
+    assert not (batch.is_stacked())
+    with batch.stacked_mode():
+        assert batch.is_stacked()
+
+    assert not (batch.is_stacked())
+
+
+def test_context_not_stack():
+    class Image(Document):
+        tensor: TorchTensor[3, 224, 224]
+
+    batch = DocumentArray[Image](
+        [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)]
+    )
+
+    batch.stack()
+
+    assert batch.is_stacked()
+    with batch.unstacked_mode():
+        assert not (batch.is_stacked())
+
+    assert batch.is_stacked()

From 7cfbb0ecf515aca1a3a03b67aa928fa8ccdbeede Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Fri, 2 Dec 2022 15:25:28 +0100
Subject: [PATCH 12/53] feat: add setter

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py                       |  9 +++++++
 docarray/array/mixins/attribute.py            | 23 +++++++++++++++---
 .../units/array/test_mixins/test_attribute.py | 24 +++++++++++++++++--
 tests/units/array/test_stack_mod.py           | 15 ++++++++++++
 4 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index f1a844235bf..660ccf1b051 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -65,6 +65,15 @@ def _property_generator(val: str):
 
         return _DocumentArrayTyped
 
+    def __getitem__(self, item):
+        if self.is_stacked():
+            doc = super().__getitem__(item)
+            for field in self._columns.keys():
+                setattr(doc, field, self._columns[field][item])
+            return doc
+        else:
+            return super().__getitem__(item)
+
     def stack(self):
 
         if not (self.is_stacked()):
diff --git a/docarray/array/mixins/attribute.py b/docarray/array/mixins/attribute.py
index ad1fb594f5a..98be374b498 100644
--- a/docarray/array/mixins/attribute.py
+++ b/docarray/array/mixins/attribute.py
@@ -2,15 +2,16 @@
 
 from docarray.array.abstract_array import AbstractDocumentArray
 from docarray.document import BaseDocument
+from docarray.typing import TorchTensor
 
 
 class GetAttributeArrayMixin(AbstractDocumentArray):
     """Helpers that provide attributes getter in bulk"""
 
-    def _get_documents_attribute(
+    def _get_array_attribute(
         self,
         field: str,
-    ) -> Union[List, AbstractDocumentArray]:
+    ) -> Union[List, AbstractDocumentArray, TorchTensor]:
         """Return all values of the fields from all docs this array contains
 
         :param field: name of the fields to extract
@@ -19,7 +20,7 @@ def _get_documents_attribute(
         """
         field_type = self.__class__.document_type._get_nested_document_class(field)
 
-        if self._columns is not None and field in self._columns.keys():
+        if self.is_stacked() and field in self._columns.keys():
             return self._columns[field]
         elif issubclass(field_type, BaseDocument):
             # calling __class_getitem__ ourselves is a hack otherwise mypy complain
@@ -30,3 +31,19 @@ def _get_documents_attribute(
             )
         else:
             return [getattr(doc, field) for doc in self]
+
+    def _set_array_attribute(
+        self,
+        field: str,
+        values: Union[List, AbstractDocumentArray, TorchTensor],
+    ):
+        """Set all document if this DocumentArray with the passed values
+
+        :param field: name of the fields to extract
+        :values: the values to set at the document array level
+        """
+        if self.is_stacked() and field in self._columns.keys():
+            self._columns[field] = values
+        else:
+            for doc, value in zip(self, values):
+                setattr(doc, field, value)
diff --git a/tests/units/array/test_mixins/test_attribute.py b/tests/units/array/test_mixins/test_attribute.py
index 1a168bcd0f0..4472a91ff57 100644
--- a/tests/units/array/test_mixins/test_attribute.py
+++ b/tests/units/array/test_mixins/test_attribute.py
@@ -16,19 +16,39 @@ class Mmdoc(BaseDocument):
         (Mmdoc(text=f'hello{i}', tensor=np.zeros((3, 224, 224))) for i in range(N))
     )
 
-    tensors = da._get_documents_attribute('tensor')
+    tensors = da._get_array_attribute('tensor')
 
     assert len(tensors) == N
     for tensor in tensors:
         assert tensor.shape == (3, 224, 224)
 
-    texts = da._get_documents_attribute('text')
+    texts = da._get_array_attribute('text')
 
     assert len(texts) == N
     for i, text in enumerate(texts):
         assert text == f'hello{i}'
 
 
+def test_set_attributes():
+    class InnerDoc(BaseDocument):
+        text: str
+
+    class Mmdoc(BaseDocument):
+        inner: InnerDoc
+
+    N = 10
+
+    da = DocumentArray[Mmdoc](
+        (Mmdoc(inner=InnerDoc(text=f'hello{i}')) for i in range(N))
+    )
+
+    list_docs = [InnerDoc(text=f'hello{i}') for i in range(N)]
+    da._set_array_attribute('inner', list_docs)
+
+    for doc, list_doc in zip(da, list_docs):
+        assert doc.inner is list_doc
+
+
 def test_get_bulk_attributes():
     class Mmdoc(BaseDocument):
         text: str
diff --git a/tests/units/array/test_stack_mod.py b/tests/units/array/test_stack_mod.py
index 43fc0ed00c1..c43fa68d6bd 100644
--- a/tests/units/array/test_stack_mod.py
+++ b/tests/units/array/test_stack_mod.py
@@ -124,3 +124,18 @@ class Image(Document):
         assert not (batch.is_stacked())
 
     assert batch.is_stacked()
+
+
+def test_stack_setter():
+    class Image(Document):
+        tensor: TorchTensor[3, 224, 224]
+
+    batch = DocumentArray[Image](
+        [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)]
+    )
+
+    batch.stack()
+
+    batch.tensor = torch.ones(10, 3, 224, 224)
+
+    assert (batch.tensor == torch.ones(10, 3, 224, 224)).all()

From 58d53aefb7ae89f7a8be336a4f08205a400c6380 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Mon, 5 Dec 2022 08:16:25 +0100
Subject: [PATCH 13/53] fix: fix getter

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 660ccf1b051..d505e6b146e 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -55,7 +55,7 @@ class _DocumentArrayTyped(cls):  # type: ignore
         for field in _DocumentArrayTyped.document_type.__fields__.keys():
 
             def _property_generator(val: str):
-                return property(lambda self: self._get_documents_attribute(val))
+                return property(lambda self: self._get_array_attribute(val))
 
             setattr(_DocumentArrayTyped, field, _property_generator(field))
             # this generates property on the fly based on the schema of the item

From 8ddf6182f071fe3c5dfa8fb971bcb58be0613969 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Mon, 5 Dec 2022 08:23:00 +0100
Subject: [PATCH 14/53] feat: add setter as property

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index d505e6b146e..2484ba15a6b 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -55,7 +55,14 @@ class _DocumentArrayTyped(cls):  # type: ignore
         for field in _DocumentArrayTyped.document_type.__fields__.keys():
 
             def _property_generator(val: str):
-                return property(lambda self: self._get_array_attribute(val))
+                def _getter(self):
+                    return self._get_array_attribute(val)
+
+                def _setter(self, value):
+                    self._set_array_attribute(val, value)
+
+                # need docstring for the property
+                return property(fget=_getter, fset=_setter)
 
             setattr(_DocumentArrayTyped, field, _property_generator(field))
             # this generates property on the fly based on the schema of the item

From 4ce3a4b73c0c20200bf2a9e47555d5f0efdb0397 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Mon, 5 Dec 2022 09:41:00 +0100
Subject: [PATCH 15/53] feat: add more test

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py                      |  1 +
 tests/integrations/array/test_torch_train.py | 35 ++++++++++++++++++++
 tests/units/array/test_stack_mod.py          | 16 +++++++++
 3 files changed, 52 insertions(+)
 create mode 100644 tests/integrations/array/test_torch_train.py

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 2484ba15a6b..16dacad069b 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -75,6 +75,7 @@ def _setter(self, value):
     def __getitem__(self, item):
         if self.is_stacked():
             doc = super().__getitem__(item)
+            # NOTE: this could be speed up by using a cache
             for field in self._columns.keys():
                 setattr(doc, field, self._columns[field][item])
             return doc
diff --git a/tests/integrations/array/test_torch_train.py b/tests/integrations/array/test_torch_train.py
new file mode 100644
index 00000000000..0b25bd190f0
--- /dev/null
+++ b/tests/integrations/array/test_torch_train.py
@@ -0,0 +1,35 @@
+from typing import Optional
+
+import torch
+
+from docarray import Document, DocumentArray
+from docarray.typing import TorchTensor
+
+
+def test_torch_train():
+    class Mmdoc(Document):
+        text: str
+        tensor: Optional[TorchTensor[3, 224, 224]]
+
+    N = 10
+
+    batch = DocumentArray[Mmdoc](Mmdoc(text=f'hello{i}') for i in range(N))
+
+    batch.tensor = torch.zeros(N, 3, 224, 224)
+    batch.stack()
+
+    class Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.conv = torch.nn.Conv2d(3, 16, 3)
+
+        def forward(self, x):
+            return self.conv(x)
+
+    model = Model()
+    opt = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
+
+    for _ in range(2):
+        loss = model(batch.tensor).sum()
+        loss.backward()
+        opt.step()
diff --git a/tests/units/array/test_stack_mod.py b/tests/units/array/test_stack_mod.py
index c43fa68d6bd..f26120416a4 100644
--- a/tests/units/array/test_stack_mod.py
+++ b/tests/units/array/test_stack_mod.py
@@ -1,3 +1,5 @@
+from typing import Optional
+
 import pytest
 import torch
 
@@ -139,3 +141,17 @@ class Image(Document):
     batch.tensor = torch.ones(10, 3, 224, 224)
 
     assert (batch.tensor == torch.ones(10, 3, 224, 224)).all()
+
+
+def test_stack_optional():
+    class Image(Document):
+        tensor: Optional[TorchTensor[3, 224, 224]]
+
+    batch = DocumentArray[Image](
+        [Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)]
+    )
+
+    batch.stack()
+
+    assert (batch._columns['tensor'] == torch.zeros(10, 3, 224, 224)).all()
+    assert (batch.tensor == torch.zeros(10, 3, 224, 224)).all()

From 9c181f6ec28352ea7fff2ab2073301316fa28027 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Mon, 5 Dec 2022 10:23:10 +0100
Subject: [PATCH 16/53] feat: add numpy to stack mod

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py             | 29 ++++++++++++++++-------------
 tests/units/array/test_stack_mod.py | 20 +++++++++++++++++++-
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 16dacad069b..66007a509db 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -3,12 +3,13 @@
 from functools import wraps
 from typing import Dict, Iterable, List, Optional, Type, Union
 
+import numpy as np
 import torch
 
 from docarray.array.abstract_array import AbstractDocumentArray
 from docarray.array.mixins import GetAttributeArrayMixin, ProtoArrayMixin
 from docarray.document import AnyDocument, BaseDocument, BaseNode
-from docarray.typing import TorchTensor
+from docarray.typing import NdArray, TorchTensor
 
 
 def _stacked_mode_blocker(func):
@@ -41,7 +42,9 @@ class DocumentArray(
     def __init__(self, docs: Iterable[BaseDocument]):
         super().__init__(doc_ for doc_ in docs)
 
-        self._columns: Optional[Dict[str, Optional[TorchTensor]]] = None
+        self._columns: Optional[
+            Dict[str, Union[TorchTensor, DocumentArray, NdArray, None]]
+        ] = None
 
     def __class_getitem__(cls, item: Type[BaseDocument]):
         if not issubclass(item, BaseDocument):
@@ -86,32 +89,33 @@ def stack(self):
 
         if not (self.is_stacked()):
 
-            self._columns: Optional[Dict[str, Optional[TorchTensor]]] = dict()
+            self._columns = dict()
 
             for field_name, field in self.document_type.__fields__.items():
-                if issubclass(field.type_, TorchTensor) or issubclass(
-                    field.type_, BaseDocument
+                if (
+                    issubclass(field.type_, TorchTensor)
+                    or issubclass(field.type_, BaseDocument)
+                    or issubclass(field.type_, NdArray)
                 ):
                     self._columns[field_name] = None
 
             columns_to_stack: Dict[
-                str, Union[List[TorchTensor], List[BaseDocument]]
+                str, Union[List[TorchTensor], List[NdArray], List[BaseDocument]]
             ] = defaultdict(list)
+
             for doc in self:
                 for field_to_stack in self._columns.keys():
                     columns_to_stack[field_to_stack].append(
                         getattr(doc, field_to_stack)
                     )
-                    if issubclass(
-                        self.document_type.__fields__[field_to_stack].type_, TorchTensor
-                    ):
-                        self._columns[field_to_stack] = None
                     setattr(doc, field_to_stack, None)
 
             for field_to_stack, to_stack in columns_to_stack.items():
                 type_ = self.document_type.__fields__[field_to_stack].type_
                 if issubclass(type_, TorchTensor):
                     self._columns[field_to_stack] = torch.stack(to_stack)
+                if issubclass(type_, NdArray):
+                    self._columns[field_to_stack] = np.stack(to_stack)
                 elif issubclass(type_, BaseDocument):
                     self._columns[field_to_stack] = DocumentArray[type_](
                         to_stack
@@ -119,9 +123,8 @@ def stack(self):
 
             for i, doc in enumerate(self):
                 for field_to_stack in self._columns.keys():
-                    if issubclass(
-                        self.document_type.__fields__[field_to_stack].type_, TorchTensor
-                    ):
+                    type_ = self.document_type.__fields__[field_to_stack].type_
+                    if issubclass(type_, TorchTensor) or issubclass(type_, NdArray):
                         setattr(doc, field_to_stack, self._columns[field_to_stack][i])
 
         return self
diff --git a/tests/units/array/test_stack_mod.py b/tests/units/array/test_stack_mod.py
index f26120416a4..35ea9ffc786 100644
--- a/tests/units/array/test_stack_mod.py
+++ b/tests/units/array/test_stack_mod.py
@@ -1,10 +1,11 @@
 from typing import Optional
 
+import numpy as np
 import pytest
 import torch
 
 from docarray import Document, DocumentArray
-from docarray.typing import TorchTensor
+from docarray.typing import NdArray, TorchTensor
 
 
 def test_stack():
@@ -155,3 +156,20 @@ class Image(Document):
 
     assert (batch._columns['tensor'] == torch.zeros(10, 3, 224, 224)).all()
     assert (batch.tensor == torch.zeros(10, 3, 224, 224)).all()
+
+
+def test_stack_numpy():
+    class Image(Document):
+        tensor: NdArray[3, 224, 224]
+
+    batch = DocumentArray[Image](
+        [Image(tensor=np.zeros((3, 224, 224))) for _ in range(10)]
+    )
+
+    batch.stack()
+
+    assert (batch._columns['tensor'] == np.zeros((10, 3, 224, 224))).all()
+    assert (batch.tensor == np.zeros((10, 3, 224, 224))).all()
+    assert batch.tensor.ctypes.data == batch._columns['tensor'].ctypes.data
+
+    batch.unstack()

From 4f73b9efe6f035fced47e853fdffe6bbc9ae57c1 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Mon, 5 Dec 2022 10:44:23 +0100
Subject: [PATCH 17/53] feat: add stack classmethod

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py                   | 14 ++++++++------
 docarray/typing/tensor/abstract_tensor.py |  7 ++++++-
 docarray/typing/tensor/ndarray.py         |  6 ++++++
 docarray/typing/tensor/torch_tensor.py    | 18 +++++++++++++++++-
 4 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 66007a509db..2de20763cb6 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -1,10 +1,7 @@
 from collections import defaultdict
 from contextlib import contextmanager
 from functools import wraps
-from typing import Dict, Iterable, List, Optional, Type, Union
-
-import numpy as np
-import torch
+from typing import Dict, Iterable, List, Optional, Type, TypeVar, Union
 
 from docarray.array.abstract_array import AbstractDocumentArray
 from docarray.array.mixins import GetAttributeArrayMixin, ProtoArrayMixin
@@ -24,6 +21,9 @@ def wrapper(self, *args, **kwargs):
     return wrapper
 
 
+T = TypeVar('T', bound='DocumentArray')
+
+
 class DocumentArray(
     list,
     ProtoArrayMixin,
@@ -113,9 +113,11 @@ def stack(self):
             for field_to_stack, to_stack in columns_to_stack.items():
                 type_ = self.document_type.__fields__[field_to_stack].type_
                 if issubclass(type_, TorchTensor):
-                    self._columns[field_to_stack] = torch.stack(to_stack)
+                    self._columns[field_to_stack] = TorchTensor.__docarray_stack__(
+                        to_stack
+                    )
                 if issubclass(type_, NdArray):
-                    self._columns[field_to_stack] = np.stack(to_stack)
+                    self._columns[field_to_stack] = NdArray.__docarray_stack__(to_stack)
                 elif issubclass(type_, BaseDocument):
                     self._columns[field_to_stack] = DocumentArray[type_](
                         to_stack
diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py
index a792f329896..286d73f3f3b 100644
--- a/docarray/typing/tensor/abstract_tensor.py
+++ b/docarray/typing/tensor/abstract_tensor.py
@@ -1,6 +1,6 @@
 import abc
 from abc import ABC
-from typing import TYPE_CHECKING, Any, Generic, Tuple, Type, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, Sequence, Tuple, Type, TypeVar
 
 from docarray.typing.abstract_type import AbstractType
 
@@ -91,3 +91,8 @@ def validate(
     def __class_getitem__(cls, item: Any):
         target_shape = cls.__validate_getitem__(item)
         return cls._create_parametrized_type(target_shape)
+
+    @classmethod
+    def __docarray_stack__(cls, seq: Sequence[T]) -> T:
+        """Stack a sequence of tensors into a single tensor."""
+        ...
diff --git a/docarray/typing/tensor/ndarray.py b/docarray/typing/tensor/ndarray.py
index 1e0ee9174fd..2fafca367a1 100644
--- a/docarray/typing/tensor/ndarray.py
+++ b/docarray/typing/tensor/ndarray.py
@@ -5,6 +5,7 @@
     Dict,
     Generic,
     List,
+    Sequence,
     Tuple,
     Type,
     TypeVar,
@@ -191,3 +192,8 @@ def _flush_tensor_to_proto(pb_msg: 'NdArrayProto', value: 'NdArray'):
         pb_msg.dense.ClearField('shape')
         pb_msg.dense.shape.extend(list(value.shape))
         pb_msg.dense.dtype = value.dtype.str
+
+    @classmethod
+    def __docarray_stack__(cls, seq: Sequence[T]) -> T:
+        """Stack a sequence of ndarray into a single ndarray."""
+        return cls.from_ndarray(np.stack(seq))
diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py
index 54253b01999..b4e17fef009 100644
--- a/docarray/typing/tensor/torch_tensor.py
+++ b/docarray/typing/tensor/torch_tensor.py
@@ -1,6 +1,17 @@
 import warnings
 from copy import copy
-from typing import TYPE_CHECKING, Any, Dict, Generic, Tuple, Type, TypeVar, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generic,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
 
 import numpy as np
 import torch  # type: ignore
@@ -212,3 +223,8 @@ def _flush_tensor_to_proto(pb_msg: 'NdArrayProto', value: 'TorchTensor'):
         pb_msg.dense.ClearField('shape')
         pb_msg.dense.shape.extend(list(value_np.shape))
         pb_msg.dense.dtype = value_np.dtype.str
+
+    @classmethod
+    def __docarray_stack__(cls, seq: Sequence[T]) -> T:
+        """Stack a sequence of ndarray into a single ndarray."""
+        return cls.from_native_torch_tensor(torch.stack(seq))

From 9cec782ef6fe99dd2d3f49425de3e882a51b725f Mon Sep 17 00:00:00 2001
From: Johannes Messner <44071807+JohannesMessner@users.noreply.github.com>
Date: Thu, 1 Dec 2022 14:54:53 +0100
Subject: [PATCH 18/53] feat: embedding type (#877)

* refactor: make abstract tensor parametrization modular

Signed-off-by: Johannes Messner <messnerjo@gmail.com>

* feat: add type for torch embedding

Signed-off-by: Johannes Messner <messnerjo@gmail.com>

* feat: embedding type for ndarray

Signed-off-by: Johannes Messner <messnerjo@gmail.com>

* fix: fix general embedding type

Signed-off-by: Johannes Messner <messnerjo@gmail.com>

* test: update tests

Signed-off-by: Johannes Messner <messnerjo@gmail.com>

* fix: typing issue

Signed-off-by: Johannes Messner <messnerjo@gmail.com>

* fix: metaclass conflict

Signed-off-by: Johannes Messner <messnerjo@gmail.com>

* docs: clarify docstring

Signed-off-by: Johannes Messner <messnerjo@gmail.com>

* test: split test into two

Signed-off-by: Johannes Messner <messnerjo@gmail.com>

Signed-off-by: Johannes Messner <messnerjo@gmail.com>
Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/typing/tensor/abstract_tensor.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py
index 286d73f3f3b..a9f070ac4d5 100644
--- a/docarray/typing/tensor/abstract_tensor.py
+++ b/docarray/typing/tensor/abstract_tensor.py
@@ -88,11 +88,17 @@ def validate(
 
         return _ParametrizedTensor
 
+<<<<<<< HEAD
     def __class_getitem__(cls, item: Any):
         target_shape = cls.__validate_getitem__(item)
         return cls._create_parametrized_type(target_shape)
 
+=======
+>>>>>>> feat: embedding type (#877)
     @classmethod
     def __docarray_stack__(cls, seq: Sequence[T]) -> T:
         """Stack a sequence of tensors into a single tensor."""
         ...
+    def __class_getitem__(cls, item: Any):
+        target_shape = cls.__validate_getitem__(item)
+        return cls._create_parametrized_type(target_shape)

From 3506501d1dc47ad1ea912b1a64bd5fe26b326c20 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Mon, 5 Dec 2022 11:36:48 +0100
Subject: [PATCH 19/53] refactor: halfway trough mypy not complaining

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/abstract_array.py          | 15 ++++++++++++++-
 docarray/array/array.py                   |  2 +-
 docarray/array/mixins/attribute.py        | 10 ++++++----
 docarray/typing/tensor/abstract_tensor.py |  5 +++--
 docarray/typing/tensor/ndarray.py         |  3 +--
 docarray/typing/tensor/torch_tensor.py    |  4 ++--
 6 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/docarray/array/abstract_array.py b/docarray/array/abstract_array.py
index 4812d69d4ca..be722f9a3e9 100644
--- a/docarray/array/abstract_array.py
+++ b/docarray/array/abstract_array.py
@@ -1,12 +1,18 @@
 from abc import abstractmethod
-from typing import Iterable, Type
+from typing import TYPE_CHECKING, Dict, Iterable, Optional, Type, Union
 
 from docarray.document import BaseDocument
 
+if TYPE_CHECKING:
+    from docarray.typing import NdArray, TorchTensor
+
 
 class AbstractDocumentArray(Iterable):
 
     document_type: Type[BaseDocument]
+    _columns: Optional[
+        Dict[str, Union['TorchTensor', 'AbstractDocumentArray', 'NdArray', None]]
+    ]
 
     @abstractmethod
     def __init__(self, docs: Iterable[BaseDocument]):
@@ -17,3 +23,10 @@ def __class_getitem__(
         cls, item: Type[BaseDocument]
     ) -> Type['AbstractDocumentArray']:
         ...
+
+    @abstractmethod
+    def is_stacked(self) -> bool:
+        """
+        Return True if the document array is in stack mode
+        """
+        ...
diff --git a/docarray/array/array.py b/docarray/array/array.py
index 2de20763cb6..e98abb7b9c2 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -43,7 +43,7 @@ def __init__(self, docs: Iterable[BaseDocument]):
         super().__init__(doc_ for doc_ in docs)
 
         self._columns: Optional[
-            Dict[str, Union[TorchTensor, DocumentArray, NdArray, None]]
+            Dict[str, Union[TorchTensor, AbstractDocumentArray, NdArray, None]]
         ] = None
 
     def __class_getitem__(cls, item: Type[BaseDocument]):
diff --git a/docarray/array/mixins/attribute.py b/docarray/array/mixins/attribute.py
index 98be374b498..5a633daa997 100644
--- a/docarray/array/mixins/attribute.py
+++ b/docarray/array/mixins/attribute.py
@@ -1,8 +1,10 @@
-from typing import List, Union
+from typing import TYPE_CHECKING, List, Union
 
 from docarray.array.abstract_array import AbstractDocumentArray
 from docarray.document import BaseDocument
-from docarray.typing import TorchTensor
+
+if TYPE_CHECKING:
+    from docarray.typing import NdArray, TorchTensor
 
 
 class GetAttributeArrayMixin(AbstractDocumentArray):
@@ -11,7 +13,7 @@ class GetAttributeArrayMixin(AbstractDocumentArray):
     def _get_array_attribute(
         self,
         field: str,
-    ) -> Union[List, AbstractDocumentArray, TorchTensor]:
+    ) -> Union[List, AbstractDocumentArray, 'TorchTensor', 'NdArray']:
         """Return all values of the fields from all docs this array contains
 
         :param field: name of the fields to extract
@@ -35,7 +37,7 @@ def _get_array_attribute(
     def _set_array_attribute(
         self,
         field: str,
-        values: Union[List, AbstractDocumentArray, TorchTensor],
+        values: Union[List, AbstractDocumentArray, 'TorchTensor', 'NdArray'],
     ):
         """Set all document if this DocumentArray with the passed values
 
diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py
index a9f070ac4d5..42d13da4497 100644
--- a/docarray/typing/tensor/abstract_tensor.py
+++ b/docarray/typing/tensor/abstract_tensor.py
@@ -1,6 +1,6 @@
 import abc
 from abc import ABC
-from typing import TYPE_CHECKING, Any, Generic, Sequence, Tuple, Type, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, List, Tuple, Type, TypeVar, Union
 
 from docarray.typing.abstract_type import AbstractType
 
@@ -96,7 +96,8 @@ def __class_getitem__(cls, item: Any):
 =======
 >>>>>>> feat: embedding type (#877)
     @classmethod
-    def __docarray_stack__(cls, seq: Sequence[T]) -> T:
+    @abc.abstractmethod
+    def __docarray_stack__(cls: Type[T], seq: Union[List[T], Tuple[T]]) -> T:
         """Stack a sequence of tensors into a single tensor."""
         ...
     def __class_getitem__(cls, item: Any):
diff --git a/docarray/typing/tensor/ndarray.py b/docarray/typing/tensor/ndarray.py
index 2fafca367a1..24509aa02ee 100644
--- a/docarray/typing/tensor/ndarray.py
+++ b/docarray/typing/tensor/ndarray.py
@@ -5,7 +5,6 @@
     Dict,
     Generic,
     List,
-    Sequence,
     Tuple,
     Type,
     TypeVar,
@@ -194,6 +193,6 @@ def _flush_tensor_to_proto(pb_msg: 'NdArrayProto', value: 'NdArray'):
         pb_msg.dense.dtype = value.dtype.str
 
     @classmethod
-    def __docarray_stack__(cls, seq: Sequence[T]) -> T:
+    def __docarray_stack__(cls: Type[T], seq: Union[List[T], Tuple[T]]) -> T:
         """Stack a sequence of ndarray into a single ndarray."""
         return cls.from_ndarray(np.stack(seq))
diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py
index b4e17fef009..29a7db41f1f 100644
--- a/docarray/typing/tensor/torch_tensor.py
+++ b/docarray/typing/tensor/torch_tensor.py
@@ -5,7 +5,7 @@
     Any,
     Dict,
     Generic,
-    Sequence,
+    List,
     Tuple,
     Type,
     TypeVar,
@@ -225,6 +225,6 @@ def _flush_tensor_to_proto(pb_msg: 'NdArrayProto', value: 'TorchTensor'):
         pb_msg.dense.dtype = value_np.dtype.str
 
     @classmethod
-    def __docarray_stack__(cls, seq: Sequence[T]) -> T:
+    def __docarray_stack__(cls: Type[T], seq: Union[List[T], Tuple[T]]) -> T:
         """Stack a sequence of ndarray into a single ndarray."""
         return cls.from_native_torch_tensor(torch.stack(seq))

From f14885717a3abed4e501e187035b52acf21997e7 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Mon, 5 Dec 2022 14:09:05 +0100
Subject: [PATCH 20/53] refactor: almost the end of mypy

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/abstract_array.py   |  9 ++++----
 docarray/array/array.py            | 36 +++++++++++++++++++++---------
 docarray/array/mixins/attribute.py | 23 ++++++++++++++++---
 3 files changed, 50 insertions(+), 18 deletions(-)

diff --git a/docarray/array/abstract_array.py b/docarray/array/abstract_array.py
index be722f9a3e9..5f21f7cdc03 100644
--- a/docarray/array/abstract_array.py
+++ b/docarray/array/abstract_array.py
@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Dict, Iterable, Optional, Type, Union
+from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Type, Union
 
 from docarray.document import BaseDocument
 
@@ -26,7 +26,8 @@ def __class_getitem__(
 
     @abstractmethod
     def is_stacked(self) -> bool:
-        """
-        Return True if the document array is in stack mode
-        """
+        ...
+
+    @abstractmethod
+    def _column_fields(self) -> List[str]:
         ...
diff --git a/docarray/array/array.py b/docarray/array/array.py
index e98abb7b9c2..4e885578d9b 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -131,6 +131,17 @@ def stack(self):
 
         return self
 
+    def unstack(self):
+        if self.is_stacked():
+
+            for field in list(self._columns.keys()):
+                # list needed here otherwise we are modifying the dict while iterating
+                del self._columns[field]
+
+            self._columns = None
+
+        return self
+
     @contextmanager
     def stacked_mode(self):
         try:
@@ -145,20 +156,23 @@ def unstacked_mode(self):
         finally:
             self.stack()
 
-    def unstack(self):
-        if self.is_stacked():
-
-            for field in list(self._columns.keys()):
-                # list needed here otherwise we are modifying the dict while iterating
-                del self._columns[field]
-
-            self._columns = None
-
-        return self
-
     def is_stacked(self) -> bool:
+        """
+        Return True if the document array is in stack mode
+        """
         return self._columns is not None
 
+    def _column_fields(self) -> List[str]:
+        """
+        return the field store as columns
+        :return: the list of keys of the columns
+        """
+        if self.is_stacked() and self._columns is not None:
+            # need to repeat myself here bc of mypy
+            return list(self._columns.keys())
+        else:
+            return []
+
     append = _stacked_mode_blocker(list.append)
     extend = _stacked_mode_blocker(list.extend)
     clear = _stacked_mode_blocker(list.clear)
diff --git a/docarray/array/mixins/attribute.py b/docarray/array/mixins/attribute.py
index 5a633daa997..83f64a6eadf 100644
--- a/docarray/array/mixins/attribute.py
+++ b/docarray/array/mixins/attribute.py
@@ -22,8 +22,20 @@ def _get_array_attribute(
         """
         field_type = self.__class__.document_type._get_nested_document_class(field)
 
-        if self.is_stacked() and field in self._columns.keys():
-            return self._columns[field]
+        if (
+            self.is_stacked()
+            and field in self._column_fields()
+            and self._columns is not None
+        ):
+            attributes = self._columns[field]
+            if attributes is not None:
+                return attributes
+            else:
+                raise ValueError(
+                    f'The column is not set for the field {field} even though '
+                    f'it is in stacked mode'
+                )
+
         elif issubclass(field_type, BaseDocument):
             # calling __class_getitem__ ourselves is a hack otherwise mypy complain
             # most likely a bug in mypy though
@@ -44,7 +56,12 @@ def _set_array_attribute(
         :param field: name of the fields to extract
         :values: the values to set at the document array level
         """
-        if self.is_stacked() and field in self._columns.keys():
+        if (
+            self.is_stacked()
+            and self._columns is not None
+            and field in self._column_fields()
+            and not isinstance(values, List)
+        ):
             self._columns[field] = values
         else:
             for doc, value in zip(self, values):

From e37608dea355fb0d903971b8fa09712fb1a365ab Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Mon, 5 Dec 2022 15:08:33 +0100
Subject: [PATCH 21/53] refactor: mypy good

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/typing/tensor/ndarray.py      | 4 +++-
 docarray/typing/tensor/torch_tensor.py | 6 ++++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/docarray/typing/tensor/ndarray.py b/docarray/typing/tensor/ndarray.py
index 24509aa02ee..ad002187328 100644
--- a/docarray/typing/tensor/ndarray.py
+++ b/docarray/typing/tensor/ndarray.py
@@ -193,6 +193,8 @@ def _flush_tensor_to_proto(pb_msg: 'NdArrayProto', value: 'NdArray'):
         pb_msg.dense.dtype = value.dtype.str
 
     @classmethod
-    def __docarray_stack__(cls: Type[T], seq: Union[List[T], Tuple[T]]) -> T:
+    def __docarray_stack__(
+        cls: Type[T], seq: Union[List['NdArray'], Tuple['NdArray']]
+    ) -> T:
         """Stack a sequence of ndarray into a single ndarray."""
         return cls.from_ndarray(np.stack(seq))
diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py
index 29a7db41f1f..47522fd834c 100644
--- a/docarray/typing/tensor/torch_tensor.py
+++ b/docarray/typing/tensor/torch_tensor.py
@@ -225,6 +225,8 @@ def _flush_tensor_to_proto(pb_msg: 'NdArrayProto', value: 'TorchTensor'):
         pb_msg.dense.dtype = value_np.dtype.str
 
     @classmethod
-    def __docarray_stack__(cls: Type[T], seq: Union[List[T], Tuple[T]]) -> T:
+    def __docarray_stack__(
+        cls: Type[T], seq: Union[Tuple['TorchTensor'], List['TorchTensor']]
+    ) -> T:
         """Stack a sequence of ndarray into a single ndarray."""
-        return cls.from_native_torch_tensor(torch.stack(seq))
+        return cls.from_native_torch_tensor(torch.stack(seq))  # type: ignore

From 94b3b08f8cb2c3da3e7a3fc98dbd9030e4a5cadf Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Mon, 5 Dec 2022 16:11:22 +0100
Subject: [PATCH 22/53] chore: add docstring for stack mode

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 108 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 108 insertions(+)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 4e885578d9b..68c53b5003a 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -86,6 +86,62 @@ def __getitem__(self, item):
             return super().__getitem__(item)
 
     def stack(self):
+        """
+        Calling this method will make the DocumentArray enter in stacked mode
+
+        This mode will stack all the Tensor and Nested fields of the documents in the
+        array into columns. This is useful when you want to perform operations on the
+        whole array at once. Especially when it involves working on Tensor directly
+        where you want to avoid for loops and work on batch directly.
+
+        By default, a DocumentArray is not in stack mode, it means that each Document
+        holds its own data and that the DocumentArray is just a python list of
+        Documents. When you access fields of the DocumentArray you actually loop over
+        the list of Documents and access the field of each Document.
+
+        In stack mode, accessing or setting fields of the DocumentArray will access or
+        set the column of the array.
+
+        When entering stack mode the DocumentArray will create a column for each field
+        of the Document that are Tensor or that are Nested Document that contains at
+        least one Tensor field.
+
+        IMPORTANT: in stacked_mode you cannot add or remove Document from the array.
+        You can only modify the fields of the DocumentArray. This is intentionally done
+        to avoid inconsistencies between the columns and the list of Documents and
+        because adding to a column is slow. If you need to do so please use the unstack
+        method to exit stacked mode.
+
+        EXAMPLE USAGE
+        .. code-block:: python
+            from docarray import Document, DocumentArray
+            from docarray.typing import NdArray
+
+
+            class Image(Document):
+                tensor: NdArray[100]
+
+
+            batch = DocumentArray[Image]([Image(tensor=np.zeros((100))) \
+                for _ in range(10)])
+
+            batch.stack()
+
+            print(batch[0].tensor[0])
+            # >>> 0
+
+            print(batch.tensor.shape)
+            # >>> (10, 3, 224, 224)
+
+            batch.tensor = np.ones((10, 100))
+
+            print(batch[0].tensor[0])
+            # >>> 1
+
+            batch.append(Image(tensor=np.zeros((100))))
+            # >>> raise RuntimeError('Cannot call append when the document array is in
+            # >>> stack mode'
+        """
 
         if not (self.is_stacked()):
 
@@ -132,6 +188,58 @@ def stack(self):
         return self
 
     def unstack(self):
+        """
+        Calling this method will make the DocumentArray enter in unstack mode.
+        This is the default mode of any DocumentArray. This means you should call this
+        only method if the DocumentArray is already in stacked mode. (Calling it while
+        being in unstacked mode will have no effect)
+
+        Unstack will unstack all the columns of the DocumentArray and put the data back
+        in each Document of the DocumentArray.
+
+        In unstacked mode DocumentArray behaves like a normal python list of Documents.
+        This means that you can call append, insert, index etc without problem.
+
+
+        You can still access DocumentArray fields in unstacked mode. This will return
+        a list of the fields of each Document in the array (or A DocumentArray if the
+        field contains Document and not type). This list/DocumentsArray will be
+        created on the fly and will not be stored in the DocumentArray. If you want to
+        do heavy operations on the fields of the DocumentArray you should use the stack
+        mode.
+
+        You can use setter as well to set the fields of the DocumentArray. This will
+        iterate over the list/DocumentsArray/Ndarray and set the field of each Document
+        accordingly.
+
+        EXAMPLE USAGE
+        .. code-block:: python
+            from docarray import Document, DocumentArray, Text
+            from docarray.typing import NdArray
+
+
+            class Image(Document):
+                tensor: NdArray[100]
+
+            batch = DocumentArray[Image]([Image(tensor=np.zeros((100))) \
+                for _ in range(10)])
+
+            batch.stack()
+            batch.unstack()
+            print(type(batch.tensor))
+            # >>> list
+            print(len(batch.tensor))
+            # >>> 10
+
+            batch.append(Image(tensor=np.zeros((100))))
+            # this work just fine
+
+            batch.tensor = np.ones((10, 100))
+            # this iterate over the ndarray and assign the row to each Document
+
+            print(batch[0].tensor[0])
+            # >>> 1
+        """
         if self.is_stacked():
 
             for field in list(self._columns.keys()):

From f3d91c7a98c4d501b71cc49f5f1aef7c244550a1 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Mon, 5 Dec 2022 17:36:28 +0100
Subject: [PATCH 23/53] chore: add docstring for document array

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 67 +++++++++++++++++++++++++++++++++++++++--
 1 file changed, 64 insertions(+), 3 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 68c53b5003a..98af5af6597 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -32,9 +32,70 @@ class DocumentArray(
     BaseNode,
 ):
     """
-    a DocumentArray is a list-like container of Document of the same schema
+    a DocumentArray is a container of Document.
 
     :param docs: iterable of Document
+
+
+    A DocumentArray can only contain Document that follow the same schema. To precise
+    this schema you can use the `DocumentArray[Document]` syntax. This will create a
+    DocumentArray that can only contain Document of the type Document. (Note that there
+    exist a special schema (AnySchema) that allow any Document to be stored in the
+    DocumentArray, but this is not recommended to use and exist here mostly for
+    serialization purpose with protobuf when the receiver does not know the schema of
+    DocumentArray but still want to deserialize it).
+
+    EXAMPLE USAGE
+    .. code-block:: python
+        from docarray import Document, DocumentArray
+        from docarray.typing import NdArray, ImageUrl
+
+        class Image(Document):
+            tensor: Optional[NdArray[100]]
+            url: ImageUrl
+
+        da = DocumentArray[Image](Image(url='http://url.com') for _ in range(10))
+
+
+    DocumentArray define setter and getter for each field of the Document schema. These
+    getter and setter are defined dynamically at runtime. This allows to access the
+    field of the Document in a natural way. For example, if you have a DocumentArray of
+    Image you can do: `da.tensor` to get the tensor of all the Image in the
+    DocumentArray. You can also do `da.tensor = np.random.random([10, 100])` to set the
+    tensor of all the Image.
+
+
+    DocumentArray can be in two mode: row and stack.
+
+
+    By default, it is in row mode (or unstacked mode), a DocumentArray is almost just a
+    list of Document. You can append Document to it, iterate over it, etc. Each Document
+    owns its data. You can see it as a row based datastructure. In this case
+    the getter and setter shown above will return a list of the field of each Document
+    (or a DocumentArray if the field is a nested Document) This list/DocumentArray will
+    be created in the fly. The setter will set the field of each Document to the value
+    of the list/DocumentArray/Tensor passed as parameters.
+
+    Nevertheless, this list-like behavior is not always optimal especially when you want
+    to process you data in batch and do operation which involves matrix computation,
+    like in deep learning or to compute the cosine similarly of embeddings, in this case
+    you want to stack all the tensor in a single batch. This is where the stack mode of
+    the DocumentArray comes in handy.
+    In stacked mode each field which are Tensor are stored as a column in a tensor of
+    batch the size of the DocumentArray. This allows to do operation on the whole batch
+    instead of iterating over the DocumentArray.
+    In this mode the Document inside in the Document don't own the data anymore but just
+    reference to the data in the tensor. But For the user they are no difference it
+    looks and feel the same.
+    In stacked mode the getter and setter just replace the tensor of the DocumentArray
+    of the given field.
+    Finally, in stacked mode operation like `da.append` are not allowed anymore because
+    they are too slow and not recommended to use. You should rather use the unstacked
+    mode.
+
+    To switch from stacked mode to unstacked mode you need to call `da.unstack()` and
+    `da.stack`. There are as well two context manager to for these modes.
+    `with da.stack_mode():` and `with da.unstack_mode():`
     """
 
     document_type: Type[BaseDocument] = AnyDocument
@@ -194,8 +255,8 @@ def unstack(self):
         only method if the DocumentArray is already in stacked mode. (Calling it while
         being in unstacked mode will have no effect)
 
-        Unstack will unstack all the columns of the DocumentArray and put the data back
-        in each Document of the DocumentArray.
+        Calling unstack will unstack all the columns of the DocumentArray and put the
+        data back in each Document of the DocumentArray.
 
         In unstacked mode DocumentArray behaves like a normal python list of Documents.
         This means that you can call append, insert, index etc without problem.

From 64b856b38aaf949e1eb2f4b7b759ed752a3362cc Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Mon, 5 Dec 2022 17:44:50 +0100
Subject: [PATCH 24/53] chore: add docstring

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 98af5af6597..97622aa6078 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -96,6 +96,8 @@ class Image(Document):
     To switch from stacked mode to unstacked mode you need to call `da.unstack()` and
     `da.stack`. There are as well two context manager to for these modes.
     `with da.stack_mode():` and `with da.unstack_mode():`
+
+    see {meth}`.stack` and {meth}`.unstack` for more information.
     """
 
     document_type: Type[BaseDocument] = AnyDocument
@@ -202,6 +204,8 @@ class Image(Document):
             batch.append(Image(tensor=np.zeros((100))))
             # >>> raise RuntimeError('Cannot call append when the document array is in
             # >>> stack mode'
+
+        see {meth}`.unstack` for more information on how to switch to unstack mode
         """
 
         if not (self.is_stacked()):
@@ -300,6 +304,8 @@ class Image(Document):
 
             print(batch[0].tensor[0])
             # >>> 1
+
+        see {meth}`.stack` for more information on how to switch to stack mode
         """
         if self.is_stacked():
 

From 4297ab210df20e278a8b51e0053706666be6123a Mon Sep 17 00:00:00 2001
From: samsja <55492238+samsja@users.noreply.github.com>
Date: Mon, 5 Dec 2022 18:06:58 +0100
Subject: [PATCH 25/53] fix: apply johannes suggestion

Co-authored-by: Johannes Messner <44071807+JohannesMessner@users.noreply.github.com>
Signed-off-by: samsja <55492238+samsja@users.noreply.github.com>
Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 97622aa6078..8df6852dbe3 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -37,13 +37,11 @@ class DocumentArray(
     :param docs: iterable of Document
 
 
-    A DocumentArray can only contain Document that follow the same schema. To precise
+    A DocumentArray can only contain Documents that follow the same schema. To precise
     this schema you can use the `DocumentArray[Document]` syntax. This will create a
     DocumentArray that can only contain Document of the type Document. (Note that there
-    exist a special schema (AnySchema) that allow any Document to be stored in the
-    DocumentArray, but this is not recommended to use and exist here mostly for
-    serialization purpose with protobuf when the receiver does not know the schema of
-    DocumentArray but still want to deserialize it).
+    exists a special schema (AnySchema) that allows any Document to be stored in the
+    DocumentArray, but this is not recommended to use).
 
     EXAMPLE USAGE
     .. code-block:: python
@@ -57,15 +55,15 @@ class Image(Document):
         da = DocumentArray[Image](Image(url='http://url.com') for _ in range(10))
 
 
-    DocumentArray define setter and getter for each field of the Document schema. These
-    getter and setter are defined dynamically at runtime. This allows to access the
+    DocumentArray defines setter and getter for each field of the Document schema. These
+    getters and setters are defined dynamically at runtime. This allows to access the
     field of the Document in a natural way. For example, if you have a DocumentArray of
     Image you can do: `da.tensor` to get the tensor of all the Image in the
     DocumentArray. You can also do `da.tensor = np.random.random([10, 100])` to set the
     tensor of all the Image.
 
 
-    DocumentArray can be in two mode: row and stack.
+    A DocumentArray can be in one of two modes: row mode and stacked mode.
 
 
     By default, it is in row mode (or unstacked mode), a DocumentArray is almost just a
@@ -150,7 +148,7 @@ def __getitem__(self, item):
 
     def stack(self):
         """
-        Calling this method will make the DocumentArray enter in stacked mode
+        Calling this method will make the DocumentArray enter stacked mode
 
         This mode will stack all the Tensor and Nested fields of the documents in the
         array into columns. This is useful when you want to perform operations on the

From 52a5e695d9c51f4222aaa66e7c73573d9ccb01f0 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Mon, 5 Dec 2022 18:13:25 +0100
Subject: [PATCH 26/53] fix: apply johannes suggestion

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/abstract_array.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docarray/array/abstract_array.py b/docarray/array/abstract_array.py
index 5f21f7cdc03..a93a4e6a288 100644
--- a/docarray/array/abstract_array.py
+++ b/docarray/array/abstract_array.py
@@ -12,7 +12,7 @@ class AbstractDocumentArray(Iterable):
     document_type: Type[BaseDocument]
     _columns: Optional[
         Dict[str, Union['TorchTensor', 'AbstractDocumentArray', 'NdArray', None]]
-    ]
+    ]  # here columns are the holder of the data in tensor modes
 
     @abstractmethod
     def __init__(self, docs: Iterable[BaseDocument]):

From 3b91d619f8129a70a18bb651838940576cf6c150 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Mon, 5 Dec 2022 18:16:36 +0100
Subject: [PATCH 27/53] fix: apply johannes suggestion

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 8df6852dbe3..1eb902598fd 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -48,10 +48,12 @@ class DocumentArray(
         from docarray import Document, DocumentArray
         from docarray.typing import NdArray, ImageUrl
 
+
         class Image(Document):
             tensor: Optional[NdArray[100]]
             url: ImageUrl
 
+
         da = DocumentArray[Image](Image(url='http://url.com') for _ in range(10))
 
 
@@ -184,7 +186,7 @@ class Image(Document):
 
 
             batch = DocumentArray[Image]([Image(tensor=np.zeros((100))) \
-                for _ in range(10)])
+                        for _ in range(10)])
 
             batch.stack()
 
@@ -235,7 +237,7 @@ class Image(Document):
                     self._columns[field_to_stack] = TorchTensor.__docarray_stack__(
                         to_stack
                     )
-                if issubclass(type_, NdArray):
+                elif issubclass(type_, NdArray):
                     self._columns[field_to_stack] = NdArray.__docarray_stack__(to_stack)
                 elif issubclass(type_, BaseDocument):
                     self._columns[field_to_stack] = DocumentArray[type_](
@@ -284,8 +286,9 @@ def unstack(self):
             class Image(Document):
                 tensor: NdArray[100]
 
+
             batch = DocumentArray[Image]([Image(tensor=np.zeros((100))) \
-                for _ in range(10)])
+                    for _ in range(10)])
 
             batch.stack()
             batch.unstack()

From 6866bc62612f9cf02d46c2940add20c6756e1fc0 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 10:12:57 +0100
Subject: [PATCH 28/53] fix: shorten docstring

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 74 +++++++++++++----------------------------
 1 file changed, 24 insertions(+), 50 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 1eb902598fd..96b0b8258ba 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -65,31 +65,28 @@ class Image(Document):
     tensor of all the Image.
 
 
-    A DocumentArray can be in one of two modes: row mode and stacked mode.
+    A DocumentArray can be in one of two modes: unstacked mode and stacked mode.
 
 
-    By default, it is in row mode (or unstacked mode), a DocumentArray is almost just a
-    list of Document. You can append Document to it, iterate over it, etc. Each Document
-    owns its data. You can see it as a row based datastructure. In this case
-    the getter and setter shown above will return a list of the field of each Document
-    (or a DocumentArray if the field is a nested Document) This list/DocumentArray will
-    be created in the fly. The setter will set the field of each Document to the value
+    **Unstacked mode (default)**:
+    In this case a DocumentArray is a list of Document and each Document owns its data.
+    The getter and setter shown above will return a list of the fields of each Document
+    (or a DocumentArray if the field is a nested Document). This list/DocumentArray will
+    be created on the fly. The setter will set the field of each Document to the value
     of the list/DocumentArray/Tensor passed as parameters.
 
     Nevertheless, this list-like behavior is not always optimal especially when you want
-    to process you data in batch and do operation which involves matrix computation,
-    like in deep learning or to compute the cosine similarly of embeddings, in this case
-    you want to stack all the tensor in a single batch. This is where the stack mode of
-    the DocumentArray comes in handy.
-    In stacked mode each field which are Tensor are stored as a column in a tensor of
-    batch the size of the DocumentArray. This allows to do operation on the whole batch
-    instead of iterating over the DocumentArray.
+    to process you data in batch and do operation which involves matrix computation,This
+    is where the stack mode of the DocumentArray comes in handy.
+
+    In **stacked mode** each field which are Tensor are stored as a column in a tensor
+    of batch the size of the DocumentArray. This allows to do operation on the whole
+    batch instead of iterating over the DocumentArray.
     In this mode the Document inside in the Document don't own the data anymore but just
-    reference to the data in the tensor. But For the user they are no difference it
-    looks and feel the same.
+    reference to the data in the tensor of the DocumentArray.
     In stacked mode the getter and setter just replace the tensor of the DocumentArray
     of the given field.
-    Finally, in stacked mode operation like `da.append` are not allowed anymore because
+    Operation like `da.append` are not allowed anymore because
     they are too slow and not recommended to use. You should rather use the unstacked
     mode.
 
@@ -150,30 +147,22 @@ def __getitem__(self, item):
 
     def stack(self):
         """
-        Calling this method will make the DocumentArray enter stacked mode
-
-        This mode will stack all the Tensor and Nested fields of the documents in the
-        array into columns. This is useful when you want to perform operations on the
-        whole array at once. Especially when it involves working on Tensor directly
-        where you want to avoid for loops and work on batch directly.
-
-        By default, a DocumentArray is not in stack mode, it means that each Document
-        holds its own data and that the DocumentArray is just a python list of
-        Documents. When you access fields of the DocumentArray you actually loop over
-        the list of Documents and access the field of each Document.
-
-        In stack mode, accessing or setting fields of the DocumentArray will access or
-        set the column of the array.
+        Calling this method will make the DocumentArray enter stacked mode. You should
+        call this method only if the DocumentArray is already in unstacked mode
+        (the default mode). (Calling it while being already in stacked mode will have
+        no effect)
 
         When entering stack mode the DocumentArray will create a column for each field
         of the Document that are Tensor or that are Nested Document that contains at
-        least one Tensor field.
+        least one Tensor field. This is useful when you want to perform operations on
+        the whole array at once. In stack mode, accessing or setting fields of the
+        DocumentArray will access or set the column of the array.
+
 
         IMPORTANT: in stacked_mode you cannot add or remove Document from the array.
         You can only modify the fields of the DocumentArray. This is intentionally done
-        to avoid inconsistencies between the columns and the list of Documents and
-        because adding to a column is slow. If you need to do so please use the unstack
-        method to exit stacked mode.
+        because extending or inserting a column is slow. You should rather use the
+        unstack mode for such operation.
 
         EXAMPLE USAGE
         .. code-block:: python
@@ -262,21 +251,6 @@ def unstack(self):
         Calling unstack will unstack all the columns of the DocumentArray and put the
         data back in each Document of the DocumentArray.
 
-        In unstacked mode DocumentArray behaves like a normal python list of Documents.
-        This means that you can call append, insert, index etc without problem.
-
-
-        You can still access DocumentArray fields in unstacked mode. This will return
-        a list of the fields of each Document in the array (or A DocumentArray if the
-        field contains Document and not type). This list/DocumentsArray will be
-        created on the fly and will not be stored in the DocumentArray. If you want to
-        do heavy operations on the fields of the DocumentArray you should use the stack
-        mode.
-
-        You can use setter as well to set the fields of the DocumentArray. This will
-        iterate over the list/DocumentsArray/Ndarray and set the field of each Document
-        accordingly.
-
         EXAMPLE USAGE
         .. code-block:: python
             from docarray import Document, DocumentArray, Text

From c80168960e1d661b5e573c2dd383c694ccbd5757 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 10:28:18 +0100
Subject: [PATCH 29/53] fix: apply johannes suggestion

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 96b0b8258ba..e60b3c86e81 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -32,7 +32,7 @@ class DocumentArray(
     BaseNode,
 ):
     """
-    a DocumentArray is a container of Document.
+    a DocumentArray is a container of Document of the same schema.
 
     :param docs: iterable of Document
 

From 93e18927621bea689b9374d4477d7586c39356f9 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 10:42:53 +0100
Subject: [PATCH 30/53] refactor: apply johannes suggestion

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index e60b3c86e81..4729d27585d 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -175,7 +175,7 @@ class Image(Document):
 
 
             batch = DocumentArray[Image]([Image(tensor=np.zeros((100))) \
-                        for _ in range(10)])
+                    for _ in range(10)])
 
             batch.stack()
 
@@ -221,17 +221,14 @@ class Image(Document):
                     setattr(doc, field_to_stack, None)
 
             for field_to_stack, to_stack in columns_to_stack.items():
+
                 type_ = self.document_type.__fields__[field_to_stack].type_
-                if issubclass(type_, TorchTensor):
-                    self._columns[field_to_stack] = TorchTensor.__docarray_stack__(
-                        to_stack
-                    )
-                elif issubclass(type_, NdArray):
-                    self._columns[field_to_stack] = NdArray.__docarray_stack__(to_stack)
-                elif issubclass(type_, BaseDocument):
+                if issubclass(type_, BaseDocument):
                     self._columns[field_to_stack] = DocumentArray[type_](
                         to_stack
                     ).stack()
+                else:
+                    self._columns[field_to_stack] = type_.__docarray_stack__(to_stack)
 
             for i, doc in enumerate(self):
                 for field_to_stack in self._columns.keys():
@@ -262,7 +259,7 @@ class Image(Document):
 
 
             batch = DocumentArray[Image]([Image(tensor=np.zeros((100))) \
-                    for _ in range(10)])
+                        for _ in range(10)])
 
             batch.stack()
             batch.unstack()

From 73d2fda13f426486ef4f4375ff40215cf074d0a6 Mon Sep 17 00:00:00 2001
From: samsja <55492238+samsja@users.noreply.github.com>
Date: Tue, 6 Dec 2022 10:54:53 +0100
Subject: [PATCH 31/53] fix: apply johannes suggestion

Co-authored-by: Johannes Messner <44071807+JohannesMessner@users.noreply.github.com>
Signed-off-by: samsja <55492238+samsja@users.noreply.github.com>
Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/mixins/attribute.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docarray/array/mixins/attribute.py b/docarray/array/mixins/attribute.py
index 83f64a6eadf..e206767f720 100644
--- a/docarray/array/mixins/attribute.py
+++ b/docarray/array/mixins/attribute.py
@@ -32,8 +32,8 @@ def _get_array_attribute(
                 return attributes
             else:
                 raise ValueError(
-                    f'The column is not set for the field {field} even though '
-                    f'it is in stacked mode'
+                    f'The DocumentArray is in stacked mode, but no stacked data is '
+                    f'present for {field}. This is inconsistent'
                 )
 
         elif issubclass(field_type, BaseDocument):

From bcaad73dbcc2815e05fca87b5d9919240ee5064c Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 11:07:07 +0100
Subject: [PATCH 32/53] feat: add warning in list like funciton

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 4729d27585d..f2c0039a39b 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -18,6 +18,11 @@ def wrapper(self, *args, **kwargs):
             )
         return func(self, *args, **kwargs)
 
+    (
+        wrapper.__doc__
+        + ' \n This method is not available in stacked mode call {meth}.unstack() '
+        'to use it'
+    )
     return wrapper
 
 
@@ -174,8 +179,9 @@ class Image(Document):
                 tensor: NdArray[100]
 
 
-            batch = DocumentArray[Image]([Image(tensor=np.zeros((100))) \
-                    for _ in range(10)])
+            batch = DocumentArray[Image](
+                [Image(tensor=np.zeros((100))) for _ in range(10)]
+            )  # noqa: E510
 
             batch.stack()
 
@@ -258,8 +264,9 @@ class Image(Document):
                 tensor: NdArray[100]
 
 
-            batch = DocumentArray[Image]([Image(tensor=np.zeros((100))) \
-                        for _ in range(10)])
+            batch = DocumentArray[Image](
+                [Image(tensor=np.zeros((100))) for _ in range(10)]
+            )  # noqa: E510
 
             batch.stack()
             batch.unstack()

From ce346896278ec3c2db3830dca8acd728dabfca9c Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 11:08:49 +0100
Subject: [PATCH 33/53] fix: fix docstring column fields

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index f2c0039a39b..6e05bde2845 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -318,7 +318,7 @@ def is_stacked(self) -> bool:
 
     def _column_fields(self) -> List[str]:
         """
-        return the field store as columns
+        return the list of fields that are columns of the DocumentArray
         :return: the list of keys of the columns
         """
         if self.is_stacked() and self._columns is not None:

From 1f57835f2b98fa91ec379e97f7e65d00f19fde2b Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 11:16:33 +0100
Subject: [PATCH 34/53] feat: add warning in list like funciton

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 6e05bde2845..aa45ca950d0 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -18,11 +18,12 @@ def wrapper(self, *args, **kwargs):
             )
         return func(self, *args, **kwargs)
 
-    (
-        wrapper.__doc__
-        + ' \n This method is not available in stacked mode call {meth}.unstack() '
-        'to use it'
-    )
+        wrapper.__doc__ = (
+            wrapper.__doc__
+            + ' \n This method is not available in stacked mode call {meth}.unstack() '
+            + 'to use it'
+        )
+
     return wrapper
 
 

From a6abd5e062799ad5c82907a8a4f67bca086e6b01 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 11:35:19 +0100
Subject: [PATCH 35/53] feat: apply jihannes suggestion

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 23 ++++++++++-------------
 1 file changed, 10 insertions(+), 13 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index aa45ca950d0..031216c3c27 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -38,16 +38,16 @@ class DocumentArray(
     BaseNode,
 ):
     """
-    a DocumentArray is a container of Document of the same schema.
+    A DocumentArray is a container of Documents.
 
     :param docs: iterable of Document
 
-
-    A DocumentArray can only contain Documents that follow the same schema. To precise
-    this schema you can use the `DocumentArray[Document]` syntax. This will create a
-    DocumentArray that can only contain Document of the type Document. (Note that there
-    exists a special schema (AnySchema) that allows any Document to be stored in the
-    DocumentArray, but this is not recommended to use).
+    A DocumentArray is a list of Documents of any schema. Nevertheless, most of the
+    features of DocumentArray are only available if the contains Documents are
+    homogeneous and follow the same schema. To precise this schema you can use the
+    `DocumentArray[MyDocument]` syntax where MyDocument is a Document class
+    (i,e a schema).This will create a DocumentArray that can only contain Document of
+    the type Document.
 
     EXAMPLE USAGE
     .. code-block:: python
@@ -63,12 +63,9 @@ class Image(Document):
         da = DocumentArray[Image](Image(url='http://url.com') for _ in range(10))
 
 
-    DocumentArray defines setter and getter for each field of the Document schema. These
-    getters and setters are defined dynamically at runtime. This allows to access the
-    field of the Document in a natural way. For example, if you have a DocumentArray of
-    Image you can do: `da.tensor` to get the tensor of all the Image in the
-    DocumentArray. You can also do `da.tensor = np.random.random([10, 100])` to set the
-    tensor of all the Image.
+    If your DocumentArray is homogeneous (i.e follow the same schema) You can access the
+    field at the DocumentArray level. Example `da.tensor`. You can set them as well.
+    by doing `da.tensor = np.random.random([10, 100])`
 
 
     A DocumentArray can be in one of two modes: unstacked mode and stacked mode.

From 834fada5833d46c7e7b6ad66b64435a0a0bd799a Mon Sep 17 00:00:00 2001
From: samsja <55492238+samsja@users.noreply.github.com>
Date: Tue, 6 Dec 2022 11:36:10 +0100
Subject: [PATCH 36/53] feat: apply aleaddine suggestion

Co-authored-by: Johannes Messner <44071807+JohannesMessner@users.noreply.github.com>
Signed-off-by: samsja <55492238+samsja@users.noreply.github.com>
Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 031216c3c27..6bf369c63fb 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -72,19 +72,18 @@ class Image(Document):
 
 
     **Unstacked mode (default)**:
-    In this case a DocumentArray is a list of Document and each Document owns its data.
+    In this case a DocumentArray is a list of Documents and each Document owns its data.
     The getter and setter shown above will return a list of the fields of each Document
     (or a DocumentArray if the field is a nested Document). This list/DocumentArray will
     be created on the fly. The setter will set the field of each Document to the value
     of the list/DocumentArray/Tensor passed as parameters.
 
-    Nevertheless, this list-like behavior is not always optimal especially when you want
-    to process you data in batch and do operation which involves matrix computation,This
+    This list-like behavior is not always optimal especially when you want
+    to process you data in batches or do operations which involves matrix computation. This
     is where the stack mode of the DocumentArray comes in handy.
 
-    In **stacked mode** each field which are Tensor are stored as a column in a tensor
-    of batch the size of the DocumentArray. This allows to do operation on the whole
-    batch instead of iterating over the DocumentArray.
+    In **stacked mode** tensor-like fields of every Document are stored on the DocumentArray level, as one stacked tensor. This enables operations on the the entire
+    batch without iterating over the DocumentArray.
     In this mode the Document inside in the Document don't own the data anymore but just
     reference to the data in the tensor of the DocumentArray.
     In stacked mode the getter and setter just replace the tensor of the DocumentArray

From bc766138bd8c5d36d80efa154205493528fdff7a Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 11:38:23 +0100
Subject: [PATCH 37/53] feat: apply jihannes suggestion

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 20 +++++++-------------
 1 file changed, 7 insertions(+), 13 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 6bf369c63fb..6cc7eac66eb 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -79,15 +79,14 @@ class Image(Document):
     of the list/DocumentArray/Tensor passed as parameters.
 
     This list-like behavior is not always optimal especially when you want
-    to process you data in batches or do operations which involves matrix computation. This
-    is where the stack mode of the DocumentArray comes in handy.
+    to process you data in batches or do operations which involves matrix computation.
+    This is where the stack mode of the DocumentArray comes in handy.
 
-    In **stacked mode** tensor-like fields of every Document are stored on the DocumentArray level, as one stacked tensor. This enables operations on the the entire
+    In **stacked mode** tensor-like fields of every Document are stored on the
+    DocumentArray level, as one stacked tensor. This enables operations on the entire
     batch without iterating over the DocumentArray.
     In this mode the Document inside in the Document don't own the data anymore but just
     reference to the data in the tensor of the DocumentArray.
-    In stacked mode the getter and setter just replace the tensor of the DocumentArray
-    of the given field.
     Operation like `da.append` are not allowed anymore because
     they are too slow and not recommended to use. You should rather use the unstacked
     mode.
@@ -149,10 +148,7 @@ def __getitem__(self, item):
 
     def stack(self):
         """
-        Calling this method will make the DocumentArray enter stacked mode. You should
-        call this method only if the DocumentArray is already in unstacked mode
-        (the default mode). (Calling it while being already in stacked mode will have
-        no effect)
+        Puts the DocumentArray into stacked mode.
 
         When entering stack mode the DocumentArray will create a column for each field
         of the Document that are Tensor or that are Nested Document that contains at
@@ -243,10 +239,8 @@ class Image(Document):
 
     def unstack(self):
         """
-        Calling this method will make the DocumentArray enter in unstack mode.
-        This is the default mode of any DocumentArray. This means you should call this
-        only method if the DocumentArray is already in stacked mode. (Calling it while
-        being in unstacked mode will have no effect)
+        Puts the DocumentArray into unstacked mode.
+
 
         Calling unstack will unstack all the columns of the DocumentArray and put the
         data back in each Document of the DocumentArray.

From 17802b338c1ce8e288f4aa5d9fdd2d941b7f7b6a Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 11:46:09 +0100
Subject: [PATCH 38/53] feat: apply jihannes suggestion

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 6cc7eac66eb..47134d22532 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -156,12 +156,6 @@ def stack(self):
         the whole array at once. In stack mode, accessing or setting fields of the
         DocumentArray will access or set the column of the array.
 
-
-        IMPORTANT: in stacked_mode you cannot add or remove Document from the array.
-        You can only modify the fields of the DocumentArray. This is intentionally done
-        because extending or inserting a column is slow. You should rather use the
-        unstack mode for such operation.
-
         EXAMPLE USAGE
         .. code-block:: python
             from docarray import Document, DocumentArray

From 052288f93b6d7851a0a62d3ad3f96edf23a6daa0 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 11:51:09 +0100
Subject: [PATCH 39/53] feat: apply jihannes suggestion

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 47134d22532..9af19459ed8 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -96,6 +96,17 @@ class Image(Document):
     `with da.stack_mode():` and `with da.unstack_mode():`
 
     see {meth}`.stack` and {meth}`.unstack` for more information.
+
+    You should use unstacked mode, if:
+
+    * You want to insert, append, delete, or shuffle Documents in your DocumentArray
+    * Intend to separate the DocumentArray into smaller batches later on
+
+    You should use stacked mode, if you:
+
+    * You want to process the entire DocumentArray as one batch
+    * Are using the DocumentArray in an ML model for training or inference
+
     """
 
     document_type: Type[BaseDocument] = AnyDocument

From b0dd31d7b06366360099db7efcc439edbb49c42d Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 11:52:56 +0100
Subject: [PATCH 40/53] fix: docstring for blocked method

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 9af19459ed8..c16eb44f66b 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -18,11 +18,11 @@ def wrapper(self, *args, **kwargs):
             )
         return func(self, *args, **kwargs)
 
-        wrapper.__doc__ = (
-            wrapper.__doc__
-            + ' \n This method is not available in stacked mode call {meth}.unstack() '
-            + 'to use it'
-        )
+    wrapper.__doc__ = (
+        wrapper.__doc__
+        + ' \n This method is not available in stacked mode call {meth}.unstack() '
+        + 'to use it'
+    )
 
     return wrapper
 

From 5fd7864a944bf1d87da0ae97127848d9020be7e7 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 11:56:34 +0100
Subject: [PATCH 41/53] feat: add better docstring for context manager

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index c16eb44f66b..4189208a9e2 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -294,6 +294,15 @@ class Image(Document):
 
     @contextmanager
     def stacked_mode(self):
+        """
+        Context manager to put the DocumentArray in stack mode and unstack it when
+        exiting the context manager.
+
+        EXAMPLE USAGE
+        .. code-block:: python
+            with da.stacked_mode():
+                ...
+        """
         try:
             yield self.stack()
         finally:
@@ -301,6 +310,15 @@ def stacked_mode(self):
 
     @contextmanager
     def unstacked_mode(self):
+        """
+        Context manager to put the DocumentArray in unstacked mode and stack it when
+        exiting the context manager.
+
+        EXAMPLE USAGE
+        .. code-block:: python
+            with da.unstacked_mode():
+                ...
+        """
         try:
             yield self.unstack()
         finally:

From c1bfa6f1abdda46c8dc7f4044052bccfc96125e3 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 13:37:59 +0100
Subject: [PATCH 42/53] fix: mypy

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 28 +++++++++++++++++++---------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 4189208a9e2..d1f3da6d132 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -1,7 +1,7 @@
 from collections import defaultdict
 from contextlib import contextmanager
 from functools import wraps
-from typing import Dict, Iterable, List, Optional, Type, TypeVar, Union
+from typing import DefaultDict, Dict, Iterable, List, Optional, Type, TypeVar, Union
 
 from docarray.array.abstract_array import AbstractDocumentArray
 from docarray.array.mixins import GetAttributeArrayMixin, ProtoArrayMixin
@@ -157,9 +157,10 @@ def __getitem__(self, item):
         else:
             return super().__getitem__(item)
 
-    def stack(self):
+    def stack(self: T) -> T:
         """
         Puts the DocumentArray into stacked mode.
+        :return: itself
 
         When entering stack mode the DocumentArray will create a column for each field
         of the Document that are Tensor or that are Nested Document that contains at
@@ -213,9 +214,11 @@ class Image(Document):
                 ):
                     self._columns[field_name] = None
 
-            columns_to_stack: Dict[
+            columns_to_stack: DefaultDict[
                 str, Union[List[TorchTensor], List[NdArray], List[BaseDocument]]
-            ] = defaultdict(list)
+            ] = defaultdict(
+                list
+            )  # type: ignore
 
             for doc in self:
                 for field_to_stack in self._columns.keys():
@@ -228,7 +231,7 @@ class Image(Document):
 
                 type_ = self.document_type.__fields__[field_to_stack].type_
                 if issubclass(type_, BaseDocument):
-                    self._columns[field_to_stack] = DocumentArray[type_](
+                    self._columns[field_to_stack] = DocumentArray[type_](  # type: ignore # noqa: E501
                         to_stack
                     ).stack()
                 else:
@@ -238,13 +241,21 @@ class Image(Document):
                 for field_to_stack in self._columns.keys():
                     type_ = self.document_type.__fields__[field_to_stack].type_
                     if issubclass(type_, TorchTensor) or issubclass(type_, NdArray):
-                        setattr(doc, field_to_stack, self._columns[field_to_stack][i])
+                        if self._columns is not None:
+                            setattr(doc, field_to_stack, self._columns[field_to_stack])
+                        else:
+                            raise ValueError(
+                                f'The DocumentArray is in stacked mode, '
+                                f'but no stacked data is '
+                                f'present for {field_to_stack}. This is inconsistent'
+                            )
 
         return self
 
-    def unstack(self):
+    def unstack(self: T) -> T:
         """
         Puts the DocumentArray into unstacked mode.
+        :return: itself
 
 
         Calling unstack will unstack all the columns of the DocumentArray and put the
@@ -282,8 +293,7 @@ class Image(Document):
 
         see {meth}`.stack` for more information on how to switch to stack mode
         """
-        if self.is_stacked():
-
+        if self.is_stacked() and self._columns:
             for field in list(self._columns.keys()):
                 # list needed here otherwise we are modifying the dict while iterating
                 del self._columns[field]

From 9a75e809645cb71965ade21c48ee045e78494ce8 Mon Sep 17 00:00:00 2001
From: samsja <55492238+samsja@users.noreply.github.com>
Date: Tue, 6 Dec 2022 13:40:52 +0100
Subject: [PATCH 43/53] feat: apply alex suggestion

Co-authored-by: Alex Cureton-Griffiths <alexcg1@users.noreply.github.com>
Signed-off-by: samsja <55492238+samsja@users.noreply.github.com>
Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 57 ++++++++++++++++++++---------------------
 1 file changed, 28 insertions(+), 29 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index d1f3da6d132..8c192dd8353 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -42,11 +42,11 @@ class DocumentArray(
 
     :param docs: iterable of Document
 
-    A DocumentArray is a list of Documents of any schema. Nevertheless, most of the
-    features of DocumentArray are only available if the contains Documents are
-    homogeneous and follow the same schema. To precise this schema you can use the
+    A DocumentArray is a list of Documents of any schema. However, many 
+    DocumentArray features are only available if these Documents are
+    homogeneous and follow the same schema. To be precise, in this schema you can use the
     `DocumentArray[MyDocument]` syntax where MyDocument is a Document class
-    (i,e a schema).This will create a DocumentArray that can only contain Document of
+    (i.e. schema). This creates a DocumentArray that can only contain Documents of
     the type Document.
 
     EXAMPLE USAGE
@@ -60,12 +60,12 @@ class Image(Document):
             url: ImageUrl
 
 
-        da = DocumentArray[Image](Image(url='http://url.com') for _ in range(10))
+        da = DocumentArray[Image](Image(url='http://url.com/foo.png') for _ in range(10))
 
 
-    If your DocumentArray is homogeneous (i.e follow the same schema) You can access the
-    field at the DocumentArray level. Example `da.tensor`. You can set them as well.
-    by doing `da.tensor = np.random.random([10, 100])`
+    If your DocumentArray is homogeneous (i.e. follows the same schema), you can access the
+    field at the DocumentArray level (for example `da.tensor`). You can also set them,
+    with `da.tensor = np.random.random([10, 100])`
 
 
     A DocumentArray can be in one of two modes: unstacked mode and stacked mode.
@@ -73,26 +73,25 @@ class Image(Document):
 
     **Unstacked mode (default)**:
     In this case a DocumentArray is a list of Documents and each Document owns its data.
-    The getter and setter shown above will return a list of the fields of each Document
-    (or a DocumentArray if the field is a nested Document). This list/DocumentArray will
-    be created on the fly. The setter will set the field of each Document to the value
+    The getter and setter shown above returns a list of the fields of each Document
+    (or a DocumentArray if the field is a nested Document). This list/DocumentArray is created on the fly. The setter sets the field of each Document to the value
     of the list/DocumentArray/Tensor passed as parameters.
 
-    This list-like behavior is not always optimal especially when you want
-    to process you data in batches or do operations which involves matrix computation.
+    This list-like behavior is not always optimal, especially when you want
+    to process data in batches or perform operations involving matrix computation.
     This is where the stack mode of the DocumentArray comes in handy.
 
     In **stacked mode** tensor-like fields of every Document are stored on the
     DocumentArray level, as one stacked tensor. This enables operations on the entire
     batch without iterating over the DocumentArray.
-    In this mode the Document inside in the Document don't own the data anymore but just
-    reference to the data in the tensor of the DocumentArray.
-    Operation like `da.append` are not allowed anymore because
-    they are too slow and not recommended to use. You should rather use the unstacked
-    mode.
-
-    To switch from stacked mode to unstacked mode you need to call `da.unstack()` and
-    `da.stack`. There are as well two context manager to for these modes.
+    In this mode the Documents in in the Document don't own the data anymore but just
+    reference the data in the DocumentArray's tensor.
+    Operations like `da.append` are not supported because
+    they are too slow. For these operations you should use unstacked
+    mode instead.
+
+    To switch from stacked to unstacked mode (or vice-versa), call `da.unstack()` or
+    `da.stack`. There are also two context manager for these modes:
     `with da.stack_mode():` and `with da.unstack_mode():`
 
     see {meth}`.stack` and {meth}`.unstack` for more information.
@@ -162,11 +161,11 @@ def stack(self: T) -> T:
         Puts the DocumentArray into stacked mode.
         :return: itself
 
-        When entering stack mode the DocumentArray will create a column for each field
-        of the Document that are Tensor or that are Nested Document that contains at
-        least one Tensor field. This is useful when you want to perform operations on
-        the whole array at once. In stack mode, accessing or setting fields of the
-        DocumentArray will access or set the column of the array.
+        When entering stacked mode the DocumentArray creates a column for each field
+        of the Document that are Tensors or nested Documents that contain at
+        least one Tensor field. This is useful when to perform operations on
+        the whole array at once. In stacked mode, accessing or setting the DocumentArray's fields
+        DocumentArray accesses or sets the column of the array.
 
         EXAMPLE USAGE
         .. code-block:: python
@@ -258,8 +257,8 @@ def unstack(self: T) -> T:
         :return: itself
 
 
-        Calling unstack will unstack all the columns of the DocumentArray and put the
-        data back in each Document of the DocumentArray.
+        Calling unstack will unstack all the columns of the DocumentArray and restore the
+        data of each Document in the DocumentArray.
 
         EXAMPLE USAGE
         .. code-block:: python
@@ -286,7 +285,7 @@ class Image(Document):
             # this work just fine
 
             batch.tensor = np.ones((10, 100))
-            # this iterate over the ndarray and assign the row to each Document
+            # this iterates over the ndarray and assigns each row to a Document
 
             print(batch[0].tensor[0])
             # >>> 1

From d0d75c027848af213cdf885aff99e2273598e619 Mon Sep 17 00:00:00 2001
From: samsja <55492238+samsja@users.noreply.github.com>
Date: Tue, 6 Dec 2022 13:42:43 +0100
Subject: [PATCH 44/53] feat: apply alex suggestion

Co-authored-by: Alex Cureton-Griffiths <alexcg1@users.noreply.github.com>
Signed-off-by: samsja <55492238+samsja@users.noreply.github.com>
Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 8c192dd8353..b35615c43bb 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -290,7 +290,7 @@ class Image(Document):
             print(batch[0].tensor[0])
             # >>> 1
 
-        see {meth}`.stack` for more information on how to switch to stack mode
+        see {meth}`.stack` for more information on switching to stack mode
         """
         if self.is_stacked() and self._columns:
             for field in list(self._columns.keys()):

From e0fec6b84c0f9a18ea1043553a27e0b3c8e2f7f0 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 13:45:01 +0100
Subject: [PATCH 45/53] feat: apply alex suggestion

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 29 ++++++++++++++++-------------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index b35615c43bb..98a1ca4de0c 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -42,12 +42,12 @@ class DocumentArray(
 
     :param docs: iterable of Document
 
-    A DocumentArray is a list of Documents of any schema. However, many 
+    A DocumentArray is a list of Documents of any schema. However, many
     DocumentArray features are only available if these Documents are
-    homogeneous and follow the same schema. To be precise, in this schema you can use the
-    `DocumentArray[MyDocument]` syntax where MyDocument is a Document class
+    homogeneous and follow the same schema. To be precise, in this schema you can use
+    the `DocumentArray[MyDocument]` syntax where MyDocument is a Document class
     (i.e. schema). This creates a DocumentArray that can only contain Documents of
-    the type Document.
+    the type 'MyDocument'.
 
     EXAMPLE USAGE
     .. code-block:: python
@@ -60,12 +60,14 @@ class Image(Document):
             url: ImageUrl
 
 
-        da = DocumentArray[Image](Image(url='http://url.com/foo.png') for _ in range(10))
+        da = DocumentArray[Image](
+            Image(url='http://url.com/foo.png') for _ in range(10)
+        )  # noqa: E510
 
 
-    If your DocumentArray is homogeneous (i.e. follows the same schema), you can access the
-    field at the DocumentArray level (for example `da.tensor`). You can also set them,
-    with `da.tensor = np.random.random([10, 100])`
+    If your DocumentArray is homogeneous (i.e. follows the same schema), you can access
+    the field at the DocumentArray level (for example `da.tensor`). You can also set
+    them, with `da.tensor = np.random.random([10, 100])`
 
 
     A DocumentArray can be in one of two modes: unstacked mode and stacked mode.
@@ -74,7 +76,8 @@ class Image(Document):
     **Unstacked mode (default)**:
     In this case a DocumentArray is a list of Documents and each Document owns its data.
     The getter and setter shown above returns a list of the fields of each Document
-    (or a DocumentArray if the field is a nested Document). This list/DocumentArray is created on the fly. The setter sets the field of each Document to the value
+    (or a DocumentArray if the field is a nested Document). This list/DocumentArray
+    is created on the fly. The setter sets the field of each Document to the value
     of the list/DocumentArray/Tensor passed as parameters.
 
     This list-like behavior is not always optimal, especially when you want
@@ -163,8 +166,8 @@ def stack(self: T) -> T:
 
         When entering stacked mode the DocumentArray creates a column for each field
         of the Document that are Tensors or nested Documents that contain at
-        least one Tensor field. This is useful when to perform operations on
-        the whole array at once. In stacked mode, accessing or setting the DocumentArray's fields
+        least one Tensor field. This is useful when to perform operations on the whole
+        array at once. In stacked mode, accessing or setting the DocumentArray's fields
         DocumentArray accesses or sets the column of the array.
 
         EXAMPLE USAGE
@@ -257,8 +260,8 @@ def unstack(self: T) -> T:
         :return: itself
 
 
-        Calling unstack will unstack all the columns of the DocumentArray and restore the
-        data of each Document in the DocumentArray.
+        Calling unstack will unstack all the columns of the DocumentArray and restore
+        the data of each Document in the DocumentArray.
 
         EXAMPLE USAGE
         .. code-block:: python

From 0b3f53896ad455591a956ba3cc4b62d38cc10bbb Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 13:54:46 +0100
Subject: [PATCH 46/53] fix: mypy

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 98a1ca4de0c..257c43b7b14 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -218,8 +218,8 @@ class Image(Document):
 
             columns_to_stack: DefaultDict[
                 str, Union[List[TorchTensor], List[NdArray], List[BaseDocument]]
-            ] = defaultdict(
-                list
+            ] = defaultdict(  # type: ignore
+                list  # type: ignore
             )  # type: ignore
 
             for doc in self:

From 228277d61d967dd7e10ab48ec81f9976463087e4 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 14:02:07 +0100
Subject: [PATCH 47/53] fix: rebase

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/typing/tensor/abstract_tensor.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py
index 42d13da4497..6a603b7699b 100644
--- a/docarray/typing/tensor/abstract_tensor.py
+++ b/docarray/typing/tensor/abstract_tensor.py
@@ -88,18 +88,16 @@ def validate(
 
         return _ParametrizedTensor
 
-<<<<<<< HEAD
     def __class_getitem__(cls, item: Any):
         target_shape = cls.__validate_getitem__(item)
         return cls._create_parametrized_type(target_shape)
 
-=======
->>>>>>> feat: embedding type (#877)
     @classmethod
     @abc.abstractmethod
     def __docarray_stack__(cls: Type[T], seq: Union[List[T], Tuple[T]]) -> T:
         """Stack a sequence of tensors into a single tensor."""
         ...
+
     def __class_getitem__(cls, item: Any):
         target_shape = cls.__validate_getitem__(item)
         return cls._create_parametrized_type(target_shape)

From 6d07939bae8871035a05c295ce000c46da5c92a7 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 14:02:29 +0100
Subject: [PATCH 48/53] fix: rebase

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/typing/tensor/abstract_tensor.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py
index 6a603b7699b..f96c075a39c 100644
--- a/docarray/typing/tensor/abstract_tensor.py
+++ b/docarray/typing/tensor/abstract_tensor.py
@@ -97,7 +97,3 @@ def __class_getitem__(cls, item: Any):
     def __docarray_stack__(cls: Type[T], seq: Union[List[T], Tuple[T]]) -> T:
         """Stack a sequence of tensors into a single tensor."""
         ...
-
-    def __class_getitem__(cls, item: Any):
-        target_shape = cls.__validate_getitem__(item)
-        return cls._create_parametrized_type(target_shape)

From 65ba015be2e6f7762c73541715440c4011831950 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 15:37:19 +0100
Subject: [PATCH 49/53] refactor: fix problem

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/abstract_array.py    |  4 +--
 docarray/array/array.py             | 51 ++++++++++++++++++-----------
 tests/units/array/test_stack_mod.py |  3 ++
 3 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/docarray/array/abstract_array.py b/docarray/array/abstract_array.py
index a93a4e6a288..da7cf397fc4 100644
--- a/docarray/array/abstract_array.py
+++ b/docarray/array/abstract_array.py
@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Sequence, Type, Union
 
 from docarray.document import BaseDocument
 
@@ -7,7 +7,7 @@
     from docarray.typing import NdArray, TorchTensor
 
 
-class AbstractDocumentArray(Iterable):
+class AbstractDocumentArray(Sequence):
 
     document_type: Type[BaseDocument]
     _columns: Optional[
diff --git a/docarray/array/array.py b/docarray/array/array.py
index 257c43b7b14..9ccb6dd12fd 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -149,7 +149,10 @@ def _setter(self, value):
 
         return _DocumentArrayTyped
 
-    def __getitem__(self, item):
+    def __list_getitem__(self, item):
+        return super().__getitem__(item)
+
+    def __getitem__(self, item):  # note this should handle slices
         if self.is_stacked():
             doc = super().__getitem__(item)
             # NOTE: this could be speed up by using a cache
@@ -157,7 +160,14 @@ def __getitem__(self, item):
                 setattr(doc, field, self._columns[field][item])
             return doc
         else:
-            return super().__getitem__(item)
+            return self.__list_getitem__(item)
+
+    def __list_iter__(self):
+        return super().__iter__()
+
+    def __iter__(self):
+        for i in range(len(self)):
+            yield self[i]
 
     def stack(self: T) -> T:
         """
@@ -204,7 +214,7 @@ class Image(Document):
         see {meth}`.unstack` for more information on how to switch to unstack mode
         """
 
-        if not (self.is_stacked()):
+        if not self.is_stacked():
 
             self._columns = dict()
 
@@ -222,8 +232,8 @@ class Image(Document):
                 list  # type: ignore
             )  # type: ignore
 
-            for doc in self:
-                for field_to_stack in self._columns.keys():
+            for doc in self.__list_iter__():
+                for field_to_stack in self._column_fields():
                     columns_to_stack[field_to_stack].append(
                         getattr(doc, field_to_stack)
                     )
@@ -239,19 +249,6 @@ class Image(Document):
                 else:
                     self._columns[field_to_stack] = type_.__docarray_stack__(to_stack)
 
-            for i, doc in enumerate(self):
-                for field_to_stack in self._columns.keys():
-                    type_ = self.document_type.__fields__[field_to_stack].type_
-                    if issubclass(type_, TorchTensor) or issubclass(type_, NdArray):
-                        if self._columns is not None:
-                            setattr(doc, field_to_stack, self._columns[field_to_stack])
-                        else:
-                            raise ValueError(
-                                f'The DocumentArray is in stacked mode, '
-                                f'but no stacked data is '
-                                f'present for {field_to_stack}. This is inconsistent'
-                            )
-
         return self
 
     def unstack(self: T) -> T:
@@ -296,7 +293,19 @@ class Image(Document):
         see {meth}`.stack` for more information on switching to stack mode
         """
         if self.is_stacked() and self._columns:
-            for field in list(self._columns.keys()):
+
+            for i, doc in enumerate(self.__list_iter__()):
+                for field in self._column_fields():
+                    val = self._columns[field]
+                    if val is not None:
+                        setattr(doc, field, val[i])
+                    else:
+                        raise RuntimeError('Internal error, one of the column is None')
+                    # NOTE: here we might need to copy the tensor
+                    # see here
+                    # https://discuss.pytorch.org/t/what-happened-to-a-view-of-a-tensor-when-the-original-tensor-is-deleted/167294 # noqa: E501
+
+            for field in self._column_fields():
                 # list needed here otherwise we are modifying the dict while iterating
                 del self._columns[field]
 
@@ -304,6 +313,10 @@ class Image(Document):
 
         return self
 
+    def a(self, x: Union[TorchTensor, NdArray, AbstractDocumentArray, None]):
+        if x is not None:
+            x[0]
+
     @contextmanager
     def stacked_mode(self):
         """
diff --git a/tests/units/array/test_stack_mod.py b/tests/units/array/test_stack_mod.py
index 35ea9ffc786..3fe20c07ddf 100644
--- a/tests/units/array/test_stack_mod.py
+++ b/tests/units/array/test_stack_mod.py
@@ -25,6 +25,9 @@ class Image(Document):
     for doc, tensor in zip(batch, batch.tensor):
         assert doc.tensor.data_ptr() == tensor.data_ptr()
 
+    for i in range(len(batch)):
+        assert batch[i].tensor.data_ptr() == batch.tensor[i].data_ptr()
+
 
 def test_stack_mod_nested_document():
     class Image(Document):

From baa7f537fe0ddaa7bc03a4e5530486444b155f9c Mon Sep 17 00:00:00 2001
From: samsja <55492238+samsja@users.noreply.github.com>
Date: Tue, 6 Dec 2022 15:40:41 +0100
Subject: [PATCH 50/53] feat: apply alex suggestion

Co-authored-by: Alex Cureton-Griffiths <alexcg1@users.noreply.github.com>
Signed-off-by: samsja <55492238+samsja@users.noreply.github.com>
---
 docarray/array/array.py            | 28 ++++++++++++++--------------
 docarray/array/mixins/attribute.py |  4 ++--
 docarray/typing/tensor/ndarray.py  |  2 +-
 3 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 9ccb6dd12fd..5f4ca0cab58 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -66,8 +66,8 @@ class Image(Document):
 
 
     If your DocumentArray is homogeneous (i.e. follows the same schema), you can access
-    the field at the DocumentArray level (for example `da.tensor`). You can also set
-    them, with `da.tensor = np.random.random([10, 100])`
+    fields at the DocumentArray level (for example `da.tensor`). You can also set
+    fields, with `da.tensor = np.random.random([10, 100])`
 
 
     A DocumentArray can be in one of two modes: unstacked mode and stacked mode.
@@ -75,8 +75,8 @@ class Image(Document):
 
     **Unstacked mode (default)**:
     In this case a DocumentArray is a list of Documents and each Document owns its data.
-    The getter and setter shown above returns a list of the fields of each Document
-    (or a DocumentArray if the field is a nested Document). This list/DocumentArray
+    The getter and setter shown above return a list of the fields of each Document
+    (or DocumentArray if the field is a nested Document). This list/DocumentArray
     is created on the fly. The setter sets the field of each Document to the value
     of the list/DocumentArray/Tensor passed as parameters.
 
@@ -84,30 +84,30 @@ class Image(Document):
     to process data in batches or perform operations involving matrix computation.
     This is where the stack mode of the DocumentArray comes in handy.
 
-    In **stacked mode** tensor-like fields of every Document are stored on the
-    DocumentArray level, as one stacked tensor. This enables operations on the entire
+    In **stacked mode**, tensor-like fields of every Document are stored on the
+    DocumentArray level as one stacked tensor. This enables operations on the entire
     batch without iterating over the DocumentArray.
-    In this mode the Documents in in the Document don't own the data anymore but just
+    In this mode the Documents in the DocumentArray no longer own the data but just
     reference the data in the DocumentArray's tensor.
     Operations like `da.append` are not supported because
     they are too slow. For these operations you should use unstacked
     mode instead.
 
     To switch from stacked to unstacked mode (or vice-versa), call `da.unstack()` or
-    `da.stack`. There are also two context manager for these modes:
+    `da.stack`. There are also two context managers for these modes:
     `with da.stack_mode():` and `with da.unstack_mode():`
 
     see {meth}`.stack` and {meth}`.unstack` for more information.
 
-    You should use unstacked mode, if:
+    You should use unstacked mode if you want to:
 
-    * You want to insert, append, delete, or shuffle Documents in your DocumentArray
-    * Intend to separate the DocumentArray into smaller batches later on
+    * Insert, append, delete, or shuffle Documents in your DocumentArray.
+    * Separate the DocumentArray into smaller batches later on.
 
-    You should use stacked mode, if you:
+    You should use stacked mode if you want to:
 
-    * You want to process the entire DocumentArray as one batch
-    * Are using the DocumentArray in an ML model for training or inference
+    * Process the entire DocumentArray as one batch.
+    * Use the DocumentArray in an ML model for training or inference.
 
     """
 
diff --git a/docarray/array/mixins/attribute.py b/docarray/array/mixins/attribute.py
index e206767f720..148111db7f5 100644
--- a/docarray/array/mixins/attribute.py
+++ b/docarray/array/mixins/attribute.py
@@ -51,10 +51,10 @@ def _set_array_attribute(
         field: str,
         values: Union[List, AbstractDocumentArray, 'TorchTensor', 'NdArray'],
     ):
-        """Set all document if this DocumentArray with the passed values
+        """Set all Documents in this DocumentArray using the passed values
 
         :param field: name of the fields to extract
-        :values: the values to set at the document array level
+        :values: the values to set at the DocumentArray level
         """
         if (
             self.is_stacked()
diff --git a/docarray/typing/tensor/ndarray.py b/docarray/typing/tensor/ndarray.py
index ad002187328..af5092f086c 100644
--- a/docarray/typing/tensor/ndarray.py
+++ b/docarray/typing/tensor/ndarray.py
@@ -196,5 +196,5 @@ def _flush_tensor_to_proto(pb_msg: 'NdArrayProto', value: 'NdArray'):
     def __docarray_stack__(
         cls: Type[T], seq: Union[List['NdArray'], Tuple['NdArray']]
     ) -> T:
-        """Stack a sequence of ndarray into a single ndarray."""
+        """Stack a sequence of ndarrays into a single ndarray."""
         return cls.from_ndarray(np.stack(seq))

From d7c0458c31c134b5ca5a76473df62d80b8fca47f Mon Sep 17 00:00:00 2001
From: samsja <55492238+samsja@users.noreply.github.com>
Date: Tue, 6 Dec 2022 15:42:57 +0100
Subject: [PATCH 51/53] feat: apply alex suggestion

Co-authored-by: Alex Cureton-Griffiths <alexcg1@users.noreply.github.com>
Signed-off-by: samsja <55492238+samsja@users.noreply.github.com>
---
 docarray/array/array.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 5f4ca0cab58..95b85649f49 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -174,11 +174,14 @@ def stack(self: T) -> T:
         Puts the DocumentArray into stacked mode.
         :return: itself
 
-        When entering stacked mode the DocumentArray creates a column for each field
-        of the Document that are Tensors or nested Documents that contain at
-        least one Tensor field. This is useful when to perform operations on the whole
+        When entering stacked mode DocumentArray creates a column for:
+        
+         - Each Tensor field of each Document
+         - Nested Documents that contain at least one Tensor field. 
+         
+        This is useful to perform operations on the whole
         array at once. In stacked mode, accessing or setting the DocumentArray's fields
-        DocumentArray accesses or sets the column of the array.
+        accesses or sets the columns of the array.
 
         EXAMPLE USAGE
         .. code-block:: python
@@ -257,7 +260,7 @@ def unstack(self: T) -> T:
         :return: itself
 
 
-        Calling unstack will unstack all the columns of the DocumentArray and restore
+        Calling unstack will unstack all columns of the DocumentArray and restore
         the data of each Document in the DocumentArray.
 
         EXAMPLE USAGE
@@ -320,7 +323,7 @@ def a(self, x: Union[TorchTensor, NdArray, AbstractDocumentArray, None]):
     @contextmanager
     def stacked_mode(self):
         """
-        Context manager to put the DocumentArray in stack mode and unstack it when
+        Context manager to put the DocumentArray in stacked mode and unstack it when
         exiting the context manager.
 
         EXAMPLE USAGE

From 0049d66e32e80d8852b821e27469e3e326231528 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 15:47:15 +0100
Subject: [PATCH 52/53] fix: apply black to suggestion

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 95b85649f49..3ac4bc4efbf 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -175,10 +175,10 @@ def stack(self: T) -> T:
         :return: itself
 
         When entering stacked mode DocumentArray creates a column for:
-        
+
          - Each Tensor field of each Document
-         - Nested Documents that contain at least one Tensor field. 
-         
+         - Nested Documents that contain at least one Tensor field.
+
         This is useful to perform operations on the whole
         array at once. In stacked mode, accessing or setting the DocumentArray's fields
         accesses or sets the columns of the array.

From 12868768b88d31376b437989de82925826914455 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 6 Dec 2022 16:10:36 +0100
Subject: [PATCH 53/53] fix: update docstring

Signed-off-by: Sami Jaghouar <sami.jaghouar@hotmail.fr>
---
 docarray/array/array.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/docarray/array/array.py b/docarray/array/array.py
index 3ac4bc4efbf..c71c8f54830 100644
--- a/docarray/array/array.py
+++ b/docarray/array/array.py
@@ -19,9 +19,7 @@ def wrapper(self, *args, **kwargs):
         return func(self, *args, **kwargs)
 
     wrapper.__doc__ = (
-        wrapper.__doc__
-        + ' \n This method is not available in stacked mode call {meth}.unstack() '
-        + 'to use it'
+        wrapper.__doc__ + ' \n This method is not available in stacked mode call'
     )
 
     return wrapper