Kohya_ss native trainer colab broken yesterday. Captured before/ after install log. Jax, Key array.

Previously the install was broken by jax key array but could be fixed by:
!pip install "jax[cuda12_pip]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Any ideas whats broken? Is it because colab updated the default package versions?

It will kill the colab when you reach start training.
Reproduce:
https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-trainer.ipynb

**Working output:**
Stream
Looking in links: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
Collecting jax==0.4.23 (from jax[cuda12_pip]==0.4.23)
Downloading jax-0.4.23-py3-none-any.whl.metadata (24 kB)
Requirement already satisfied: ml-dtypes>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from jax==0.4.23->jax[cuda12_pip]==0.4.23) (0.4.1)
Requirement already satisfied: numpy>=1.22 in /usr/local/lib/python3.10/dist-packages (from jax==0.4.23->jax[cuda12_pip]==0.4.23) (1.26.4)
Requirement already satisfied: opt-einsum in /usr/local/lib/python3.10/dist-packages (from jax==0.4.23->jax[cuda12_pip]==0.4.23) (3.3.0)
Requirement already satisfied: scipy>=1.9 in /usr/local/lib/python3.10/dist-packages (from jax==0.4.23->jax[cuda12_pip]==0.4.23) (1.13.1)
Collecting jaxlib==0.4.23+cuda12.cudnn89 (from jax[cuda12_pip]==0.4.23)
Downloading https://storage.googleapis.com/jax-releases/cuda12/jaxlib-0.4.23%2Bcuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl (131.8 MB)
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 131.8/131.8 MB 8.8 MB/s eta 0:00:00
[?25hCollecting nvidia-cublas-cu12>=12.2.5.6 (from jax[cuda12_pip]==0.4.23)
Downloading nvidia_cublas_cu12-12.6.1.4-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)
Collecting nvidia-cuda-cupti-cu12>=12.2.142 (from jax[cuda12_pip]==0.4.23)
Downloading nvidia_cuda_cupti_cu12-12.6.68-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)
Collecting nvidia-cuda-nvcc-cu12>=12.2.140 (from jax[cuda12_pip]==0.4.23)
Downloading nvidia_cuda_nvcc_cu12-12.6.68-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)
Collecting nvidia-cuda-runtime-cu12>=12.2.140 (from jax[cuda12_pip]==0.4.23)
Downloading nvidia_cuda_runtime_cu12-12.6.68-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)
Requirement already satisfied: nvidia-cudnn-cu12>=8.9 in /usr/local/lib/python3.10/dist-packages (from jax[cuda12_pip]==0.4.23) (8.9.2.26)
Collecting nvidia-cufft-cu12>=11.0.8.103 (from jax[cuda12_pip]==0.4.23)
Downloading nvidia_cufft_cu12-11.2.6.59-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)
Collecting nvidia-cusolver-cu12>=11.5.2 (from jax[cuda12_pip]==0.4.23)
Downloading nvidia_cusolver_cu12-11.6.4.69-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)
Collecting nvidia-cusparse-cu12>=12.1.2.141 (from jax[cuda12_pip]==0.4.23)
Downloading nvidia_cusparse_cu12-12.5.3.3-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)
Collecting nvidia-nccl-cu12>=2.18.3 (from jax[cuda12_pip]==0.4.23)
Using cached nvidia_nccl_cu12-2.23.4-py3-none-manylinux2014_x86_64.whl.metadata (1.8 kB)
Requirement already satisfied: nvidia-nvjitlink-cu12>=12.2 in /usr/local/lib/python3.10/dist-packages (from jax[cuda12_pip]==0.4.23) (12.6.68)
Downloading jax-0.4.23-py3-none-any.whl (1.7 MB)
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.7/1.7 MB 33.9 MB/s eta 0:00:00
[?25hDownloading nvidia_cublas_cu12-12.6.1.4-py3-none-manylinux2014_x86_64.whl (378.9 MB)
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 378.9/378.9 MB 2.3 MB/s eta 0:00:00
[?25hDownloading nvidia_cuda_cupti_cu12-12.6.68-py3-none-manylinux2014_x86_64.whl (8.9 MB)
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 8.9/8.9 MB 113.5 MB/s eta 0:00:00
[?25hDownloading nvidia_cuda_nvcc_cu12-12.6.68-py3-none-manylinux2014_x86_64.whl (21.1 MB)
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 21.1/21.1 MB 86.4 MB/s eta 0:00:00
[?25hDownloading nvidia_cuda_runtime_cu12-12.6.68-py3-none-manylinux2014_x86_64.whl (897 kB)
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 897.7/897.7 kB 50.5 MB/s eta 0:00:00
[?25hDownloading nvidia_cufft_cu12-11.2.6.59-py3-none-manylinux2014_x86_64.whl (192.5 MB)
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 192.5/192.5 MB 5.8 MB/s eta 0:00:00
[?25hDownloading nvidia_cusolver_cu12-11.6.4.69-py3-none-manylinux2014_x86_64.whl (130.5 MB)
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 130.5/130.5 MB 8.6 MB/s eta 0:00:00
[?25hDownloading nvidia_cusparse_cu12-12.5.3.3-py3-none-manylinux2014_x86_64.whl (216.5 MB)
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 216.5/216.5 MB 3.0 MB/s eta 0:00:00
[?25hUsing cached nvidia_nccl_cu12-2.23.4-py3-none-manylinux2014_x86_64.whl (199.0 MB)
Installing collected packages: nvidia-nccl-cu12, nvidia-cusparse-cu12, nvidia-cufft-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvcc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, nvidia-cusolver-cu12, jaxlib, jax
Attempting uninstall: nvidia-nccl-cu12
Found existing installation: nvidia-nccl-cu12 2.18.1
Uninstalling nvidia-nccl-cu12-2.18.1:
Successfully uninstalled nvidia-nccl-cu12-2.18.1
Attempting uninstall: nvidia-cusparse-cu12
Found existing installation: nvidia-cusparse-cu12 12.1.0.106
Uninstalling nvidia-cusparse-cu12-12.1.0.106:
Successfully uninstalled nvidia-cusparse-cu12-12.1.0.106
Attempting uninstall: nvidia-cufft-cu12
Found existing installation: nvidia-cufft-cu12 11.0.2.54
Uninstalling nvidia-cufft-cu12-11.0.2.54:
Successfully uninstalled nvidia-cufft-cu12-11.0.2.54
Attempting uninstall: nvidia-cuda-runtime-cu12
Found existing installation: nvidia-cuda-runtime-cu12 12.1.105
Uninstalling nvidia-cuda-runtime-cu12-12.1.105:
Successfully uninstalled nvidia-cuda-runtime-cu12-12.1.105
Attempting uninstall: nvidia-cuda-cupti-cu12
Found existing installation: nvidia-cuda-cupti-cu12 12.1.105
Uninstalling nvidia-cuda-cupti-cu12-12.1.105:
Successfully uninstalled nvidia-cuda-cupti-cu12-12.1.105
Attempting uninstall: nvidia-cublas-cu12
Found existing installation: nvidia-cublas-cu12 12.1.3.1
Uninstalling nvidia-cublas-cu12-12.1.3.1:
Successfully uninstalled nvidia-cublas-cu12-12.1.3.1
Attempting uninstall: nvidia-cusolver-cu12
Found existing installation: nvidia-cusolver-cu12 11.4.5.107
Uninstalling nvidia-cusolver-cu12-11.4.5.107:
Successfully uninstalled nvidia-cusolver-cu12-11.4.5.107
Attempting uninstall: jaxlib
Found existing installation: jaxlib 0.4.26+cuda12.cudnn89
Uninstalling jaxlib-0.4.26+cuda12.cudnn89:
Successfully uninstalled jaxlib-0.4.26+cuda12.cudnn89
Attempting uninstall: jax
Found existing installation: jax 0.4.26
Uninstalling jax-0.4.26:
Successfully uninstalled jax-0.4.26
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
orbax-checkpoint 0.6.4 requires jax>=0.4.26, but you have jax 0.4.23 which is incompatible.
torch 2.1.0 requires nvidia-cublas-cu12==12.1.3.1; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cublas-cu12 12.6.1.4 which is incompatible.
torch 2.1.0 requires nvidia-cuda-cupti-cu12==12.1.105; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cuda-cupti-cu12 12.6.68 which is incompatible.
torch 2.1.0 requires nvidia-cuda-runtime-cu12==12.1.105; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cuda-runtime-cu12 12.6.68 which is incompatible.
torch 2.1.0 requires nvidia-cufft-cu12==11.0.2.54; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cufft-cu12 11.2.6.59 which is incompatible.
torch 2.1.0 requires nvidia-cusolver-cu12==11.4.5.107; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cusolver-cu12 11.6.4.69 which is incompatible.
torch 2.1.0 requires nvidia-cusparse-cu12==12.1.0.106; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cusparse-cu12 12.5.3.3 which is incompatible.
torch 2.1.0 requires nvidia-nccl-cu12==2.18.1; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-nccl-cu12 2.23.4 which is incompatible.
torchaudio 2.4.1+cu121 requires torch==2.4.1, but you have torch 2.1.0 which is incompatible.
Successfully installed jax-0.4.23 jaxlib-0.4.23+cuda12.cudnn89 nvidia-cublas-cu12-12.6.1.4 nvidia-cuda-cupti-cu12-12.6.68 nvidia-cuda-nvcc-cu12-12.6.68 nvidia-cuda-runtime-cu12-12.6.68 nvidia-cufft-cu12-11.2.6.59 nvidia-cusolver-cu12-11.6.4.69 nvidia-cusparse-cu12-12.5.3.3 nvidia-nccl-cu12-2.23.4

**Broken output**
Stream
Looking in links: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
Collecting jax==0.4.23 (from jax[cuda12_pip]==0.4.23)
Downloading jax-0.4.23-py3-none-any.whl.metadata (24 kB)
Requirement already satisfied: ml-dtypes>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from jax==0.4.23->jax[cuda12_pip]==0.4.23) (0.4.1)
Requirement already satisfied: numpy>=1.22 in /usr/local/lib/python3.10/dist-packages (from jax==0.4.23->jax[cuda12_pip]==0.4.23) (1.26.4)
Requirement already satisfied: opt-einsum in /usr/local/lib/python3.10/dist-packages (from jax==0.4.23->jax[cuda12_pip]==0.4.23) (3.3.0)
Requirement already satisfied: scipy>=1.9 in /usr/local/lib/python3.10/dist-packages (from jax==0.4.23->jax[cuda12_pip]==0.4.23) (1.13.1)
Collecting jaxlib==0.4.23+cuda12.cudnn89 (from jax[cuda12_pip]==0.4.23)
Downloading https://storage.googleapis.com/jax-releases/cuda12/jaxlib-0.4.23%2Bcuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl (131.8 MB)
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 131.8/131.8 MB 7.9 MB/s eta 0:00:00
[?25hCollecting nvidia-cublas-cu12>=12.2.5.6 (from jax[cuda12_pip]==0.4.23)
Using cached nvidia_cublas_cu12-12.6.1.4-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)
Collecting nvidia-cuda-cupti-cu12>=12.2.142 (from jax[cuda12_pip]==0.4.23)
Using cached nvidia_cuda_cupti_cu12-12.6.68-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)
Requirement already satisfied: nvidia-cuda-nvcc-cu12>=12.2.140 in /usr/local/lib/python3.10/dist-packages (from jax[cuda12_pip]==0.4.23) (12.6.68)
Collecting nvidia-cuda-runtime-cu12>=12.2.140 (from jax[cuda12_pip]==0.4.23)
Using cached nvidia_cuda_runtime_cu12-12.6.68-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)
Requirement already satisfied: nvidia-cudnn-cu12>=8.9 in /usr/local/lib/python3.10/dist-packages (from jax[cuda12_pip]==0.4.23) (8.9.2.26)
Collecting nvidia-cufft-cu12>=11.0.8.103 (from jax[cuda12_pip]==0.4.23)
Using cached nvidia_cufft_cu12-11.2.6.59-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)
Collecting nvidia-cusolver-cu12>=11.5.2 (from jax[cuda12_pip]==0.4.23)
Using cached nvidia_cusolver_cu12-11.6.4.69-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)
Collecting nvidia-cusparse-cu12>=12.1.2.141 (from jax[cuda12_pip]==0.4.23)
Using cached nvidia_cusparse_cu12-12.5.3.3-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)
Collecting nvidia-nccl-cu12>=2.18.3 (from jax[cuda12_pip]==0.4.23)
Using cached nvidia_nccl_cu12-2.23.4-py3-none-manylinux2014_x86_64.whl.metadata (1.8 kB)
Requirement already satisfied: nvidia-nvjitlink-cu12>=12.2 in /usr/local/lib/python3.10/dist-packages (from jax[cuda12_pip]==0.4.23) (12.6.68)
Downloading jax-0.4.23-py3-none-any.whl (1.7 MB)
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.7/1.7 MB 29.0 MB/s eta 0:00:00
[?25hUsing cached nvidia_cublas_cu12-12.6.1.4-py3-none-manylinux2014_x86_64.whl (378.9 MB)
Using cached nvidia_cuda_cupti_cu12-12.6.68-py3-none-manylinux2014_x86_64.whl (8.9 MB)
Using cached nvidia_cuda_runtime_cu12-12.6.68-py3-none-manylinux2014_x86_64.whl (897 kB)
Using cached nvidia_cufft_cu12-11.2.6.59-py3-none-manylinux2014_x86_64.whl (192.5 MB)
Using cached nvidia_cusolver_cu12-11.6.4.69-py3-none-manylinux2014_x86_64.whl (130.5 MB)
Using cached nvidia_cusparse_cu12-12.5.3.3-py3-none-manylinux2014_x86_64.whl (216.5 MB)
Using cached nvidia_nccl_cu12-2.23.4-py3-none-manylinux2014_x86_64.whl (199.0 MB)
Installing collected packages: nvidia-nccl-cu12, nvidia-cusparse-cu12, nvidia-cufft-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, nvidia-cusolver-cu12, jaxlib, jax
Attempting uninstall: nvidia-nccl-cu12
Found existing installation: nvidia-nccl-cu12 2.18.1
Uninstalling nvidia-nccl-cu12-2.18.1:
Successfully uninstalled nvidia-nccl-cu12-2.18.1
Attempting uninstall: nvidia-cusparse-cu12
Found existing installation: nvidia-cusparse-cu12 12.1.0.106
Uninstalling nvidia-cusparse-cu12-12.1.0.106:
Successfully uninstalled nvidia-cusparse-cu12-12.1.0.106
Attempting uninstall: nvidia-cufft-cu12
Found existing installation: nvidia-cufft-cu12 11.0.2.54
Uninstalling nvidia-cufft-cu12-11.0.2.54:
Successfully uninstalled nvidia-cufft-cu12-11.0.2.54
Attempting uninstall: nvidia-cuda-runtime-cu12
Found existing installation: nvidia-cuda-runtime-cu12 12.1.105
Uninstalling nvidia-cuda-runtime-cu12-12.1.105:
Successfully uninstalled nvidia-cuda-runtime-cu12-12.1.105
Attempting uninstall: nvidia-cuda-cupti-cu12
Found existing installation: nvidia-cuda-cupti-cu12 12.1.105
Uninstalling nvidia-cuda-cupti-cu12-12.1.105:
Successfully uninstalled nvidia-cuda-cupti-cu12-12.1.105
Attempting uninstall: nvidia-cublas-cu12
Found existing installation: nvidia-cublas-cu12 12.1.3.1
Uninstalling nvidia-cublas-cu12-12.1.3.1:
Successfully uninstalled nvidia-cublas-cu12-12.1.3.1
Attempting uninstall: nvidia-cusolver-cu12
Found existing installation: nvidia-cusolver-cu12 11.4.5.107
Uninstalling nvidia-cusolver-cu12-11.4.5.107:
Successfully uninstalled nvidia-cusolver-cu12-11.4.5.107
Attempting uninstall: jaxlib
Found existing installation: jaxlib 0.4.33
Uninstalling jaxlib-0.4.33:
Successfully uninstalled jaxlib-0.4.33
Attempting uninstall: jax
Found existing installation: jax 0.4.33
Uninstalling jax-0.4.33:
Successfully uninstalled jax-0.4.33
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
flax 0.8.5 requires jax>=0.4.27, but you have jax 0.4.23 which is incompatible.
optax 0.2.3 requires jax>=0.4.27, but you have jax 0.4.23 which is incompatible.
optax 0.2.3 requires jaxlib>=0.4.27, but you have jaxlib 0.4.23+cuda12.cudnn89 which is incompatible.
orbax-checkpoint 0.6.4 requires jax>=0.4.26, but you have jax 0.4.23 which is incompatible.
torch 2.1.0 requires nvidia-cublas-cu12==12.1.3.1; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cublas-cu12 12.6.1.4 which is incompatible.
torch 2.1.0 requires nvidia-cuda-cupti-cu12==12.1.105; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cuda-cupti-cu12 12.6.68 which is incompatible.
torch 2.1.0 requires nvidia-cuda-runtime-cu12==12.1.105; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cuda-runtime-cu12 12.6.68 which is incompatible.
torch 2.1.0 requires nvidia-cufft-cu12==11.0.2.54; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cufft-cu12 11.2.6.59 which is incompatible.
torch 2.1.0 requires nvidia-cusolver-cu12==11.4.5.107; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cusolver-cu12 11.6.4.69 which is incompatible.
torch 2.1.0 requires nvidia-cusparse-cu12==12.1.0.106; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cusparse-cu12 12.5.3.3 which is incompatible.
torch 2.1.0 requires nvidia-nccl-cu12==2.18.1; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-nccl-cu12 2.23.4 which is incompatible.
torchaudio 2.4.1+cu121 requires torch==2.4.1, but you have torch 2.1.0 which is incompatible.
Successfully installed jax-0.4.23 jaxlib-0.4.23+cuda12.cudnn89 nvidia-cublas-cu12-12.6.1.4 nvidia-cuda-cupti-cu12-12.6.68 nvidia-cuda-runtime-cu12-12.6.68 nvidia-cufft-cu12-11.2.6.59 nvidia-cusolver-cu12-11.6.4.69 nvidia-cusparse-cu12-12.5.3.3 nvidia-nccl-cu12-2.23.4

Open kohya-trainer.ipynb from the linked repository and reproduce the installation in a fresh Colab runtime. Compare the working and broken install logs, then run through the start-training step to identify the JAX or CUDA package conflict. Done means the notebook installs consistently and reaches training without killing the Colab runtime.

Written by the indexing model from the issue text.