docarray · anna-charlotte · Apr 11, 2023 · Apr 6, 2023 · Apr 6, 2023 · Apr 6, 2023
diff --git a/docarray/typing/__init__.py b/docarray/typing/__init__.py
@@ -71,6 +71,7 @@
  'AudioTensorFlowTensor',
  'VideoTensorFlowTensor',
 ]
+__all_test__ = __all__ + _torch_tensors
 
 
 def __getattr__(name: str):
@@ -86,7 +87,6 @@ def __getattr__(name: str):
  import docarray.typing.tensor
 
  tensor_cls = getattr(docarray.typing.tensor, name)
-
  if name not in __all__:
  __all__.append(name)
 

diff --git a/docarray/typing/bytes/audio_bytes.py b/docarray/typing/bytes/audio_bytes.py
@@ -46,7 +46,8 @@ def _to_node_protobuf(self: T) -> 'NodeProto':
 
     def load(self) -> Tuple[AudioNdArray, int]:
         """
-        Load the Audio from the AudioBytes into an AudioNdArray
+        Load the Audio from the [`AudioBytes`][docarray.typing.AudioBytes] into an
+        [`AudioNdArray`][docarray.typing.AudioNdArray].
 
         ---
 
@@ -75,8 +76,8 @@ class MyAudio(BaseDoc):
         ```
 
         ---
-        :return: tuple of an AudioNdArray representing the audio bytes content,
-            and an integer representing the frame rate.
+        :return: tuple of an [`AudioNdArray`][docarray.typing.AudioNdArray] representing the
+ audio bytes content, and an integer representing the frame rate.
         """
         pydub = import_library('pydub', raise_error=True)  # noqa: F841
         from pydub import AudioSegment

diff --git a/docarray/typing/bytes/image_bytes.py b/docarray/typing/bytes/image_bytes.py
@@ -84,7 +84,8 @@ def load(
         axis_layout: Tuple[str, str, str] = ('H', 'W', 'C'),
     ) -> ImageNdArray:
         """
-        Load the image from the ImageBytes into an ImageNdArray
+        Load the image from the [`ImageBytes`][docarray.typing.ImageBytes] into an
+        [`ImageNdArray`][docarray.typing.ImageNdArray].
 
         ---
 
@@ -119,7 +120,7 @@ class MyDoc(BaseDoc):
         :param height: height of the image tensor.
         :param axis_layout: ordering of the different image axes.
             'H' = height, 'W' = width, 'C' = color channel
-        :return: ImageNdArray representing the image as RGB values
+        :return: [`ImageNdArray`][docarray.typing.ImageNdArray] representing the image as RGB values
         """
         raw_img = self.load_pil()
 

diff --git a/docarray/typing/bytes/video_bytes.py b/docarray/typing/bytes/video_bytes.py
@@ -51,10 +51,11 @@ def _to_node_protobuf(self: T) -> 'NodeProto':
 
     def load(self, **kwargs) -> VideoLoadResult:
         """
-        Load the video from the bytes into a VideoLoadResult object consisting of a
-        VideoNdArray (`VideoLoadResult.video`), an AudioNdArray
-        (`VideoLoadResult.audio`) and an NdArray containing the key frame indices
-        (`VideoLoadResult.key_frame_indices`).
+        Load the video from the bytes into a VideoLoadResult object consisting of:
+
+        - a [`VideoNdArray`][docarray.typing.VideoNdArray] (`VideoLoadResult.video`)
+        - an [`AudioNdArray`][docarray.typing.AudioNdArray] (`VideoLoadResult.audio`)
+        - an [`NdArray`][docarray.typing.NdArray] containing the key frame indices (`VideoLoadResult.key_frame_indices`).
 
         ---
 
@@ -82,7 +83,7 @@ class MyDoc(BaseDoc):
 
         :param kwargs: supports all keyword arguments that are being supported by
             av.open() as described [here](https://pyav.org/docs/stable/api/_globals.html?highlight=open#av.open)
-        :return: a VideoLoadResult instance with video, audio and keyframe indices
+        :return: a `VideoLoadResult` instance with video, audio and keyframe indices
         """
         if TYPE_CHECKING:
             import av

diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py
@@ -114,7 +114,6 @@ def _to_node_protobuf(self: T) -> 'NodeProto':
  """Convert itself into a NodeProto protobuf message. This function should
         be called when the Document is nested into another Document that need to be
         converted into a protobuf
-        :param field: field in which to store the content in the node proto
         :return: the nested item protobuf message
         """
  from docarray.proto import NodeProto
@@ -128,12 +127,15 @@ def __docarray_validate_shape__(cls, t: T, shape: Tuple[Union[int, str], ...]) -
         enable syntax of the form AnyTensor[shape].
         It is called when a tensor is assigned to a field of this type.
         i.e. when a tensor is passed to a Document field of type AnyTensor[shape].
+
         The intended behaviour is as follows:
+
         - If the shape of `t` is equal to `shape`, return `t`.
         - If the shape of `t` is not equal to `shape`,
             but can be reshaped to `shape`, return `t` reshaped to `shape`.
         - If the shape of `t` is not equal to `shape`
             and cannot be reshaped to `shape`, raise a ValueError.
+
         :param t: The tensor to validate.
         :param shape: The shape to validate against.
         :return: The validated tensor.
@@ -197,7 +199,7 @@ def __docarray_validate_shape__(cls, t: T, shape: Tuple[Union[int, str], ...]) -
 
  @classmethod
  def __docarray_validate_getitem__(cls, item: Any) -> Tuple[int]:
- """This method validates the input to __class_getitem__.
+ """This method validates the input to `AbstractTensor.__class_getitem__`.
 
         It is called at "class creation time",
         i.e. when a class is created with syntax of the form AnyTensor[shape].
@@ -206,11 +208,13 @@ def __docarray_validate_getitem__(cls, item: Any) -> Tuple[int]:
         A subclass can override this method to implement custom validation logic.
 
         The output of this is eventually passed to
-        {ref}`AbstractTensor.__validate_shape__` as its `shape` argument.
+        [`AbstractTensor.__docarray_validate_shape__`]
+        [docarray.typing.tensor.abstract_tensor.AbstractTensor.__docarray_validate_shape__]
+        as its `shape` argument.
 
         Raises `ValueError` if the input `item` does not pass validation.
 
-        :param item: The item to validate, passed to __class_getitem__ (`Tensor[item]`).
+        :param item: The item to validate, passed to `__class_getitem__` (`Tensor[item]`).
         :return: The validated item == the target shape of this tensor.
         """
  if isinstance(item, int):

diff --git a/docarray/typing/tensor/audio/audio_ndarray.py b/docarray/typing/tensor/audio/audio_ndarray.py
@@ -6,7 +6,7 @@
 @_register_proto(proto_type_name='audio_ndarray')
 class AudioNdArray(AbstractAudioTensor, NdArray):
     """
-    Subclass of NdArray, to represent an audio tensor.
+    Subclass of [`NdArray`][docarray.typing.NdArray], to represent an audio tensor.
     Adds audio-specific features to the tensor.
 
 
@@ -33,7 +33,7 @@ class MyAudioDoc(BaseDoc):
         audio_tensor=np.random.rand(1000, 2),
     )
 
-    doc_1.audio_tensor.save(file_path='/tmp/file_1.wav')
+ # doc_1.audio_tensor.save(file_path='/tmp/file_1.wav')
     doc_1.bytes_ = doc_1.audio_tensor.to_bytes()
 
     # from url
@@ -43,7 +43,7 @@ class MyAudioDoc(BaseDoc):
     )
 
     doc_2.audio_tensor, _ = doc_2.url.load()
-    doc_2.audio_tensor.save(file_path='/tmp/file_2.wav')
+ # doc_2.audio_tensor.save(file_path='/tmp/file_2.wav')
     doc_2.bytes_ = doc_1.audio_tensor.to_bytes()
     ```
 

diff --git a/docarray/typing/tensor/audio/audio_tensorflow_tensor.py b/docarray/typing/tensor/audio/audio_tensorflow_tensor.py
@@ -12,18 +12,15 @@ class AudioTensorFlowTensor(
     AbstractAudioTensor, TensorFlowTensor, metaclass=metaTensorFlow
 ):
     """
-    Subclass of TensorFlowTensor, to represent an audio tensor.
-    Adds audio-specific features to the tensor.
-
-
-
-    ```
+    Subclass of [`TensorFlowTensor`][docarray.typing.TensorFlowTensor],
+    to represent an audio tensor. Adds audio-specific features to the tensor.
 
     ---
