使用协作分支上的算法
git clone https://github.com/linjing-lab/jax.git
cd jax
git checkout linjing-lab
cd examples
在命令行预览方法
牛顿方法:
cat newton_method.py
拟牛顿法:
cat bfgs_method.py
在命令行运行程序
python newton_method.py
python bfgs_method.py
建议
安装jax,一定要顺带把jaxlib的多后端版本都配置成功了。如果需要多后端加速(GPU),使用@jax.jit
包装整个方法函数或者定位到某行程序。