jax-ml/jax · critical · RuntimeError

jax.distributed.initialize() must be called before any JAX c

Error message

jax.distributed.initialize() must be called before any JAX calls that might initialise the XLA backend. This includes any computation, but also calls to jax.devices, jax.device_put, and others.

What it means

jax.distributed.initialize() must run before anything that initializes the XLA backend, because collective communication needs to configure the backend with the distributed process topology up front. Once any backend is initialized, distributed setup can no longer take effect.

Source

Thrown at jax/_src/distributed.py:399

    RuntimeError: If :func:`~jax.distributed.initialize` is called more than once
      or if called after the backend is already initialized.

  Examples:

  Suppose there are two GPU processes, and process 0 is the designated coordinator
  with address ``10.0.0.1:1234``. To initialize the GPU cluster, run the
  following commands before anything else.

  On process 0:

  >>> jax.distributed.initialize(coordinator_address='10.0.0.1:1234', num_processes=2, process_id=0)  # doctest: +SKIP

  On process 1:

  >>> jax.distributed.initialize(coordinator_address='10.0.0.1:1234', num_processes=2, process_id=1)  # doctest: +SKIP
  """
  if xla_bridge.backends_are_initialized():
    raise RuntimeError("jax.distributed.initialize() must be called before "
                        "any JAX calls that might initialise the XLA backend. "
                        "This includes any computation, but also calls to jax.devices, jax.device_put, and others.")
  if partition_index is None:
    if slice_index is not None:
      # Deprecation added 2025-08-05. Should be removed after 3 months.
      warnings.warn(
          '`slice_index` has been deprecated. Please use `partition_index` instead.',
          DeprecationWarning,
      )
    partition_index = slice_index
  global_state.initialize(coordinator_address, num_processes, process_id,
                          local_device_ids, cluster_detection_method,
                          initialization_timeout, coordinator_bind_address,
                          heartbeat_timeout_seconds=heartbeat_timeout_seconds,
                          shutdown_timeout_seconds=shutdown_timeout_seconds,
                          partition_index=partition_index,
                          mtls_cert_file=mtls_cert_file,
                          mtls_key_file=mtls_key_file,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Move jax.distributed.initialize() to the very top of the entry-point script, before any other jax call or import that uses JAX
2. Audit imports for code that calls jax.devices()/device_put/computation at import time
3. In notebooks, restart the kernel, run initialize in the first cell

Example fix

# before
import jax
print(jax.devices())
jax.distributed.initialize(...)  # RuntimeError

# after
import jax
jax.distributed.initialize(...)
print(jax.devices())

import jax
from jax._src import xla_bridge
assert not xla_bridge.backends_are_initialized(), 'call jax.distributed.initialize first'

When it happens

Trigger: Calling jax.devices(), jax.device_put(), or any computation (jit, grad, etc.) before jax.distributed.initialize(); initialize must be the first JAX call in the process.

Common situations: Importing a module that eagerly queries devices or does warm-up compute; logging device info at startup before distributed init; notebook cells run out of order; profiler/tracer imports that touch the backend.

Related errors

• Mapped away dimension of inputs passed to vmap should be sha
• The 'sharding' attribute is not available on {self._error_re
• The is_fully_addressable property was called on {self._error
• mTLS for the JAX distributed service requires jaxlib 0.11.2
• coordinator_address should be defined.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/e6d2e649e2053e2b. Report an issue: GitHub.