Skip to content

refactor: comp backend #1030

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
samsja merged 7 commits into from
Jan 18, 2023
Merged

refactor: comp backend #1030

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
1cd2ccc
refactor: get_comp_backend return a obj not a class
samsja Jan 17, 2023
d714b1f
wip
samsja Jan 17, 2023
4f77ab5
wip
samsja Jan 17, 2023
a5ada44
refactor: no need for decorator
samsja Jan 17, 2023
5b1dfd0
fix: remove comment
samsja Jan 18, 2023
7b9c145
Merge branch 'feat-rewrite-v2' into refactor-comp-backend
samsja Jan 18, 2023
89fee7c
Merge branch 'feat-rewrite-v2' into refactor-comp-backend
samsja Jan 18, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 3 additions & 8 deletions docarray/array/array_stacked.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,14 +95,9 @@ def to(self: T, device: str) -> T:
for field in self._columns.keys():
col = self._columns[field]
if isinstance(col, AbstractTensor):
# the casting below is arbitrary, in reality `col` could be of any
# subclass of AbstractTensor. But to make mypy happy we have to cast
# it to a concrete subclass thereof
# see mypy issue: https://github.com/python/mypy/issues/14421
col_ = cast('TorchTensor', col)
self._columns[field] = col_.get_comp_backend().to_device(col_, device)
elif isinstance(col, NdArray):
self._columns[field] = col.get_comp_backend().to_device(col, device)
self._columns[field] = col.__class__._docarray_from_native(
col.get_comp_backend().to_device(col, device)
)
else: # recursive call
col_docarray = cast(T, col)
col_docarray.to(device)
Expand Down
35 changes: 3 additions & 32 deletions docarray/computation/abstract_comp_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,18 +1,17 @@
import typing
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, TypeVar, Union, overload
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, TypeVar, Union

if TYPE_CHECKING:
import numpy as np

# In practice all of the below will be the same type
TTensor = TypeVar('TTensor')
TAbstractTensor = TypeVar('TAbstractTensor')
TTensorRetrieval = TypeVar('TTensorRetrieval')
TTensorMetrics = TypeVar('TTensorMetrics')


class AbstractComputationalBackend(ABC, typing.Generic[TTensor, TAbstractTensor]):
class AbstractComputationalBackend(ABC, typing.Generic[TTensor]):
"""
Abstract base class for computational backends.
Every supported tensor/ML framework (numpy, torch etc.) should define its own
Expand Down Expand Up @@ -73,37 +72,9 @@ def shape(tensor: 'TTensor') -> Tuple[int, ...]:
"""Get shape of tensor"""
...

@overload
@staticmethod
def reshape(tensor: 'TAbstractTensor', shape: Tuple[int, ...]) -> 'TAbstractTensor':
"""
Gives a new shape to tensor without changing its data.

:param tensor: tensor to be reshaped
:param shape: the new shape
:return: a tensor with the same data and number of elements as tensor
but with the specified shape.
"""
...

@overload
@staticmethod
def reshape(tensor: 'TTensor', shape: Tuple[int, ...]) -> 'TTensor':
"""
Gives a new shape to tensor without changing its data.

:param tensor: tensor to be reshaped
:param shape: the new shape
:return: a tensor with the same data and number of elements as tensor
but with the specified shape.
"""
...

@staticmethod
@abstractmethod
def reshape(
tensor: Union['TTensor', 'TAbstractTensor'], shape: Tuple[int, ...]
) -> Union['TTensor', 'TAbstractTensor']:
def reshape(tensor: 'TTensor', shape: Tuple[int, ...]) -> 'TTensor':
"""
Gives a new shape to tensor without changing its data.

Expand Down
47 changes: 2 additions & 45 deletions docarray/computation/numpy_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
import warnings
from typing import Any, List, Optional, Tuple, Union, overload
from typing import Any, List, Optional, Tuple, Union

