def force_multi_devices(num_cpu_devices: int):
  """Run with set number of CPU devices."""
  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 +
        ' --xla_force_host_platform_device_count={}'.format(num_cpu_devices))
  # Clear any cached backends so new CPU backend will pick up the env var.
  xla_bridge.get_backend.cache_clear()
Run with set number of CPU devices.
language/mentionmemory/utils/test_utils.py