调用jax.random.PGRNKey时Jupyter内核崩溃求助（M2 Mac环境）

在搭载M2芯片的MacBook的Jupyter Lab环境中，调用jax.random.PRNGKey时出现内核崩溃。当前环境信息：

• Python 3.10.10
• jax 0.4.9
• jaxlib 0.4.9

已尝试升级相关依赖包，但问题仍未解决。

• 安装Apple Silicon适配的jaxlib版本：M2芯片属于arm64架构，需确保jaxlib为对应架构版本。卸载现有包后重新安装：

  pip uninstall -y jax jaxlib
  # 仅CPU版本
  pip install "jax[cpu]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
  # 或支持Metal GPU加速的版本
  pip install jax-metal jax
  

• 调整Python版本：Python 3.10.10与jax 0.4.9在Apple Silicon上可能存在兼容性冲突，尝试降级到Python 3.9.x系列，再重新安装jax依赖。
• 验证Jupyter内核环境：确认Jupyter Lab使用的内核对应安装了jax的Python环境，可通过jupyter kernelspec list查看内核路径，确保与目标虚拟环境一致。
• 测试终端运行情况：先在终端的Python环境中执行import jax; jax.random.PRNGKey(42)，若同样崩溃则问题出在jax与系统的兼容性，而非Jupyter；若终端运行正常，再排查Jupyter的配置或插件冲突。
• 调整Jupyter配置：修改Jupyter配置文件（~/.jupyter/jupyter_notebook_config.py），添加以下参数避免IO或线程问题：

  c.NotebookApp.max_buffer_size = 1024*1024*1024
  c.NotebookApp.iopub_data_rate_limit = 10000000
  

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