import numpy as np

from docarray.computation import AbstractComputationalBackend
from docarray.typing import NdArray


def _expand_if_single_axis(*matrices: np.ndarray) -> List[np.ndarray]:
Expand All @@ -30,7 +29,7 @@ def _expand_if_scalar(arr: np.ndarray) -> np.ndarray:
return arr


class NumpyCompBackend(AbstractComputationalBackend[np.ndarray, NdArray]):
class NumpyCompBackend(AbstractComputationalBackend[np.ndarray]):
"""
Computational backend for Numpy.
"""
Expand All @@ -41,22 +40,8 @@ def stack(
) -> 'np.ndarray':
return np.stack(tensors, axis=dim)

@overload
@staticmethod
def to_device(tensor: 'NdArray', device: str) -> 'NdArray':
"""Move the tensor to the specified device."""
...

@overload
@staticmethod
def to_device(tensor: 'np.ndarray', device: str) -> 'np.ndarray':
"""Move the tensor to the specified device."""
...

@staticmethod
def to_device(
tensor: Union['np.ndarray', 'NdArray'], device: str
) -> Union['np.ndarray', 'NdArray']:
"""Move the tensor to the specified device."""
raise NotImplementedError('Numpy does not support devices (GPU).')

Expand Down Expand Up @@ -88,39 +73,11 @@ def shape(array: 'np.ndarray') -> Tuple[int, ...]:
"""Get shape of array"""
return array.shape

@overload
@staticmethod
def reshape(array: 'NdArray', shape: Tuple[int, ...]) -> 'NdArray':
"""
Gives a new shape to array without changing its data.

:param array: array to be reshaped
:param shape: the new shape
:return: a array with the same data and number of elements as array
but with the specified shape.
"""
...

@overload
@staticmethod
def reshape(array: 'np.ndarray', shape: Tuple[int, ...]) -> 'np.ndarray':
"""
Gives a new shape to array without changing its data.

:param array: array to be reshaped
:param shape: the new shape
:return: a array with the same data and number of elements as array
but with the specified shape.
"""
...

@staticmethod
def reshape(
array: Union['np.ndarray', 'NdArray'], shape: Tuple[int, ...]
) -> Union['np.ndarray', 'NdArray']:
"""
Gives a new shape to array without changing its data.

:param array: array to be reshaped
:param shape: the new shape
:return: a array with the same data and number of elements as array
Expand Down
47 changes: 2 additions & 45 deletions docarray/computation/torch_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,10 @@
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union, overload
from typing import Any, List, Optional, Tuple, Union

import numpy as np
import torch

from docarray.computation.abstract_comp_backend import AbstractComputationalBackend

if TYPE_CHECKING:
from docarray.typing import TorchTensor


def _unsqueeze_if_single_axis(*matrices: torch.Tensor) -> List[torch.Tensor]:
"""Unsqueezes tensors that only have one axis, at dim 0.
Expand All @@ -32,7 +29,7 @@ def _usqueeze_if_scalar(t: torch.Tensor):
return t


class TorchCompBackend(AbstractComputationalBackend[torch.Tensor, 'TorchTensor']):
class TorchCompBackend(AbstractComputationalBackend[torch.Tensor]):
"""
Computational backend for PyTorch.
"""
Expand All @@ -43,22 +40,9 @@ def stack(
) -> 'torch.Tensor':
return torch.stack(tensors, dim=dim)

@overload
@staticmethod
def to_device(tensor: 'TorchTensor', device: str) -> 'TorchTensor':
"""Move the tensor to the specified device."""
...

@overload
@staticmethod
def to_device(tensor: 'torch.Tensor', device: str) -> 'torch.Tensor':
"""Move the tensor to the specified device."""
...

@staticmethod
def to_device(
tensor: Union['torch.Tensor', 'TorchTensor'], device: str
) -> Union['torch.Tensor', 'TorchTensor']:
return tensor.to(device)

