Skip to content

Commit

Permalink
Add numba impl for CheckAndRaise
Browse files Browse the repository at this point in the history
  • Loading branch information
Adrian Seyboldt committed Nov 15, 2022
1 parent 0bf63eb commit e8eed6c
Showing 1 changed file with 16 additions and 0 deletions.
16 changes: 16 additions & 0 deletions aesara/link/numba/dispatch/extra_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
Unique,
UnravelIndex,
)
from aesara.raise_op import CheckAndRaise


@numba_funcify.register(Bartlett)
Expand Down Expand Up @@ -372,3 +373,18 @@ def broadcast_to(x, *shape):
return np.broadcast_to(x, scalars_shape)

return broadcast_to


@numba_funcify.register(CheckAndRaise)
def numba_funcify_CheckAndRaise(op, node, **kwargs):
error = op.exc_type
msg = op.msg

@numba_basic.numba_njit
def check_and_raise(x, *conditions):
for cond in conditions:
if not cond:
raise error(msg)
return x

return check_and_raise

0 comments on commit e8eed6c

Please sign in to comment.