jax-ml/jax · warning

JAX_SLICE_INDEX has been deprecated. Please use JAX_PARTITIO

Error message

JAX_SLICE_INDEX has been deprecated. Please use JAX_PARTITION_INDEX instead.

What it means

jax.distributed.initialize() detected the JAX_SLICE_INDEX environment variable, which was renamed to JAX_PARTITION_INDEX (deprecation added 2025-08-05, removal after ~3 months). It warns and uses the old variable's value as the partition index.

Source

Thrown at jax/_src/distributed.py:229

        process_id,
        init_timeout=initialization_timeout,
        use_compression=True,
        heartbeat_timeout=heartbeat_timeout_seconds,
        **mtls_kwargs,
    )
    logger.info('Connecting to JAX distributed service on %s', coordinator_address)
    self.client.connect()

    self.initialize_preemption_sync_manager()

    if partition_index is None:
      jax_partition_index = os.environ.get('JAX_PARTITION_INDEX')
      jax_slice_index = os.environ.get('JAX_SLICE_INDEX')
      if jax_partition_index is not None:
        partition_index = int(jax_partition_index)
      elif jax_slice_index is not None:
        # Deprecation added 2025-08-05. Should be removed after 3 months.
        warnings.warn(
            'JAX_SLICE_INDEX has been deprecated. Please use'
            ' JAX_PARTITION_INDEX instead.',
            DeprecationWarning,
        )
        partition_index = int(jax_slice_index)
    self.partition_index = partition_index

  def shutdown(self):
    if self.preemption_sync_manager:
      # It's important to shut down the preemption sync manager before the
      # client because the preemption sync manager depends on the client.
      self.preemption_sync_manager.shutdown()
      self.preemption_sync_manager = None
    if self.client:
      self.client.shutdown()
      self.client = None
    if self.service:
      self.service.shutdown()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Rename the env var in launch scripts: export JAX_PARTITION_INDEX instead of JAX_SLICE_INDEX.
2. If both are set, JAX_PARTITION_INDEX wins — ensure they don't disagree.
3. Purge JAX_SLICE_INDEX from k8s manifests, Dockerfiles, and slurm scripts to avoid breakage when the deprecation is removed.

Example fix

# before
export JAX_SLICE_INDEX=0
# after
export JAX_PARTITION_INDEX=0

import os
assert os.environ.get('JAX_SLICE_INDEX') is None, 'rename to JAX_PARTITION_INDEX'

with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter('always', DeprecationWarning)
    jax.distributed.initialize(...)
old_var_used = any('JAX_SLICE_INDEX' in str(i.message) for i in w)

When it happens

Trigger: Running multi-process JAX with JAX_SLICE_INDEX set in the environment when JAX_PARTITION_INDEX is unset; common in orchestrator configs, k8s manifests, or container images built for older JAX.

Common situations: Upgrading a distributed training pipeline to newer JAX while launch scripts still export JAX_SLICE_INDEX; mixed-version clusters where the scheduler injects the old var.

Related errors

• `slice_index` has been deprecated. Please use `partition_ind
• The return value of the policies should be a boolean. Got: {
• Mapped away dimension of inputs passed to vmap should be sha
• invalid truth value {val!r} for environment {varname!r}
• Invalid value "{default}" for JAX flag {name}

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