Google Colab无法切换Python版本，旧jax/numpyro无法正常使用

核心问题 #

你切换了系统默认Python到3.9，但调用pip时还是指向Colab预装Python3.10的pip，导致安装的包不在3.9的环境目录里，导入时自然还是调用3.10下的版本。

解决步骤 #

• 先安装Python3.9对应的pip工具：

  !sudo apt-get install python3.9-distutils
  !wget https://bootstrap.pypa.io/get-pip.py
  !python3.9 get-pip.py
  

• 用Python3.9专属的pip安装指定版本的jax和numpyro：

  !python3.9 -m pip install jax==0.2.17 numpyro==0.7.1
  

• 验证安装：重启Colab运行时（菜单栏→Runtime→Restart runtime），然后执行以下代码确认版本：

  import jax
  import numpyro
  print(jax.__version__)  # 应输出0.2.17
  print(numpyro.__version__)  # 应输出0.7.1
  

关键提醒 #

• 绝对不要直接用pip命令，Colab默认的pip绑定的是Python3.10，必须用python3.9 -m pip明确指定用3.9的pip安装包。
• 安装完成后必须重启运行时，否则Colab不会加载新的Python环境和对应包。

内容的提问来源于stack exchange，提问作者Cristian Calin