embarassingly parallel loop across numpy arrays

loop.py

def loop(X, func, n_jobs=-1, *args, **kwargs):
    """"will apply func(x, *args) for x in X in parallel"""
    from joblib import Parallel, delayed, cpu_count
    max_jobs = cpu_count()
    n_jobs = max_jobs if n_jobs==-1 else n_jobs
    n_jobs = max_jobs if n_jobs>max_jobs else n_jobs
    n_jobs = min(n_jobs, len(X))
    
    parallel = Parallel(n_jobs=n_jobs)
    p_func = delayed(_loop)
    X_splits = np.array_split(X, n_jobs, axis=0)
    out = parallel(p_func(func, x, *args, **kwargs) for x in X_splits)
    out = np.array(out)
    if out.ndim > 1:
        out = np.concatenate(out, axis=0)
    return out
    

def _loop(func, X, *args, **kwargs):
    out = list()
    for x in tqdm(X):
        out.append(func(x, *args, **kwargs))
    return out