Skip to content

Commit e5e1a99

Browse files
committed
Add in scipy/sparse to minmax benchmark
1 parent 16168d1 commit e5e1a99

File tree

2 files changed

+35
-5
lines changed

2 files changed

+35
-5
lines changed

numpy/minmax.py

+20-5
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,14 @@
1-
import numpy as np
1+
mport numpy as np
22
from scipy.sparse import random, csr_matrix
33
import sparse
44
import pytest
55
import os
6-
from util import MinMaxPydataSparseTensorLoader
6+
from util import MinMaxPydataSparseTensorLoader, MinMaxScipySparseTensorLoader
77

88
@pytest.mark.parametrize("dims", [1, 3, 5])
99
def bench_pydata_minmax(tacoBench, dims):
1010
loader = MinMaxPydataSparseTensorLoader()
1111
dims_list = [20] + [20] + [43 for ele in range(dims)]
12-
#FIXME: loader.random is always between 0 and 1, need to be larger. multiply by some value and then store to tns file
13-
#TODO: matrix shouldn't be completely random. it should have blocks of dense values (to simulate pruning)
14-
# and not just sparse uniform sampling
1512

1613
matrix = loader.tensor(dims_list)
1714
extra_info = dict()
@@ -26,6 +23,24 @@ def bench():
2623
return reduced
2724
tacoBench(bench, extra_info, True)
2825

26+
@pytest.mark.parametrize("dims", [1, 3, 5])
27+
def bench_scipy_minmax(tacoBench, dims):
28+
loader = MinMaxScipySparseTensorLoader()
29+
dims_list = [20] + [20] + [43 for ele in range(dims)]
30+
31+
matrix = loader.tensor(dims_list)
32+
extra_info = dict()
33+
extra_info["nnz"] = matrix.nnz
34+
def bench():
35+
reduced = matrix
36+
for m in range(len(dims_list)):
37+
if m % 2 == 0:
38+
reduced = reduced.min(-1)
39+
else:
40+
reduced = reduced.max(-1)
41+
return reduced
42+
tacoBench(bench, extra_info, True)
43+
2944
@pytest.mark.skip(reason="Only to get matrix statistics")
3045
@pytest.mark.parametrize("dims", [1, 3, 5])
3146
def bench_minmax_statistics(tacoBench, dims):

numpy/util.py

+15
Original file line numberDiff line numberDiff line change
@@ -512,3 +512,18 @@ def tensor(self, shape, variant=None):
512512

513513
def generate_crds(self, shape):
514514
return generate_crds_helper(shape, 0, [[0]])
515+
516+
class MinMaxScipySparseTensorLoader:
517+
def __init__(self):
518+
self.loader = ScipySparseTensorLoader("csr")
519+
520+
def tensor(self, shape, variant=None):
521+
key = construct_minmax_tensor_key(shape)
522+
# If a tensor with these properties exists already, then load it.
523+
if os.path.exists(key):
524+
return self.loader.load(key)
525+
else:
526+
# Otherwise, we must create a random tensor with the desired properties,
527+
# dump it to the output file, then return it.
528+
raise NotImplementedError
529+

0 commit comments

Comments
 (0)