jax.lax.optimization_barrier

jax.lax.optimization_barrier#

jax.lax.optimization_barrier(operand, /)[source]#

阻止编译器跨越屏障移动操作。

优化屏障有多种可能的用途

  • 优化屏障确保在任何依赖于屏障输出的操作符之前评估所有输入。这可用于强制执行特定的操作顺序。

  • 优化屏障阻止公共子表达式消除。JAX 使用它来实现重新计算。

  • 优化屏障阻止编译器融合。也就是说，编译器可能不会将屏障之前的操作与屏障之后的操作融合到同一个内核中。

JAX 没有为优化屏障定义导数或批处理规则。

优化屏障在编译函数之外没有作用。

参数：

operand – JAX 值的 pytree。

返回值：

一个 JAX 值的 pytree，其结构和内容与 operand 相同。

示例

阻止对 sin 的两次调用的公共子表达式消除

>>> def f(x):
...   return jax.lax.optimization_barrier(jax.lax.sin(x)) + jax.lax.sin(x)
>>> jax.jit(f)(0.)
Array(0., dtype=float32, weak_type=True)
内容