+
+    ```python
     from typing import Optional
 
     import tensorflow as tf
-    from pydantic import parse_obj_as
 
     from docarray import BaseDoc
     from docarray.typing import AudioBytes, AudioTensorFlowTensor, AudioUrl
@@ -41,19 +38,16 @@ class MyAudioDoc(BaseDoc):
         audio_tensor=tf.random.normal((1000, 2)),
     )
 
-    doc_1.audio_tensor.save(file_path='path/to/file_1.wav')
+ # doc_1.audio_tensor.save(file_path='file_1.wav')
     doc_1.bytes_ = doc_1.audio_tensor.to_bytes()
 
-
     doc_2 = MyAudioDoc(
         title='my_second_audio_doc',
         url='https://www.kozco.com/tech/piano2.wav',
     )
 
-    doc_2.audio_tensor = doc_2.url.load()
-    doc_2.audio_tensor.save(file_path='path/to/file_2.wav')
+    doc_2.audio_tensor, _ = doc_2.url.load()
     doc_2.bytes_ = doc_1.audio_tensor.to_bytes()
-
     ```
 
     ---

diff --git a/docarray/typing/tensor/audio/audio_torch_tensor.py b/docarray/typing/tensor/audio/audio_torch_tensor.py
@@ -6,10 +6,9 @@
 @_register_proto(proto_type_name='audio_torch_tensor')
 class AudioTorchTensor(AbstractAudioTensor, TorchTensor, metaclass=metaTorchAndNode):
     """
-    Subclass of TorchTensor, to represent an audio tensor.
+    Subclass of [`TorchTensor`][docarray.typing.TorchTensor], to represent an audio tensor.
     Adds audio-specific features to the tensor.
 
-
     ---
 
     ```python
@@ -33,7 +32,7 @@ class MyAudioDoc(BaseDoc):
         audio_tensor=torch.zeros(1000, 2),
     )
 
-    doc_1.audio_tensor.save(file_path='/tmp/file_1.wav')
+ # doc_1.audio_tensor.save(file_path='/tmp/file_1.wav')
     doc_1.bytes_ = doc_1.audio_tensor.to_bytes()
 
     doc_2 = MyAudioDoc(
@@ -42,7 +41,7 @@ class MyAudioDoc(BaseDoc):
     )
 
     doc_2.audio_tensor, _ = doc_2.url.load()
-    doc_2.audio_tensor.save(file_path='/tmp/file_2.wav')
+ # doc_2.audio_tensor.save(file_path='/tmp/file_2.wav')
     doc_2.bytes_ = doc_1.audio_tensor.to_bytes()
     ```
 

diff --git a/docarray/typing/tensor/embedding/torch.py b/docarray/typing/tensor/embedding/torch.py
@@ -18,7 +18,7 @@ class TorchEmbedding(TorchTensor, EmbeddingMixin, metaclass=metaTorchAndEmbeddin
 
  def new_empty(self, *args, **kwargs):
  """
-        This method enables the deepcopy of TorchEmbedding by returning another instance of this subclass.
+        This method enables the deepcopy of `TorchEmbedding` by returning another instance of this subclass.
         If this function is not implemented, the deepcopy will throw an RuntimeError from Torch.
         """
  return self.__class__(TorchTensor.new_empty(self, *args, **kwargs))
diff --git a/docarray/typing/tensor/image/image_ndarray.py b/docarray/typing/tensor/image/image_ndarray.py
@@ -8,10 +8,10 @@
 @_register_proto(proto_type_name='image_ndarray')
 class ImageNdArray(AbstractImageTensor, NdArray):
  """
-    Subclass of NdArray, to represent an image tensor.
+    Subclass of [`NdArray`][docarray.typing.NdArray], to represent an image tensor.
     Adds image-specific features to the tensor.
     For instance the ability convert the tensor back to image bytes which are
-    optimized to send over the wire
+    optimized to send over the wire.
 
 
     ---

