3. LAX

LAX介绍¶

LAX是对一条或多条XLA HLO操作指令的封装

jax.lax是一个加强库，类似于jax.numpy。jax.numpy源代码在jax/_src/numpy中。

JAX代码有一个原语（primitives，具体参见：jax.lax package — JAX documentation）的概念，每一个原语可以理解为一个基本数据操作（算子），该操作是对XLA操作指令的简单封装。使用JAX编程时，可以直接调用这些原语，还有一些较复杂的函数是由原语构建起来的，也可以在编程中使用（比如jax.numpy包中的函数）。两者本质上都是相同的，都对应着一系列的XLA操作指令（HLO）。

https://jax.readthedocs.io/en/latest/jax.lax.html

官方建议尽量用jax.numpy而不是直接使用jax.lax，因为jax.numpy更稳定。jax.numpy是由jax.lax实现的，接口更加稳定。