`slice_index` has been deprecated. Please use `partition_ind

Error message

`slice_index` has been deprecated. Please use `partition_index` instead.

What it means

jax.distributed.initialize() was called with the slice_index keyword argument, renamed to partition_index (deprecation added 2025-08-05). JAX warns and forwards the value to partition_index.

Source

Thrown at jax/_src/distributed.py:405

  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,
                          mtls_ca_file=mtls_ca_file,
                          mtls_peer_uri_prefix=mtls_peer_uri_prefix,
                          verify_secure_credentials=verify_secure_credentials)


def is_initialized() -> bool:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Rename the kwarg: jax.distributed.initialize(partition_index=...).
Update any wrapper utilities that forward slice_index via **kwargs.
Grep the codebase for 'slice_index' to catch all call sites before the deprecation is accelerated.

Example fix

# before
jax.distributed.initialize(coordinator_address=..., num_processes=8, process_id=rank, slice_index=idx)
# after
jax.distributed.initialize(coordinator_address=..., num_processes=8, process_id=rank, partition_index=idx)

Defensive patterns

Strategy: validation

Validate before calling

import inspect, jax
params = inspect.signature(jax.distributed.initialize).parameters
assert 'partition_index' in params and 'slice_index' in params  # both accepted for now

Try / catch

with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter('always', DeprecationWarning)
    jax.distributed.initialize(coordinator_address=..., num_processes=n, process_id=i, slice_index=i)
# rename kwarg once warning observed

Prevention

Use only documented kwargs; avoid **kwargs forwarding into initialize.
Grep codebase for 'slice_index' when upgrading JAX past 2025-08.

When it happens

Trigger: Explicitly calling jax.distributed.initialize(slice_index=...) in code; the value is used as partition_index after the warning.

Common situations: Existing multi-host training scripts written against older JAX APIs; libraries/tutorials that demonstrate initialize with slice_index.

Related errors

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