def get_backend(platform=None):
# TODO(mattjj,skyewm): remove this input polymorphism after we clean up how
# 'backend' values are handled
if not isinstance(platform, (type(None), str)):
return platform
with _backend_lock:
backend_factory = _backends.get(FLAGS.jax_xla_backend)
if backend_factory is None:
msg = 'Unknown jax_xla_backend value "{}".'
raise ValueError(msg.format(FLAGS.jax_xla_backend))
backend = backend_factory(platform)
util.distributed_debug_log(("Initialized backend", backend.platform),
("process_index", backend.process_index()),
("device_count", backend.device_count()),
("local_devices", backend.local_devices()))
return backend
def backends():
global _backends
global _backends_errors
global _default_backend
with _backend_lock:
if _backends is not None:
return _backends
default_priority = -1000
_backends = {}
_backends_errors = {}
for name, (factory, priority) in _backend_factories.items():
logging.vlog(1, "Initializing backend '%s'" % name)
try:
backend = factory()
if backend is not None:
if backend.device_count() > 0:
_backends[name] = backend
util.distributed_debug_log(
("Initialized backend", backend.platform),
("process_index", backend.process_index()),
("device_count", backend.device_count()),
("local_devices", backend.local_devices()))
logging.vlog(1, "Backend '%s' initialized" % name)
if priority > default_priority:
_default_backend = backend
default_priority = priority
except Exception as err:
if name in ('cpu', 'interpreter'):
# We always expect the CPU and interpreter backends to initialize
# successfully.
raise
else:
# If the backend isn't built into the binary, or if it has no devices,
# we expect a RuntimeError.
logging.info("Unable to initialize backend '%s': %s" %
(name, err))
_backends_errors[name] = str(err)
continue
if _default_backend.platform == "cpu" and FLAGS.jax_platform_name != 'cpu':
logging.warning(
'No GPU/TPU found, falling back to CPU. '
'(Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)')
return _backends
def _init_backend(platform):
factory, unused_priority = _backend_factories.get(platform, (None, None))
if factory is None:
raise RuntimeError(f"Unknown backend '{platform}'")
logging.vlog(1, "Initializing backend '%s'" % platform)
backend = factory()
# TODO(skye): consider raising more descriptive errors directly from backend
# factories instead of returning None.
if backend is None:
raise RuntimeError(f"Could not initialize backend '{platform}'")
if backend.device_count() == 0:
raise RuntimeError(f"Backend '{platform}' provides no devices.")
util.distributed_debug_log(("Initialized backend", backend.platform),
("process_index", backend.process_index()),
("device_count", backend.device_count()),
("local_devices", backend.local_devices()))
logging.vlog(1, "Backend '%s' initialized" % platform)
return backend
def backends():
global _backends
with _backend_lock:
if _backends is not None:
return _backends
_backends = {}
for name, factory in _backend_factories.items():
logging.vlog(1, "Initializing backend '%s'" % name)
try:
backend = factory()
if backend is not None:
if backend.device_count() > 0:
_backends[name] = backend
util.distributed_debug_log(
("Initialized backend", backend.platform),
("process_index", backend.process_index()),
("device_count", backend.device_count()),
("local_devices", backend.local_devices()))
logging.vlog(1, "Backend '%s' initialized" % name)
except (RuntimeError, ImportError) as err:
if name in ('cpu', 'interpreter'):
# We always expect the CPU and interpreter backends to initialize
# successfully.
raise
else:
# If the backend isn't built into the binary, or if it has no devices,
# we expect a RuntimeError.
logging.info("Unable to initialize backend '%s': %s" %
(name, err))
continue
if list(_backends.keys())[-1] == "cpu":
logging.warning(
'No GPU/TPU found, falling back to CPU. '
'(Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)')
return _backends