Skip to content

Commit fb0affe

Browse files
authored
test: Converting tensordot benchmark to run with CodSpeed (#761)
1 parent 1c56a0b commit fb0affe

File tree

1 file changed

+69
-0
lines changed

1 file changed

+69
-0
lines changed

benchmarks/test_tensordot.py

+69
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
import itertools
2+
3+
import sparse
4+
5+
import pytest
6+
7+
import numpy as np
8+
9+
DENSITY = 0.01
10+
11+
12+
def get_sides_ids(param):
13+
m, n, p, q = param
14+
return f"{m=}-{n=}-{p=}-{q=}"
15+
16+
17+
@pytest.fixture(
18+
params=itertools.product([10, 50], [10, 20], [20, 50], [10, 50]),
19+
ids=get_sides_ids,
20+
scope="function",
21+
)
22+
def sides(request):
23+
m, n, p, q = request.param
24+
return m, n, p, q
25+
26+
27+
def get_tensor_ids(param):
28+
left_index, right_index, left_format, right_format = param
29+
return f"{left_index=}-{right_index=}-{left_format=}-{right_format=}"
30+
31+
32+
@pytest.fixture(
33+
params=([(1, 2, "dense", "coo"), (1, 2, "coo", "coo"), (1, 1, "coo", "dense")]),
34+
ids=get_tensor_ids,
35+
scope="function",
36+
)
37+
def tensordot_args(request, sides, seed, max_size):
38+
m, n, p, q = sides
39+
if m * n * p * q >= max_size:
40+
pytest.skip()
41+
left_index, right_index, left_format, right_format = request.param
42+
rng = np.random.default_rng(seed=seed)
43+
44+
t = rng.random((m, n))
45+
46+
if left_format == "dense" and right_format == "coo":
47+
left_tensor = t
48+
right_tensor = sparse.random((m, p, n, q), density=DENSITY, format=right_format, random_state=rng)
49+
50+
if left_format == "coo" and right_format == "coo":
51+
left_tensor = sparse.random((m, p), density=DENSITY, format=left_format, random_state=rng)
52+
right_tensor = sparse.random((m, n, p, q), density=DENSITY, format=right_format, random_state=rng)
53+
54+
if left_format == "coo" and right_format == "dense":
55+
left_tensor = sparse.random((m, n, p, q), density=DENSITY, format=left_format, random_state=rng)
56+
right_tensor = t
57+
58+
return left_index, right_index, left_tensor, right_tensor
59+
60+
61+
@pytest.mark.parametrize("return_type", [np.ndarray, sparse.COO])
62+
def test_tensordot(benchmark, return_type, tensordot_args):
63+
left_index, right_index, left_tensor, right_tensor = tensordot_args
64+
65+
sparse.tensordot(left_tensor, right_tensor, axes=([0, left_index], [0, right_index]), return_type=return_type)
66+
67+
@benchmark
68+
def bench():
69+
sparse.tensordot(left_tensor, right_tensor, axes=([0, left_index], [0, right_index]), return_type=return_type)

0 commit comments

Comments
 (0)