{"record":{"id":"f12d4cf2c0c4fde8","repo":"jax-ml/jax","slug":"jax-slice-index-has-been-deprecated-please-use-ja","errorCode":null,"errorMessage":"JAX_SLICE_INDEX has been deprecated. Please use JAX_PARTITION_INDEX instead.","messagePattern":"JAX_SLICE_INDEX has been deprecated\\. Please use JAX_PARTITION_INDEX instead\\.","errorType":"console","errorClass":null,"httpStatus":null,"severity":"warning","filePath":"jax/_src/distributed.py","lineNumber":229,"sourceCode":"        process_id,\n        init_timeout=initialization_timeout,\n        use_compression=True,\n        heartbeat_timeout=heartbeat_timeout_seconds,\n        **mtls_kwargs,\n    )\n    logger.info('Connecting to JAX distributed service on %s', coordinator_address)\n    self.client.connect()\n\n    self.initialize_preemption_sync_manager()\n\n    if partition_index is None:\n      jax_partition_index = os.environ.get('JAX_PARTITION_INDEX')\n      jax_slice_index = os.environ.get('JAX_SLICE_INDEX')\n      if jax_partition_index is not None:\n        partition_index = int(jax_partition_index)\n      elif jax_slice_index is not None:\n        # Deprecation added 2025-08-05. Should be removed after 3 months.\n        warnings.warn(\n            'JAX_SLICE_INDEX has been deprecated. Please use'\n            ' JAX_PARTITION_INDEX instead.',\n            DeprecationWarning,\n        )\n        partition_index = int(jax_slice_index)\n    self.partition_index = partition_index\n\n  def shutdown(self):\n    if self.preemption_sync_manager:\n      # It's important to shut down the preemption sync manager before the\n      # client because the preemption sync manager depends on the client.\n      self.preemption_sync_manager.shutdown()\n      self.preemption_sync_manager = None\n    if self.client:\n      self.client.shutdown()\n      self.client = None\n    if self.service:\n      self.service.shutdown()","sourceCodeStart":211,"sourceCodeEnd":247,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/distributed.py#L211-L247","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Rename the env var in launch scripts: export JAX_PARTITION_INDEX instead of JAX_SLICE_INDEX.","If both are set, JAX_PARTITION_INDEX wins — ensure they don't disagree.","Purge JAX_SLICE_INDEX from k8s manifests, Dockerfiles, and slurm scripts to avoid breakage when the deprecation is removed."],"exampleFix":"# before\nexport JAX_SLICE_INDEX=0\n# after\nexport JAX_PARTITION_INDEX=0","handlingStrategy":"validation","validationCode":"import os\nassert os.environ.get('JAX_SLICE_INDEX') is None, 'rename to JAX_PARTITION_INDEX'","typeGuard":null,"tryCatchPattern":"with warnings.catch_warnings(record=True) as w:\n    warnings.simplefilter('always', DeprecationWarning)\n    jax.distributed.initialize(...)\nold_var_used = any('JAX_SLICE_INDEX' in str(i.message) for i in w)","preventionTips":["Centralize distributed env vars in one launch script; audit on JAX upgrades.","Grep manifests for JAX_SLICE_INDEX and rename before the 3-month removal window closes."],"tags":["jax","distributed","environment-variable","deprecation"],"backgroundTag":"renamed-env-var","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}