JAX is in GPU Image, but not in CPU or TPU Images

https://github.com/Kaggle/docker-python/blob/master/gpu.Dockerfile#L83 line 83 installs JAX, and I see it's usable in a GPU enabled kaggle notebook.

https://github.com/Kaggle/docker-python/blob/master/Dockerfile JAX is not installed here, I'm not able to use it on either a CPU or TPU notebook.

Can we get JAX on the CPU and TPU instances too?