Jax, Numpy & Scipy

Google开源的Jax据说很厉害,主要是用来解决numpy在GPU上支持度不太友好的开源包(PyTorch,Cupy,Numba等均同此),今天小试一下,确实有优秀的地方

测试(仅限很少部分的numpy api,文档中有Scipy的api)

  • Version
    Jax, Numpy & Scipy_第1张图片

  • 随机数 & jax替换numpy的点乘
    Jax, Numpy & Scipy_第2张图片

  • jit
    Jax, Numpy & Scipy_第3张图片
    在这里插入图片描述

  • 函数
    Jax, Numpy & Scipy_第4张图片

  • 梯度计算
    Jax, Numpy & Scipy_第5张图片

Reference

  • 官方文档
  • Github

你可能感兴趣的:(#,python)