JAX学习笔记

• 1.安装环境
• 2.基础知识

2.我这里用python3.7的虚拟环境

conda create -n jax python=3.7
conda activate jax

2.下载安装jax和jaxlib等相关依赖包，然后进入虚拟环境安装

cd C:UsersdzDownloadsjax-0.2.9_and_jaxlib-0.1.61-cp37-win_amd64
pip install jaxlib-0.1.61-cp37-none-win_amd64.whl
pip install jax==0.2.9
pip install matplotlib

3.写个py文件测试环境是否安装成功

import jax.numpy as jnp
import matplotlib.pyplot as plt
x_jnp=jnp.linspace(0,10,1000)
y_jnp=jnp.sin(x_jnp)*jnp.cos(x_jnp)
print(x_jnp,y_jnp)
plt.plot(x_jnp,y_jnp)
plt.show()

如下，环境安装成功