Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Benchmarking dask linalg inv on nd arrays
import numpy as np
import dask.array as da
x = da.random.random((120,500,500), chunks=(40, 500, 500)) # 80 MB chunks
X = x.compute()
@da.as_gufunc(signature="(n,n)->(n,n)", output_dtypes=float, vectorize=True)
def gufunc(x):
return np.linalg.inv(x)
def loop_asarray(x):
return da.asarray([np.linalg.inv(arr) for arr in x])
def loop_from_array(x):
# this one seems to compute() during the from_array call, ergo not lazy evaluation
return da.from_array([np.linalg.inv(arr) for arr in x])
def map_blocks(x):
return da.map_blocks(np.linalg.inv, x, chunks=x.chunks)
print("Dask:")
%timeit gufunc(x).compute()
%timeit loop_asarray(x).compute()
%timeit loop_from_array(x).compute()
%timeit map_blocks(x).compute()
print("NumPy:")
%timeit np.linalg.inv(X)
%timeit np.asarray([np.linalg.inv(n) for n in X])
#Dask:
#1.19 s ± 53 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
#10 s ± 543 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
#8.32 s ± 64.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
#1.14 s ± 6.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
#NumPy:
#655 ms ± 27.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
#813 ms ± 91.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment