diff --git a/numba_dpex/experimental/decorators.py b/numba_dpex/experimental/decorators.py index 22565eb1cb..d87ae14771 100644 --- a/numba_dpex/experimental/decorators.py +++ b/numba_dpex/experimental/decorators.py @@ -8,6 +8,7 @@ import inspect from numba.core import sigutils +from numba.core.target_extension import jit_registry, target_registry from .kernel_dispatcher import KernelDispatcher @@ -78,3 +79,6 @@ def _specialized_kernel_dispatcher(pyfunc): "the return type as void explicitly." ) return _kernel_dispatcher(func) + + +jit_registry[target_registry["dpex_kernel"]] = kernel diff --git a/numba_dpex/tests/experimental/test_exec_queue_inference.py b/numba_dpex/tests/experimental/test_exec_queue_inference.py index 35dc300ae0..793b863b02 100644 --- a/numba_dpex/tests/experimental/test_exec_queue_inference.py +++ b/numba_dpex/tests/experimental/test_exec_queue_inference.py @@ -6,6 +6,7 @@ import dpctl import dpnp import pytest +from numba.core import config import numba_dpex.experimental as exp_dpex from numba_dpex import Range @@ -34,14 +35,16 @@ def test_successful_execution_queue_inference(): c = dpnp.zeros_like(a, sycl_queue=q) r = Range(100) - # FIXME: This test fails unexpectedly if the NUMBA_CAPTURED_ERRORS is set - # to "new_style". - # Refer: https://github.com/IntelPython/numba-dpex/issues/1195 + current_captured_error_style = config.CAPTURED_ERRORS + config.CAPTURED_ERRORS = "new_style" + try: exp_dpex.call_kernel(add, r, a, b, c) except: pytest.fail("Unexpected error when calling kernel") + config.CAPTURED_ERRORS = current_captured_error_style + assert c[0] == b[0] + a[0] @@ -59,8 +62,6 @@ def test_execution_queue_inference_error(): c = dpnp.zeros_like(a, sycl_queue=q1) r = Range(100) - from numba.core import config - current_captured_error_style = config.CAPTURED_ERRORS config.CAPTURED_ERRORS = "new_style"