jit_error

Contents

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)
Parameters:
  • pred (bool, Array) – The boolean prediction.

  • err_fun (callable) – The error function, which raise errors.

  • err_arg (any) – The arguments which passed into err_f.

  • scope (str) – The scope of the error message. Can be None, ‘all’ or ‘any’.