jit_error#
- class braincore.transform.jit_error(pred, err_fun, err_arg=None, scope='any')#
Check errors in a jit function.
>>> def error(arg): >>> raise ValueError(f'error {arg}') >>> x = jax.random.uniform(jax.random.PRNGKey(0), (10,)) >>> jit_error(x.sum() < 5., error, err_arg=x)