diff --git a/docarray/typing/tensor/image/image_tensorflow_tensor.py b/docarray/typing/tensor/image/image_tensorflow_tensor.py
@@ -12,10 +12,10 @@ class ImageTensorFlowTensor(
  TensorFlowTensor, AbstractImageTensor, metaclass=metaTensorFlow
 ):
  """
-    Subclass of TensorFlowTensor, to represent an image tensor.
-    Adds image-specific features to the tensor.
+    Subclass of [`TensorFlowTensor`][docarray.typing.TensorFlowTensor],
+ to represent an image tensor. Adds image-specific features to the tensor.
     For instance the ability convert the tensor back to image bytes which are
-    optimized to send over the wire
+    optimized to send over the wire.
 
 
     ---

diff --git a/docarray/typing/tensor/image/image_torch_tensor.py b/docarray/typing/tensor/image/image_torch_tensor.py
@@ -10,10 +10,10 @@
 @_register_proto(proto_type_name='image_torch_tensor')
 class ImageTorchTensor(AbstractImageTensor, TorchTensor, metaclass=metaTorchAndNode):
  """
-    Subclass of TorchTensor, to represent an image tensor.
+    Subclass of [`TorchTensor`][docarray.typing.TorchTensor], to represent an image tensor.
     Adds image-specific features to the tensor.
     For instance the ability convert the tensor back to image bytes which are
-    optimized to send over the wire
+    optimized to send over the wire.
 
 
     ---

diff --git a/docarray/typing/tensor/ndarray.py b/docarray/typing/tensor/ndarray.py
@@ -40,9 +40,9 @@ class metaNumpy(AbstractTensor.__parametrized_meta__, tensor_base):  # type: ign
 @_register_proto(proto_type_name='ndarray')
 class NdArray(np.ndarray, AbstractTensor, Generic[ShapeT]):
     """
-    Subclass of np.ndarray, intended for use in a Document.
+    Subclass of `np.ndarray`, intended for use in a Document.
     This enables (de)serialization from/to protobuf and json, data validation,
-    and coersion from compatible types like torch.Tensor.
+    and coersion from compatible types like `torch.Tensor`.
 
     This type can also be used in a parametrized way, specifying the shape of the array.
 
@@ -143,7 +143,7 @@ def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None:
 
     def _docarray_to_json_compatible(self) -> np.ndarray:
         """
-        Convert tensor into a json compatible object
+        Convert `NdArray` into a json compatible object
         :return: a representation of the tensor compatible with orjson
         """
         return self.unwrap()
@@ -152,30 +152,33 @@ def unwrap(self) -> np.ndarray:
         """
         Return the original ndarray without any memory copy.
 
-        The original view rest intact and is still a Document NdArray
-        but the return object is a pure np.ndarray but both object share
+        The original view rest intact and is still a Document `NdArray`
+        but the return object is a pure `np.ndarray` but both object share
         the same memory layout.
 
-        EXAMPLE USAGE
-        .. code-block:: python
-            from docarray.typing import NdArray
-            import numpy as np
+        ---
 
-            t1 = NdArray.validate(np.zeros((3, 224, 224)), None, None)
-            # here t1 is a docarray NdArray
-            t2 = t.unwrap()
-            # here t2 is a pure np.ndarray but t1 is still a Docarray NdArray
-            # But both share the same underlying memory
+        ```python
+        from docarray.typing import NdArray
+        import numpy as np
 
+        t1 = NdArray.validate(np.zeros((3, 224, 224)), None, None)
+        # here t1 is a docarray NdArray
+        t2 = t1.unwrap()
+        # here t2 is a pure np.ndarray but t1 is still a Docarray NdArray
+        # But both share the same underlying memory
+        ```
 
-        :return: a numpy ndarray
+        ---
+
+        :return: a `numpy.ndarray`
         """
         return self.view(np.ndarray)
 
     @classmethod
     def from_protobuf(cls: Type[T], pb_msg: 'NdArrayProto') -> 'T':
         """
- read ndarray from a proto msg
+ Read ndarray from a proto msg
         :param pb_msg:
         :return: a numpy array
         """
@@ -190,7 +193,7 @@ def from_protobuf(cls: Type[T], pb_msg: 'NdArrayProto') -> 'T':
 
     def to_protobuf(self) -> 'NdArrayProto':
         """
- transform self into a NdArrayProto protobuf message
+ Transform self into a NdArrayProto protobuf message
         """
         from docarray.proto import NdArrayProto