# %%
import numpy as np
import ot
from joblib import Parallel, delayed
# %%
n = 1000
k = 50
M = np.random.rand(k, n, n)
a = np.ones(n) / n
# %% loop
def emd(M, axes=None):
 return ot.emd(a, a, M)
R_loop = np.zeros((k, n, n))
ot.tic()
for i in range(k):
 R_loop[i] = emd(M[i])
ot.toc()
# %% numpy take/stack
def apply_across_axis(func, M, axis=0):
 return np.stack([
 func(M.take(i, axis))
 for i in range(M.shape[axis])
 ], axis=axis)
ot.tic()
R_numpy = apply_across_axis(emd, M, 0)
ot.toc()
# %% joblib?
def apply_across_axis_joblib(func, M, axis=0, n_jobs=4):
 res = Parallel(n_jobs=n_jobs, max_nbytes=None)(
 delayed(func)(M.take(i, axis))
 for i in range(M.shape[axis])
 )
 return np.stack(res, axis=axis)
R_joblib = apply_across_axis_joblib(emd, M[:4], 0)
ot.tic()
R_joblib = apply_across_axis_joblib(emd, M, 0)
ot.toc()