jax.disable_jit

jax.disable_jit(disable=True)

上下文管理器，在其动态上下文中禁用 jit() 的行为。

对于调试，有一个在动态上下文中禁用所有 jit() 的机制非常有用。请注意，这不仅禁用用户显式使用的 jit()，还会删除 JAX 库使用的任何隐式 JIT 编译：这包括传递给诸如 scan() 和 while_loop() 等高级原语的 body 和 cond 函数的隐式 JIT 计算，jax.numpy 函数的实现中使用的 JIT，以及 API 实现中使用 jit() 的任何其他情况。但请注意，即使在 disable_jit 下，单个原始操作仍然会像正常的逐个操作的即时执行一样由 XLA 编译。

具有对 JIT 函数参数的数据依赖性的值会被跟踪和抽象化。例如，抽象值可能是 ShapedArray 实例，表示具有给定形状和 dtype 的所有可能数组的集合，但不表示具有特定值的具体数组。如果您在 JIT 函数中使用良性副作用操作（例如 print），您可能会注意到这些。

>>> import jax
>>>
>>> @jax.jit
... def f(x):
...   y = x * 2
...   print("Value of y is", y)
...   return y + 3
...
>>> print(f(jax.numpy.array([1, 2, 3])))  
Value of y is Traced<ShapedArray(int32[3])>with<DynamicJaxprTrace...>
[5 7 9]

这里，y 已被 jit() 抽象为 ShapedArray，它表示具有固定形状和类型但具有任意值的数组。y 的值也被跟踪。如果我们在调试时想看到具体的值，并且也想避免跟踪器，我们可以使用 disable_jit() 上下文管理器。

>>> import jax
>>>
>>> with jax.disable_jit():
...   print(f(jax.numpy.array([1, 2, 3])))
...
Value of y is [2 4 6]
[5 7 9]

参数:

disable (bool)