Intra-op parallelism - limit number of cores

[elchorro]

As noted in this thread XLA may use multiple CPU cores, even when one doesn't use pmap. Is there any way to control this behaviour?

I use vmap in my code and notice that CPU utilization reaches 900%. Is there a way to expose only certain number of cores to XLA?

I work on a shared server and I am afraid I will get an angry email if I end up using up all CPU resources. Sadly, the server doesn't have a GPU (which would be really useful) and the data can't leave the server so I am forced to use CPUs only.

[mattiasmar]

I'd try this:

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' # Use 8 CPU devices

From https://jax.readthedocs.io/en/latest/notebooks/xmap_tutorial.html#from-positions-to-names-in-a-toy-neural-network