谷歌开源计算框架JAX:比Numpy快30倍,还可在TPU上运行!
新智元报道
新智元报道
来源:外媒
编辑:keyu
【新智元导读】大家有了解过JAX吗?JAX是一种可在CPU、GPU和TPU上运行的“Numpy”,专门针对机器学习研究,并提供高性能自微分计算能力,速度要比纯用Numpy快几十倍!
import numpy as np # 使用标准numpy,运算将在CPU上执行。
x = np.random.random([5000, 5000]).astype(np.float32)
%timeit np.matmul(x, x)
import jax.numpy as np # 使用"JAX版"的numpy from jax import random # 注意JAX下随机数API有所不同 x = random.uniform(random.PRNGKey(0), [5000, 5000]) %timeit np.matmul(x, x)
快速入门链接:https://jax.readthedocs.io/en/latest/notebooks/quickstart.html
Github项目地址:https://github.com/google/jax
https://roberttlange.github.io/posts/2021/02/cma-es-jax/
https://roberttlange.github.io/posts/2020/03/blog-post-10/
https://jax.readthedocs.io/en/latest/notebooks/quickstart.html
https://www.zhihu.com/question/306496943/answer/1041519580
推荐阅读:
关注公众号:拾黑(shiheibook)了解更多
[广告]赞助链接:
四季很好,只要有你,文娱排行榜:https://www.yaopaiming.com/
让资讯触达的更精准有趣:https://www.0xu.cn/