|
| 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