{"record":{"id":"e6d2e649e2053e2b","repo":"jax-ml/jax","slug":"jax-distributed-initialize-must-be-called-before","errorCode":null,"errorMessage":"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.","messagePattern":"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\\.","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"critical","filePath":"jax/_src/distributed.py","lineNumber":399,"sourceCode":"    RuntimeError: If :func:`~jax.distributed.initialize` is called more than once\n      or if called after the backend is already initialized.\n\n  Examples:\n\n  Suppose there are two GPU processes, and process 0 is the designated coordinator\n  with address ``10.0.0.1:1234``. To initialize the GPU cluster, run the\n  following commands before anything else.\n\n  On process 0:\n\n  >>> jax.distributed.initialize(coordinator_address='10.0.0.1:1234', num_processes=2, process_id=0)  # doctest: +SKIP\n\n  On process 1:\n\n  >>> jax.distributed.initialize(coordinator_address='10.0.0.1:1234', num_processes=2, process_id=1)  # doctest: +SKIP\n  \"\"\"\n  if xla_bridge.backends_are_initialized():\n    raise RuntimeError(\"jax.distributed.initialize() must be called before \"\n                        \"any JAX calls that might initialise the XLA backend. \"\n                        \"This includes any computation, but also calls to jax.devices, jax.device_put, and others.\")\n  if partition_index is None:\n    if slice_index is not None:\n      # Deprecation added 2025-08-05. Should be removed after 3 months.\n      warnings.warn(\n          '`slice_index` has been deprecated. Please use `partition_index` instead.',\n          DeprecationWarning,\n      )\n    partition_index = slice_index\n  global_state.initialize(coordinator_address, num_processes, process_id,\n                          local_device_ids, cluster_detection_method,\n                          initialization_timeout, coordinator_bind_address,\n                          heartbeat_timeout_seconds=heartbeat_timeout_seconds,\n                          shutdown_timeout_seconds=shutdown_timeout_seconds,\n                          partition_index=partition_index,\n                          mtls_cert_file=mtls_cert_file,\n                          mtls_key_file=mtls_key_file,","sourceCodeStart":381,"sourceCodeEnd":417,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/distributed.py#L381-L417","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Move jax.distributed.initialize() to the very top of the entry-point script, before any other jax call or import that uses JAX","Audit imports for code that calls jax.devices()/device_put/computation at import time","In notebooks, restart the kernel, run initialize in the first cell"],"exampleFix":"# before\nimport jax\nprint(jax.devices())\njax.distributed.initialize(...)  # RuntimeError\n\n# after\nimport jax\njax.distributed.initialize(...)\nprint(jax.devices())","handlingStrategy":"validation","validationCode":"import jax\nfrom jax._src import xla_bridge\nassert not xla_bridge.backends_are_initialized(), 'call jax.distributed.initialize first'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Make jax.distributed.initialize() the first JAX call in main()","Avoid jax.devices()/device_put at import time in shared modules","Restart kernels before distributed notebook runs"],"tags":["jax","distributed","xla-backend","ordering"],"backgroundTag":"library-initialization-order","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}