Skip to content
This repository was archived by the owner on Sep 9, 2026. It is now read-only.

Commit 20f634d

Browse files
JohannesMessnerJohannesMessner
authored
refactor(v2): computational backends (#941)
* refactor: make abstract tensor parametrization modular Signed-off-by: Johannes Messner <[email protected]> * feat: add type for torch embedding Signed-off-by: Johannes Messner <[email protected]> * feat: embedding type for ndarray Signed-off-by: Johannes Messner <[email protected]> * fix: fix general embedding type Signed-off-by: Johannes Messner <[email protected]> * test: update tests Signed-off-by: Johannes Messner <[email protected]> * feat: find function Signed-off-by: Johannes Messner <[email protected]> * feat: batched query input Signed-off-by: Johannes Messner <[email protected]> * fix: fix metrics and add tests for them Signed-off-by: Johannes Messner <[email protected]> * test: add tests for find Signed-off-by: Johannes Messner <[email protected]> * test: add test for topk Signed-off-by: Johannes Messner <[email protected]> * feat: add batched find Signed-off-by: Johannes Messner <[email protected]> * test: add more tests Signed-off-by: Johannes Messner <[email protected]> * docs: improve docstrings Signed-off-by: Johannes Messner <[email protected]> * docs: improve docstrings Signed-off-by: Johannes Messner <[email protected]> * fix: mypy and some comments Signed-off-by: Johannes Messner <[email protected]> * refactor: add computational backends Signed-off-by: Johannes Messner <[email protected]> * refactor: use comp backends in find function Signed-off-by: Johannes Messner <[email protected]> * refactor: clean up file structure Signed-off-by: Johannes Messner <[email protected]> * fix: some typing issues Signed-off-by: Johannes Messner <[email protected]> * fix: reduce mypy errors Signed-off-by: Johannes Messner <[email protected]> * fix: mypy Signed-off-by: Johannes Messner <[email protected]> * fix: mypy Signed-off-by: Johannes Messner <[email protected]> * refactor: remove is_tensor flag Signed-off-by: Johannes Messner <[email protected]> Signed-off-by: Johannes Messner <[email protected]> Signed-off-by: Johannes Messner <[email protected]>
1 parent 4d85119 commit 20f634d

28 files changed

Lines changed: 834 additions & 665 deletions

‎docarray/computation/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
from docarray.computation.abstract_comp_backend import AbstractComputationalBackend
2+
3+
__all__ = ['AbstractComputationalBackend']

‎docarray/computation/abstract_comp_backend.py‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
import typing
2+
from abc import ABC, abstractmethod
3+
from typing import List, Optional, Tuple, TypeVar, Union
4+
5+
# In practice all of the below will be the same type
6+
TTensor = TypeVar('TTensor')
7+
TTensorRetrieval = TypeVar('TTensorRetrieval')
8+
TTensorMetrics = TypeVar('TTensorMetrics')
9+
10+
11+
class AbstractComputationalBackend(ABC, typing.Generic[TTensor]):
12+
"""
13+
Abstract base class for computational backends.
14+
Every supported tensor/ML framework (numpy, torch etc.) should define its own
15+
computational backend exposing common functionality expressed in that framework.
16+
That way, DocArray can leverage native implementations from all frameworks.
17+
"""
18+
19+
@staticmethod
20+
@abstractmethod
21+
def stack(
22+
tensors: Union[List['TTensor'], Tuple['TTensor']], dim: int = 0
23+
) -> 'TTensor':
24+
"""
25+
Stack a list of tensors along a new axis.
26+
"""
27+
...
28+
29+
class Retrieval(ABC, typing.Generic[TTensorRetrieval]):
30+
"""
31+
Abstract class for retrieval and ranking functionalities
32+
"""
33+
34+
@staticmethod
35+
@abstractmethod
36+
def top_k(
37+
values: 'TTensorRetrieval',
38+
k: int,
39+
descending: bool = False,
40+
device: Optional[str] = None,
41+
) -> Tuple['TTensorRetrieval', 'TTensorRetrieval']:
42+
"""
43+
Retrieves the top k smallest values in `values`,
44+
and returns them alongside their indices in the input `values`.
45+
Can also be used to retrieve the top k largest values,
46+
by setting the `descending` flag to True.
47+
48+
:param values: Tensor of values to rank.
49+
Should be of shape (n_queries, n_values_per_query).
50+
Inputs of shape (n_values_per_query,) will be expanded
51+
to (1, n_values_per_query).
52+
:param k: number of values to retrieve
53+
:param descending: retrieve largest values instead of smallest values
54+
:param device: the computational device to use.
55+
:return: Tuple containing the retrieved values, and their indices.
56+
Both ar of shape (n_queries, k)
57+
"""
58+
...
59+
60+
class Metrics(ABC, typing.Generic[TTensorMetrics]):
61+
"""
62+
Abstract base class for metrics (distances and similarities).
63+
"""
64+
65+
@staticmethod
66+
@abstractmethod
67+
def cosine_sim(
68+
x_mat: 'TTensorMetrics',
69+
y_mat: 'TTensorMetrics',
70+
eps: float = 1e-7,
71+
device: Optional[str] = None,
72+
) -> 'TTensorMetrics':
73+
"""Pairwise cosine similarities between all vectors in x_mat and y_mat.
74+
75+
:param x_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the
76+
number of vectors and n_dim is the number of dimensions of each example.
77+
:param y_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the
78+
number of vectors and n_dim is the number of dimensions of each example.
79+
:param eps: a small jitter to avoid divde by zero
80+
:param device: the device to use for computations.
81+
If not provided, the devices of x_mat and y_mat are used.
82+
:return: Tensor of shape (n_vectors, n_vectors) containing all pairwise
83+
cosine distances.
84+
The index [i_x, i_y] contains the cosine distance between
85+
x_mat[i_x] and y_mat[i_y].
86+
"""
87+
...
88+
89+
@staticmethod
90+
@abstractmethod
91+
def euclidean_dist(
92+
x_mat: 'TTensorMetrics',
93+
y_mat: 'TTensorMetrics',
94+
device: Optional[str] = None,
95+
) -> 'TTensorMetrics':
96+
"""Pairwise Euclidian distances between all vectors in x_mat and y_mat.
97+
98+
:param x_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the
99+
number of vectors and n_dim is the number of dimensions of each example.
100+
:param y_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the
101+
number of vectors and n_dim is the number of dimensions of each example.
102+
:param device: the device to use for pytorch computations.
103+
If not provided, the devices of x_mat and y_mat are used.
104+
:return: Tensor of shape (n_vectors, n_vectors) containing all pairwise
105+
euclidian distances.
106+
The index [i_x, i_y] contains the euclidian distance between
107+
x_mat[i_x] and y_mat[i_y].
108+
"""
109+
...
110+
111+
@staticmethod
112+
@abstractmethod
113+
def sqeuclidean_dist(
114+
x_mat: 'TTensorMetrics',
115+
y_mat: 'TTensorMetrics',
116+
device: Optional[str] = None,
117+
) -> 'TTensorMetrics':
118+
"""Pairwise Squared Euclidian distances between all vectors
119+
in x_mat and y_mat.
120+
121+
:param x_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the
122+
number of vectors and n_dim is the number of dimensions of each
123+
example.
124+
:param y_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the
125+
number of vectors and n_dim is the number of dimensions of each
126+
example.
127+
:param device: the device to use for pytorch computations.
128+
If not provided, the devices of x_mat and y_mat are used.
129+
:return: Tensor of shape (n_vectors, n_vectors) containing all pairwise
130+
euclidian distances.
131+
The index [i_x, i_y] contains the euclidian distance between
132+
x_mat[i_x] and y_mat[i_y].
133+
"""
134+
...

‎docarray/computation/numpy_backend.py‎

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
import warnings
2+
from typing import List, Optional, Tuple, Union
3+
4+
import numpy as np
5+
6+
from docarray.computation import AbstractComputationalBackend
7+
8+
9+
def _expand_if_single_axis(*matrices: np.ndarray) -> List[np.ndarray]:
10+
"""Expands arrays that only have one axis, at dim 0.
11+
This ensures that all outputs can be treated as matrices, not vectors.
12+
13+
:param matrices: Matrices to be expanded
14+
:return: List of the input matrices,
15+
where single axis matrices are expanded at dim 0.
16+
"""
17+
expanded = []
18+
for m in matrices:
19+
if len(m.shape) == 1:
20+
expanded.append(np.expand_dims(m, axis=0))
21+
else:
22+
expanded.append(m)
23+
return expanded
24+
25+
26+
def _expand_if_scalar(arr: np.ndarray) -> np.ndarray:
27+
if len(arr.shape) == 0: # avoid scalar output
28+
arr = np.expand_dims(arr, axis=0)
29+
return arr
30+
31+
32+
class NumpyCompBackend(AbstractComputationalBackend[np.ndarray]):
33+
"""
34+
Computational backend for Numpy.
35+
"""
36+
37+
@staticmethod
38+
def stack(
39+
tensors: Union[List['np.ndarray'], Tuple['np.ndarray']], dim: int = 0
40+
) -> 'np.ndarray':
41+
return np.stack(tensors, axis=dim)
42+
43+
class Retrieval(AbstractComputationalBackend.Retrieval[np.ndarray]):
44+
"""
45+
Abstract class for retrieval and ranking functionalities
46+
"""
47+
48+
@staticmethod
49+
def top_k(
50+
values: 'np.ndarray',
51+
k: int,
52+
descending: bool = False,
53+
device: Optional[str] = None,
54+
) -> Tuple['np.ndarray', 'np.ndarray']:
55+
"""
56+
Retrieves the top k smallest values in `values`,
57+
and returns them alongside their indices in the input `values`.
58+
Can also be used to retrieve the top k largest values,
59+
by setting the `descending` flag.
60+
61+
:param values: Torch tensor of values to rank.
62+
Should be of shape (n_queries, n_values_per_query).
63+
Inputs of shape (n_values_per_query,) will be expanded
64+
to (1, n_values_per_query).
65+
:param k: number of values to retrieve
66+
:param descending: retrieve largest values instead of smallest values
67+
:param device: Not supported for this backend
68+
:return: Tuple containing the retrieved values, and their indices.
69+
Both ar of shape (n_queries, k)
70+
"""
71+
if device is not None:
72+
warnings.warn('`device` is not supported for numpy operations')
73+
74+
if len(values.shape) == 1:
75+
values = np.expand_dims(values, axis=0)
76+
77+
if descending:
78+
values = -values
79+
80+
if k >= values.shape[1]:
81+
idx = values.argsort(axis=1)[:, :k]
82+
values = np.take_along_axis(values, idx, axis=1)
83+
else:
84+
idx_ps = values.argpartition(kth=k, axis=1)[:, :k]
85+
values = np.take_along_axis(values, idx_ps, axis=1)
86+
idx_fs = values.argsort(axis=1)
87+
idx = np.take_along_axis(idx_ps, idx_fs, axis=1)
88+
values = np.take_along_axis(values, idx_fs, axis=1)
89+
90+
if descending:
91+
values = -values
92+
93+
return values, idx
94+
95+
class Metrics(AbstractComputationalBackend.Metrics[np.ndarray]):
96+
"""
97+
Abstract base class for metrics (distances and similarities).
98+
"""
99+
100+
@staticmethod
101+
def cosine_sim(
102+
x_mat: np.ndarray,
103+
y_mat: np.ndarray,
104+
eps: float = 1e-7,
105+
device: Optional[str] = None,
106+
) -> np.ndarray:
107+
"""Pairwise cosine similarities between all vectors in x_mat and y_mat.
108+
109+
:param x_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is
110+
the number of vectors and n_dim is the number of dimensions of each
111+
example.
112+
:param y_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is
113+
the number of vectors and n_dim is the number of dimensions of each
114+
example.
115+
:param eps: a small jitter to avoid divde by zero
116+
:param device: Not supported for this backend
117+
:return: np.ndarray of shape (n_vectors, n_vectors) containing all
118+
pairwise cosine distances.
119+
The index [i_x, i_y] contains the cosine distance between
120+
x_mat[i_x] and y_mat[i_y].
121+
"""
122+
if device is not None:
123+
warnings.warn('`device` is not supported for numpy operations')
124+
125+
x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat)
126+
127+
sims = np.clip(
128+
(np.dot(x_mat, y_mat.T) + eps)
129+
/ (
130+
np.outer(
131+
np.linalg.norm(x_mat, axis=1), np.linalg.norm(y_mat, axis=1)
132+
)
133+
+ eps
134+
),
135+
-1,
136+
1,
137+
).squeeze()
138+
return _expand_if_scalar(sims)
139+
140+
@classmethod
141+
def euclidean_dist(
142+
cls, x_mat: np.ndarray, y_mat: np.ndarray, device: Optional[str] = None
143+
) -> np.ndarray:
144+
"""Pairwise Euclidian distances between all vectors in x_mat and y_mat.
145+
146+
:param x_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is
147+
the number of vectors and n_dim is the number of dimensions of each
148+
example.
149+
:param y_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is
150+
the number of vectors and n_dim is the number of dimensions of each
151+
example.
152+
:param eps: a small jitter to avoid divde by zero
153+
:param device: Not supported for this backend
154+
:return: np.ndarray of shape (n_vectors, n_vectors) containing all
155+
pairwise euclidian distances.
156+
The index [i_x, i_y] contains the euclidian distance between
157+
x_mat[i_x] and y_mat[i_y].
158+
"""
159+
if device is not None:
160+
warnings.warn('`device` is not supported for numpy operations')
161+
162+
x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat)
163+
164+
return _expand_if_scalar(
165+
np.sqrt(cls.sqeuclidean_dist(x_mat, y_mat)).squeeze()
166+
)
167+
168+
@staticmethod
169+
def sqeuclidean_dist(
170+
x_mat: np.ndarray,
171+
y_mat: np.ndarray,
172+
device: Optional[str] = None,
173+
) -> np.ndarray:
174+
"""Pairwise Squared Euclidian distances between all vectors in
175+
x_mat and y_mat.
176+
177+
:param x_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is
178+
the number of vectors and n_dim is the number of dimensions of each
179+
example.
180+
:param y_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is
181+
the number of vectors and n_dim is the number of dimensions of each
182+
example.
183+
:param device: Not supported for this backend
184+
:return: np.ndarray of shape (n_vectors, n_vectors) containing all
185+
pairwise Squared Euclidian distances.
186+
The index [i_x, i_y] contains the cosine Squared Euclidian between
187+
x_mat[i_x] and y_mat[i_y].
188+
"""
189+
eps: float = 1e-7 # avoid problems with numerical inaccuracies
190+
191+
if device is not None:
192+
warnings.warn('`device` is not supported for numpy operations')
193+
194+
x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat)
195+
196+
dists = (
197+
np.sum(y_mat**2, axis=1)
198+
+ np.sum(x_mat**2, axis=1)[:, np.newaxis]
199+
- 2 * np.dot(x_mat, y_mat.T)
200+
).squeeze()
201+
202+
# remove numerical artifacts
203+
dists = np.where(np.logical_and(dists < 0, dists > -eps), 0, dists)
204+
return _expand_if_scalar(dists)

0 commit comments

Comments
 (0)