How to set the number of devices during a test/function on CPU?

[cgarciae]

You can use this decorator to set the number of devices for a specific test or function:

def device_count(n: int):
  def decorator(f):
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
      prev_xla_flags = os.getenv('XLA_FLAGS')
      flags_str = prev_xla_flags or ''
      # Don't override user-specified device count, or other XLA flags.
      if 'xla_force_host_platform_device_count' not in flags_str:
        os.environ['XLA_FLAGS'] = (
          flags_str + f' --xla_force_host_platform_device_count={n}'
        )
      xla_bridge.get_backend.cache_clear()

      try:
        return f(*args, **kwargs)
      finally:
        if prev_xla_flags is None:
          del os.environ['XLA_FLAGS']
        else:
          os.environ['XLA_FLAGS'] = prev_xla_flags
        xla_bridge.get_backend.cache_clear()

    return wrapper

  return decorator

Example usage:

class TestFoo:
  @device_count(6)
  def test_bar(self):
    assert len(jax.devices()) == 6

[yashk2810]

Note that this probably won't work. xla_force_host_platform_device_count needs to be called before backend is initialized. Calling it after should not affect anything AFAIK.