Performs error handling on a "checkified" compiled function during the call.
| 8 | |
| 9 | |
| 10 | class _CheckifyCompiledFnWrapper: |
| 11 | """Performs error handling on a "checkified" compiled function during the call.""" |
| 12 | |
| 13 | def __init__(self, compiled: jax.stages.Compiled): |
| 14 | self._compiled = compiled |
| 15 | |
| 16 | def __call__(self, *args, **kwargs) -> Any: |
| 17 | """Calls the compiled function and raises on detected checkify error.""" |
| 18 | err, result = self._compiled(*args, **kwargs) |
| 19 | checkify.check_error(err) |
| 20 | return result |
| 21 | |
| 22 | |
| 23 | class _CheckifyLoweredFnWrapper: |