How to use JAX pmap with CPU cores

[Deepakgthomas]

I am trying to use JAX pmap but I am getting the error that XLA devices aren't visible -
Here's my code -

import jax.numpy as jnp
import os
from jax import pmap
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'

out = pmap(lambda x: x ** 2)(jnp.arange(8))
print(out)

Here's the error -

Traceback (most recent call last):
  File "new.py", line 6, in <module>
    out = pmap(lambda x: x ** 2)(jnp.arange(8))
  File "/home/thoma/anaconda3/envs/tbd/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/thoma/anaconda3/envs/tbd/lib/python3.8/site-packages/jax/_src/api.py", line 1779, in cache_miss
    execute = pxla.xla_pmap_impl_lazy(fun_, *tracers, **params)
  File "/home/thoma/anaconda3/envs/tbd/lib/python3.8/site-packages/jax/_src/interpreters/pxla.py", line 411, in xla_pmap_impl_lazy
    compiled_fun, fingerprint = parallel_callable(
  File "/home/thoma/anaconda3/envs/tbd/lib/python3.8/site-packages/jax/_src/linear_util.py", line 345, in memoized_fun
    ans = call(fun, *args)
  File "/home/thoma/anaconda3/envs/tbd/lib/python3.8/site-packages/jax/_src/interpreters/pxla.py", line 682, in parallel_callable
    pmap_executable = pmap_computation.compile()
  File "/home/thoma/anaconda3/envs/tbd/lib/python3.8/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/home/thoma/anaconda3/envs/tbd/lib/python3.8/site-packages/jax/_src/interpreters/pxla.py", line 923, in compile
    executable = UnloadedPmapExecutable.from_hlo(
  File "/home/thoma/anaconda3/envs/tbd/lib/python3.8/site-packages/jax/_src/interpreters/pxla.py", line 993, in from_hlo
    raise ValueError(msg.format(shards.num_global_shards,
jax._src.traceback_util.UnfilteredStackTrace: ValueError: compiling computation that requires 8 logical devices, but only 1 XLA devices are available (num_replicas=8)

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 "new.py", line 6, in <module>
    out = pmap(lambda x: x ** 2)(jnp.arange(8))
ValueError: compiling computation that requires 8 logical devices, but only 1 XLA devices are available (num_replicas=8)

Based on this discussion, I did this os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'. Therefore I am not sure what to do now.

[jakevdp]

Thanks for the question. The code you pasted is correct, and works when I run it in a clean environment (i.e. one where no jax operations have previously been executed). Can ensure you're running this in a clean environment, and if it still doesn't work, can you give us more information about your system? What is the output of jax.print_environment_info()?

[Deepakgthomas]

Thanks for the reply. Yes, it is a clean environment, because I am running this code as a program using PyCharm and not as a Jupyter Notebook.

When I do - print(jax.print_environment_info()), the output I get is None

Please let me know if you have any questions. I added a detailed description of the problem here .

[jakevdp]

Thanks for providing those runtime details - the issue is that you're running on a system where the default device is GPU, and the --xla_force_host_platform_device_count=8 only affects the host (CPU) device count. You can force pmap to run on the non-default CPU devices using the devices argument:

out = pmap(lambda x: x ** 2, devices=jax.devices('cpu'))(jnp.arange(8))

[wingos80]

Thanks for the clarification on using pmap on GPU vs CPU. It sounds like that it would not be possible to use pmap across multiple gpu cores, but instead you can only use pmap across multiple gpu devices, would that be correct?