Using mp.Pool() hangs while using jax library

[mpvanderschelling]

I've got the following issue
I'm using a multiprocessing library (pathos.multiprocessing to be exact), and I'm parallelizing a function run_optimization

def run_optimization(
    optimizer: Optimizer,
    function: Function,
    sampler: Sampler,
    iterations: int,
    seed: int,
    number_of_samples: int = 30,
) -> ExperimentData:

Optimizer, Function and Sampler are all custom object. I'm rerunning this function-call with a multitude of seeds. I've parallized this with the mp.Pool() function:

seed = np.random.randint(low=0, high=1e5)

if parallelization:
    args = [
        (optimizer, function, sampler, iterations,
         seed + index, number_of_samples)
        for index, _ in enumerate(range(realizations))
    ]

    with mp.Pool() as pool:
        results = pool.starmap(run_optimization, args)

This all worked fine, until I ported a jax-based optimizer (CMA_ES from the library evosax).
Now, the call to pool.starmap makes the program hang, and it seems like it is waiting on resources:

File ~/miniconda3/envs/l2co/lib/python3.8/site-packages/multiprocess/pool.py:765, in ApplyResult.get(self, timeout)
    764 def get(self, timeout=None):
--> 765     self.wait(timeout)
    766     if not self.ready():
    767         raise TimeoutError

I think that the problem is that evosax (and thus jax) is using multiple cores to do it's internal computation, and thus no resources for multithreading with pathos are available, thus it hangs.

What I've tried solutions in #743 and #2685 by setting environ variables when I first import jax:

import jax
from evosax import CMA_ES
from jax._src.typing import Array

# Limit ourselves to single-threaded jax/xla operations to avoid thrashing. See
# https://github.com/google/jax/issues/743.
os.environ["XLA_FLAGS"] = ("--xla_cpu_multi_thread_eigen=false "
                           "intra_op_parallelism_threads=1")
os.environ['MKL_NUM_THREADS'] = '1'  # set number of MKL threads to run in parallel
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['OMP_NUM_THREADS'] = '1'

os.environ["NUM_INTER_THREADS"] = "1"
os.environ["NUM_INTRA_THREADS"] = "1"

# set XLA threads and parallelism
os.environ["XLA_FLAGS"] = "--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1"
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "0"

However I still get the same results: the call to mp.Pool() hangs because it cannot find the necessary resources.
Can someone point me in the right direction to fix this?