jax.experimental.checkify.check_error#

jax.experimental.checkify.check_error(error)[源代码]#

如果 error 表示失败,则引发异常。通过 checkify() 函数化。

此函数的语义等效于

>>> def check_error(err: Error) -> None:
...   err.throw()  # can raise ValueError

但与该实现不同,check_error 可以使用 checkify() 转换进行函数化。

此函数类似于 check(),但签名不同:check() 将布尔谓词和新的错误消息字符串作为参数,而此函数将 Error 值作为参数。check() 和此函数都会在失败时引发 Python 异常(副作用),因此不能通过 jit()pmap()scan() 等进行阶段性输出。两者也可以通过使用 checkify() 进行函数化。

但与 check() 不同,此函数类似于 checkify() 的直接逆运算:checkify() 将可以引发 Python 异常的函数作为输入,并生成一个没有该效果但会生成 Error 值作为输出的新函数,而此 check_error 函数可以接受 Error 值作为输入,并且可以产生引发异常的副作用。也就是说,虽然 checkify() 从可函数化的异常效应转到错误值,但此 check_error 从错误值转到可函数化的异常效应。

当您想要将由 Error 值表示的检查(通过 checkify() 函数化 checks 产生)转换为 Python 异常时,check_error 很有用。

参数:

error (Error) – 要检查的错误。

返回类型:

例如,您可能希望通过 checkify 函数化程序的一部分,通过 jit() 阶段性输出您的函数化代码,然后在 jit() 外部重新注入您的错误值

>>> import jax
>>> from jax.experimental import checkify
>>> def f(x):
...   checkify.check(x>0, "must be positive!")
...   return x
>>> def with_inner_jit(x):
...   checked_f = checkify.checkify(f)
...   # a checkified function can be jitted
...   error, out = jax.jit(checked_f)(x)
...   checkify.check_error(error)
...   return out
>>> _ = with_inner_jit(1)  # no failed check
>>> with_inner_jit(-1)  
Traceback (most recent call last):
  ...
jax._src.JaxRuntimeError: must be positive!
>>> # can re-checkify
>>> error, _ = checkify.checkify(with_inner_jit)(-1)