TensorFlow有了替代品，竟然还是谷歌自己做出来的？
tensorflow有了替代品，竟然还是谷歌自己做出来的？这其实是tensorflow的一个简化库，名为jax，可以支持部分tensorflow的功能，但是比tensorflow更加简洁易用。
什么？tensorflow 有了替代品？什么？竟然还是谷歌自己做出来的？先别慌，从各种意义上来说，这个所谓的 “替代品” 其实是 tensorflow 的一个简化库，名为jax，结合 autograd 和 xla，可以支持部分 tensorflow 的功能，但是比 tensorflow 更加简洁易用。
虽然还不至于替代 tensorflow，但已经有 reddit 网友对 jax 寄予厚望，并表示“早就期待能有一个可以直接调用 numpy api 接口的库了！”，“希望它可以取代 tensorflow！”。
jax 结合了 autograd 和 xla，是专为高性能机器学习研究打造的产品。
有了新版本的autograd，jax 能够自动对 python 和 numpy 的自带函数求导，支持循环、分支、递归、闭包函数求导，而且可以求三阶导数。它支持自动模式反向求导（也就是反向传播）和正向求导，且二者可以任意组合成任何顺序。
jax 的创新之处在于，它基于xla在 gpu 和 tpu 上编译和运行 numpy 程序。默认情况下，编译是在底层进行的，库调用能够及时编译和执行。但是 jax 还允许使用单一函数 api jit将自己的 python 函数及时编译成经过 xla 优化的内核。编译和自动求导可以任意组合，因此可以在不脱离 python 环境的情况下实现复杂算法并获得最优性能。
jax 最初由 matt johnson、roy frostig、dougal maclaurin 和 chris leary 发起，他们均任职于谷歌大脑团队。在 github 的说明文档中，作者明确表示：jax 目前还只是一个研究项目，不是谷歌的官方产品，因此可能会有一些 bug。从作者的 github 简介来看，这应该是谷歌大脑正在尝试的新项目，在同一个 github 目录下的开源项目还包括 8 月份在业内引起热议的强化学习框架 dopamine。
以下是 jax 的简单使用示例。
github 项目传送门：https://github.com/google/jax
有关具体的安装和简单的入门指导大家可以在 github 中自行查看，在此不做过多赘述。
jax 库的实现原理
机器学习中的编程是关于函数的表达和转换。转换包括自动微分、加速器编译和自动批处理。像 python 这样的高级语言非常适合表达函数，但是通常使用者只能应用它们。我们无法访问它们的内部结构，因此无法执行转换。
jax 可以用于专门化高级python+numpy函数，并将其转换为可转换的表示形式，然后再提升为 python 函数。
jax 通过跟踪专门处理 python 函数。跟踪一个函数意味着：监视应用于其输入，以产生其输出的所有基本操作，并在有向无环图 (dag) 中记录这些操作及其之间的数据流。为了执行跟踪，jax 包装了基本的操作，就像基本的数字内核一样，这样一来，当调用它们时，它们就会将自己添加到执行的操作列表以及输入和输出中。为了跟踪这些原语之间的数据流，跟踪的值被包装在 tracer 类的实例中。
当 python 函数被提供给 grad 或 jit 时，它被包装起来以便跟踪并返回。当调用包装的函数时，我们将提供的具体参数抽象到 abstractvalue 类的实例中，将它们框起来用于跟踪跟踪器类的实例，并对它们调用函数。
抽象参数表示一组可能的值，而不是特定的值：例如，jit 将 ndarray 参数抽象为抽象值，这些值表示具有相同形状和数据类型的所有 ndarray。相反，grad 抽象 ndarray 参数来表示底层值的无穷小邻域。通过在这些抽象值上跟踪 python 函数，我们确保它足够专门化，以便转换是可处理的，并且它仍然足够通用，以便转换后的结果是有用的，并且可能是可重用的。然后将这些转换后的函数提升回 python 可调用函数，这样就可以根据需要跟踪并再次转换它们。
jax 跟踪的基本函数大多与 xla hlo 1:1 对应，并在 lax.py 中定义。这种 1:1 的对应关系使得到 xla 的大多数转换基本上都很简单，并且确保我们只有一小组原语来覆盖其他转换，比如自动微分。 jax.numpy 层是用纯 python 编写的，它只是用 lax 函数 (以及我们已经编写的其他 numpy 函数) 表示 numpy 函数。这使得 jax.numpy 易于延展。
当你使用 jax.numpy 时，底层 lax 原语是在后台进行 jit 编译的，允许你在加速器上执行每个原语操作的同时编写不受限制的 python+ numpy 代码。
但是 jax 可以做更多的事情：你可以在越来越大的函数上使用jit来进行端到端编译和优化，而不仅仅是编译和调度到一组固定的单个原语。例如，可以编译整个网络，或者编译整个梯度计算和优化器更新步骤，而不仅仅是编译和调度卷积运算。
折衷之处是，jit 函数必须满足一些额外的专门化需求：因为我们希望编译专门针对形状和数据类型的跟踪，但不是专门针对具体值的跟踪，所以 jit 装饰器下的 python 代码必须适用于抽象值。如果我们尝试在一个抽象的 x 上求 x >0 的值，结果是一个抽象的值，表示集合 {true, false}，所以 python 分支就像 if x > 0 会引起报错。
有关使用 jit 的更多要求，请参见：https://github.com/google/jax#whats-supported
好消息是，jit 是可选的：jax 库在后台对单个操作和函数使用 jit，允许编写不受限制的 python+numpy，同时仍然使用硬件加速器。但是，当你希望最大化性能时，通常可以在自己的代码中使用 jit 编译和端到端优化更大的函数。
后续计划
目前项目小组还将对以下几项做更多尝试和更新：
完善说明文档
支持 cloud tpu
支持多 gpu 和多 tpu
支持完整的 numpy 功能和部分 scipy 功能
全面支持 vmap
加速
降低 xla 函数调度开销
线性代数例程（cpu 上的 mkl 和 gpu 上的 magma）
高效自动微分原语cond和while
有关 jax 库的介绍大致如此。
新大陆NLS-NVF260工业条码扫描器的优势
小米5G新机入网，支持SA/NSA双模5G
如何制作一个能玩的起重机
低噪声、高耐压、双路高音质音频运算放大器“NL8802”发布上市
亮钻由由强强合作，“AI称重”大有前途
TensorFlow有了替代品，竟然还是谷歌自己做出来的？
日韩企业垄断CMOS芯片，中国垄断摄像头模组市场
IPv4报文格式各字段的含义
微软展示多屏显示技术 或将被用于VR头显
FPGA的基本组成结构
基于MXene的全柔性自供电压力传感系统
中低阶智能型手机生物识别大战已开打，指纹识别或将被淘汰
判断两个字符串中的字母是否一致
如何正确选择TRINAMIC运动控制芯片
关于不锈钢板激光焊接变形的研究
vivo x9、oppo r9s、荣耀v9，买哪一款更好？
苹果中国台湾主要供应商9月营收387亿美元，同比下降16.4%
电动车电瓶修复技术——笔记本电脑电池的维护技巧（网友问答）
6 种技术趋势正在改变我们的出行方式
RGB7段显示器的制作
什么？tensorflow 有了替代品？什么？竟然还是谷歌自己做出来的？先别慌，从各种意义上来说，这个所谓的 “替代品” 其实是 tensorflow 的一个简化库，名为jax，结合 autograd 和 xla，可以支持部分 tensorflow 的功能，但是比 tensorflow 更加简洁易用。
虽然还不至于替代 tensorflow，但已经有 reddit 网友对 jax 寄予厚望，并表示“早就期待能有一个可以直接调用 numpy api 接口的库了！”，“希望它可以取代 tensorflow！”。
jax 结合了 autograd 和 xla，是专为高性能机器学习研究打造的产品。
有了新版本的autograd，jax 能够自动对 python 和 numpy 的自带函数求导，支持循环、分支、递归、闭包函数求导，而且可以求三阶导数。它支持自动模式反向求导（也就是反向传播）和正向求导，且二者可以任意组合成任何顺序。
jax 的创新之处在于，它基于xla在 gpu 和 tpu 上编译和运行 numpy 程序。默认情况下，编译是在底层进行的，库调用能够及时编译和执行。但是 jax 还允许使用单一函数 api jit将自己的 python 函数及时编译成经过 xla 优化的内核。编译和自动求导可以任意组合，因此可以在不脱离 python 环境的情况下实现复杂算法并获得最优性能。
jax 最初由 matt johnson、roy frostig、dougal maclaurin 和 chris leary 发起，他们均任职于谷歌大脑团队。在 github 的说明文档中，作者明确表示：jax 目前还只是一个研究项目，不是谷歌的官方产品，因此可能会有一些 bug。从作者的 github 简介来看，这应该是谷歌大脑正在尝试的新项目，在同一个 github 目录下的开源项目还包括 8 月份在业内引起热议的强化学习框架 dopamine。
以下是 jax 的简单使用示例。
github 项目传送门：https://github.com/google/jax
有关具体的安装和简单的入门指导大家可以在 github 中自行查看，在此不做过多赘述。
jax 库的实现原理
机器学习中的编程是关于函数的表达和转换。转换包括自动微分、加速器编译和自动批处理。像 python 这样的高级语言非常适合表达函数，但是通常使用者只能应用它们。我们无法访问它们的内部结构，因此无法执行转换。
jax 可以用于专门化高级python+numpy函数，并将其转换为可转换的表示形式，然后再提升为 python 函数。
jax 通过跟踪专门处理 python 函数。跟踪一个函数意味着：监视应用于其输入，以产生其输出的所有基本操作，并在有向无环图 (dag) 中记录这些操作及其之间的数据流。为了执行跟踪，jax 包装了基本的操作，就像基本的数字内核一样，这样一来，当调用它们时，它们就会将自己添加到执行的操作列表以及输入和输出中。为了跟踪这些原语之间的数据流，跟踪的值被包装在 tracer 类的实例中。
当 python 函数被提供给 grad 或 jit 时，它被包装起来以便跟踪并返回。当调用包装的函数时，我们将提供的具体参数抽象到 abstractvalue 类的实例中，将它们框起来用于跟踪跟踪器类的实例，并对它们调用函数。
抽象参数表示一组可能的值，而不是特定的值：例如，jit 将 ndarray 参数抽象为抽象值，这些值表示具有相同形状和数据类型的所有 ndarray。相反，grad 抽象 ndarray 参数来表示底层值的无穷小邻域。通过在这些抽象值上跟踪 python 函数，我们确保它足够专门化，以便转换是可处理的，并且它仍然足够通用，以便转换后的结果是有用的，并且可能是可重用的。然后将这些转换后的函数提升回 python 可调用函数，这样就可以根据需要跟踪并再次转换它们。
jax 跟踪的基本函数大多与 xla hlo 1:1 对应，并在 lax.py 中定义。这种 1:1 的对应关系使得到 xla 的大多数转换基本上都很简单，并且确保我们只有一小组原语来覆盖其他转换，比如自动微分。 jax.numpy 层是用纯 python 编写的，它只是用 lax 函数 (以及我们已经编写的其他 numpy 函数) 表示 numpy 函数。这使得 jax.numpy 易于延展。
当你使用 jax.numpy 时，底层 lax 原语是在后台进行 jit 编译的，允许你在加速器上执行每个原语操作的同时编写不受限制的 python+ numpy 代码。
但是 jax 可以做更多的事情：你可以在越来越大的函数上使用jit来进行端到端编译和优化，而不仅仅是编译和调度到一组固定的单个原语。例如，可以编译整个网络，或者编译整个梯度计算和优化器更新步骤，而不仅仅是编译和调度卷积运算。
折衷之处是，jit 函数必须满足一些额外的专门化需求：因为我们希望编译专门针对形状和数据类型的跟踪，但不是专门针对具体值的跟踪，所以 jit 装饰器下的 python 代码必须适用于抽象值。如果我们尝试在一个抽象的 x 上求 x >0 的值，结果是一个抽象的值，表示集合 {true, false}，所以 python 分支就像 if x > 0 会引起报错。
有关使用 jit 的更多要求，请参见：https://github.com/google/jax#whats-supported
好消息是，jit 是可选的：jax 库在后台对单个操作和函数使用 jit，允许编写不受限制的 python+numpy，同时仍然使用硬件加速器。但是，当你希望最大化性能时，通常可以在自己的代码中使用 jit 编译和端到端优化更大的函数。
后续计划
目前项目小组还将对以下几项做更多尝试和更新：
完善说明文档
支持 cloud tpu
支持多 gpu 和多 tpu
支持完整的 numpy 功能和部分 scipy 功能
全面支持 vmap
加速
降低 xla 函数调度开销
线性代数例程（cpu 上的 mkl 和 gpu 上的 magma）
高效自动微分原语cond和while
有关 jax 库的介绍大致如此。
新大陆NLS-NVF260工业条码扫描器的优势
小米5G新机入网，支持SA/NSA双模5G
如何制作一个能玩的起重机
低噪声、高耐压、双路高音质音频运算放大器“NL8802”发布上市
亮钻由由强强合作，“AI称重”大有前途
TensorFlow有了替代品，竟然还是谷歌自己做出来的？
日韩企业垄断CMOS芯片，中国垄断摄像头模组市场
IPv4报文格式各字段的含义
微软展示多屏显示技术 或将被用于VR头显
FPGA的基本组成结构
基于MXene的全柔性自供电压力传感系统
中低阶智能型手机生物识别大战已开打，指纹识别或将被淘汰
判断两个字符串中的字母是否一致
如何正确选择TRINAMIC运动控制芯片
关于不锈钢板激光焊接变形的研究
vivo x9、oppo r9s、荣耀v9，买哪一款更好？
苹果中国台湾主要供应商9月营收387亿美元，同比下降16.4%
电动车电瓶修复技术——笔记本电脑电池的维护技巧（网友问答）
6 种技术趋势正在改变我们的出行方式
RGB7段显示器的制作