Feature request: Slurm env detection

Hello, I was wondering if it’s possible to detect slurm environments automatically like `jax.distributed.initialize()`. Example:

The relevant jax code is here: https://github.com/google/jax/blob/main/jax/_src/clusters/slurm_cluster.py

```
#!/bin/bash
#SBATCH --job-name=trl
#SBATCH --partition=production-cluster
#SBATCH --gpus-per-task=8
#SBATCH --cpus-per-gpu=10
#SBATCH --ntasks=2 # 1 nodes
#SBATCH --output=slurm/logs/%x_%j.out
srun python jax_test.py
```
```python
import jax
import os
jax.distributed.initialize()
print(os.environ.get("CUDA_VISIBLE_DEVICES"))
print("jax.device_count()", jax.device_count())
print("jax.local_device_count()", jax.local_device_count())
print("jax.devices()", jax.devices())
```
yields
```
0,1,2,3,4,5,6,7
jax.device_count() 16
jax.local_device_count() 8
jax.devices() [StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=1, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=2, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=3, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=4, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=5, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=6, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=7, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=8, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=9, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=10, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=11, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=12, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=13, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=14, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=15, process_index=1, slice_index=1)]
0,1,2,3,4,5,6,7
jax.device_count() 16
jax.local_device_count() 8
jax.devices() [StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=1, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=2, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=3, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=4, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=5, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=6, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=7, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=8, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=9, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=10, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=11, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=12, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=13, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=14, process_index=1, slice_index=1), StreamExecutorGpuDevice(id=15, process_index=1, slice_index=1)]
```