jax.experimental.enable_x64

jax.experimental.enable_x64(new_val=True)

临时启用 X64 模式的实验性上下文管理器。

警告

此上下文管理器仍然是实验性的，因为它从根本上是损坏的，并且可能导致意外行为，尤其是在与 JAX 转换（如 jax.jit()、 jax.vmap()、 jax.grad() 等）结合使用时。 有关详细信息，请参见 jax-ml/jax#5982。

用法

>>> x = np.arange(5, dtype='float64')
>>> with enable_x64():
...   print(jnp.asarray(x).dtype)
...
float64

另请参阅

jax.experimental.disable_x64

临时禁用 X64 模式。

参数:

new_val (bool)