@staticmethod
Expand Down Expand Up @@ -92,36 +76,9 @@ def none_value() -> Any:
def shape(tensor: 'torch.Tensor') -> Tuple[int, ...]:
return tuple(tensor.shape)

@overload
@staticmethod
def reshape(tensor: 'TorchTensor', shape: Tuple[int, ...]) -> 'TorchTensor':
"""
Gives a new shape to tensor without changing its data.

:param tensor: tensor to be reshaped
:param shape: the new shape
:return: a tensor with the same data and number of elements as tensor
but with the specified shape.
"""
...

@overload
@staticmethod
def reshape(tensor: 'torch.Tensor', shape: Tuple[int, ...]) -> 'torch.Tensor':
"""
Gives a new shape to tensor without changing its data.

:param tensor: tensor to be reshaped
:param shape: the new shape
:return: a tensor with the same data and number of elements as tensor
but with the specified shape.
"""
...

@staticmethod
def reshape(
tensor: Union['torch.Tensor', 'TorchTensor'], shape: Tuple[int, ...]
) -> Union['torch.Tensor', 'TorchTensor']:
"""
Gives a new shape to tensor without changing its data.

Expand Down
4 changes: 2 additions & 2 deletions docarray/typing/tensor/abstract_tensor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ def __docarray_validate_shape__(cls, t: T, shape: Tuple[Union[int, str]]) -> T:
:param shape: The shape to validate against.
:return: The validated tensor.
"""
comp_be = t.get_comp_backend()() # mypy Generics require instantiation
comp_be = t.get_comp_backend()
tshape = comp_be.shape(t)
if tshape == shape:
return t
Expand Down Expand Up @@ -202,7 +202,7 @@ def _docarray_from_native(cls: Type[T], value: Any) -> T:

@staticmethod
@abc.abstractmethod
def get_comp_backend() -> Type[AbstractComputationalBackend[TTensor, T]]:
def get_comp_backend() -> AbstractComputationalBackend:
"""The computational backend compatible with this tensor type."""
...

Expand Down
4 changes: 2 additions & 2 deletions docarray/typing/tensor/ndarray.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,11 +202,11 @@ def to_protobuf(self) -> 'NdArrayProto':
return nd_proto

@staticmethod
def get_comp_backend() -> Type['NumpyCompBackend']:
def get_comp_backend() -> 'NumpyCompBackend':
"""Return the computational backend of the tensor"""
from docarray.computation.numpy_backend import NumpyCompBackend

return NumpyCompBackend
return NumpyCompBackend()
Copy link
Copy Markdown
Member

@JohannesMessner JohannesMessner Jan 18, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it a bit of an anti pattern to crate all of these instances if we don't really need them? would it make sense to maintain a singleton instance?

Copy link
Copy Markdown
Contributor

@Jackmin801 Jackmin801 Jan 18, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh that is a good idea actually. I support this. Lets go singleton with this.

Copy link
Copy Markdown
Member Author

@samsja samsja Jan 18, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

okay good idea

Copy link
Copy Markdown
Member Author

@samsja samsja Jan 18, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

actually doing the singleton is not that trivial in python ... I think it adds complexity for few real advantages at the end. Only conceptual advantage.


def __class_getitem__(cls, item: Any, *args, **kwargs):
# see here for mypy bug: https://github.com/python/mypy/issues/14123
Expand Down
4 changes: 2 additions & 2 deletions docarray/typing/tensor/torch_tensor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -216,11 +216,11 @@ def to_protobuf(self) -> 'NdArrayProto':
return nd_proto

@staticmethod
def get_comp_backend() -> Type['TorchCompBackend']:
def get_comp_backend() -> 'TorchCompBackend':
"""Return the computational backend of the tensor"""
from docarray.computation.torch_backend import TorchCompBackend

return TorchCompBackend
return TorchCompBackend()

@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):
Expand Down