基于JAX的二阶优化方法的实践

使用协作分支上的算法

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包装整个方法函数或者定位到某行程序。