Parallelizing code to CPU cores

[marcosrdac]

Hello, huys,

I'm trying to make my code parallel at the level of my CPU cores. #5506 and #2714 talk about this environment variable that needs to be set so that jax treats each CPU core as a separate device. I tried doing it both from shell, with an export statement, and from python's os module (approach written below). This is a simple code that runs fine when I set a core for work:

from os import environ
from jax import numpy as jnp, random, pmap

CORES = 4
environ['XLA_FLAGS'] = f'--xla_force_host_platform_device_count={CORES}'

# function to generate a normal distributions for each core used 
pnormal = pmap(random.normal, static_broadcasted_argnums=1)

keys = random.split(random.PRNGKey(0), CORES)
dists = pnormal(keys, (3, 3))  # already gives me an error for more than one core

result = pmap(lambda x: x.T@x)(dists)
print(result)

but when I try to use more than one of them (I'm using an i5, with 4 cores) I keep getting an error, which is the same I get when no flag is set at all:

ValueError: compiling computation that requires 4 logical devices, but only 1 XLA devices are available (num_replicas=4, num_partitions=1)

Full error message:

Traceback (most recent call last):
  File "pmap0.py", line 10, in <module>
    dists = pnormal(keys, (3, 3))
  File "/home/marcosrdac/.local/share/python_envs/m/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 183, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/marcosrdac/.local/share/python_envs/m/lib/python3.8/site-packages/jax/_src/api.py", line 1647, in f_pmapped
    out = pxla.xla_pmap(
  File "/home/marcosrdac/.local/share/python_envs/m/lib/python3.8/site-packages/jax/core.py", line 1620, in bind
    return call_bind(self, fun, *args, **params)
  File "/home/marcosrdac/.local/share/python_envs/m/lib/python3.8/site-packages/jax/core.py", line 1551, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/home/marcosrdac/.local/share/python_envs/m/lib/python3.8/site-packages/jax/core.py", line 1623, in process
    return trace.process_map(self, fun, tracers, params)
  File "/home/marcosrdac/.local/share/python_envs/m/lib/python3.8/site-packages/jax/core.py", line 606, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/home/marcosrdac/.local/share/python_envs/m/lib/python3.8/site-packages/jax/interpreters/pxla.py", line 624, in xla_pmap_impl
    compiled_fun, fingerprint = parallel_callable(fun, backend, axis_name, axis_size,
  File "/home/marcosrdac/.local/share/python_envs/m/lib/python3.8/site-packages/jax/linear_util.py", line 262, in memoized_fun
    ans = call(fun, *args)
  File "/home/marcosrdac/.local/share/python_envs/m/lib/python3.8/site-packages/jax/interpreters/pxla.py", line 809, in parallel_callable
    raise ValueError(msg.format(num_global_shards, xb.device_count(backend),
jax._src.traceback_util.UnfilteredStackTrace: ValueError: compiling computation that requires 4 logical devices, but only 1 XLA devices are available (num_replicas=4, num_partitions=1)

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "pmap0.py", line 10, in <module>
    dists = pnormal(keys, (3, 3))
  File "/home/marcosrdac/.local/share/python_envs/m/lib/python3.8/site-packages/jax/interpreters/pxla.py", line 809, in parallel_callable
    raise ValueError(msg.format(num_global_shards, xb.device_count(backend),
ValueError: compiling computation that requires 4 logical devices, but only 1 XLA devices are available (num_replicas=4, num_partitions=1)

I wonder what I'm doing wrong. Do you have any idea?

---

EDIT: What's yet more odd is... I just tried it on a i7 and it works fine there.

[PhilipVinc]

I think the flag must be set before you load jax

[marcosrdac]

I actually tried both ways to set variables (from shell, via an export, and from inside the Python code, via the os module). It is odd that both of them work on my i7, but none of them do on my i5.