Limit jax multithreading

By default jax appears to multithread most operations, eg.

x = jr.normal(jrkey, shape=(50000, 50000))
x @ x

will run across all available cores. This is great in general, and matches numpy’s behavior. But it presents problems when trying to run a bunch of small operations in parallel, eg. running the same script initialized with 4 different random seeds on a 4-core machine.

Is there any option in jax to cap the number of threads that it uses? Something like https://stackoverflow.com/questions/17053671/python-how-do-you-stop-numpy-from-multithreading?

[samuela]

For my own (and other’s) future googling, my current approach looks like

from multiprocessing import get_context
import os

# 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")

def job(random_seed: int):
  # jax jax jax

if __name__ == "__main__":
  # See https://codewithoutrules.com/2018/09/04/python-multiprocessing/.
  with get_context("spawn").Pool() as pool:
    pool.imap_unordered(job, range(100))

There may be a way better way, but it seems to work 🤷‍♀️

[mattjj]

Did you try setting these environment variables? (My comment didn’t explain this very well.)

XLA_FLAGS="--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1" python my_file.py

That seems to work for me in a test, at least for a big matmul.