jax.experimental.enable_x64#

jax.experimental.enable_x64(new_val=True)[源代码]#

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

警告

此上下文管理器已从 JAX v0.8.0 开始弃用,并将在 JAX v0.9.0 中移除。请改用 jax.enable_x64()

用法

>>> import jax
>>> x = np.arange(5, dtype='float64')
>>> with _enable_x64(True):
...   print(jnp.asarray(x).dtype)
...
float64

另请参阅

jax.experimental.disable_x64

临时禁用 X64 模式。

参数:

new_val (bool)