fix: rvs in jit (#110)

Closes #108 dist.rvs now also works in a compiled context. --------- Co-authored-by: Jan wagner <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Hans Dembinski <[email protected]>
HDembinski · Sep 24, 2024 · 8b3d04a · 8b3d04a
1 parent dc2b36b
commit 8b3d04a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/src/numba_stats/_util.py b/src/numba_stats/_util.py
@@ -187,7 +187,7 @@ def {fname}({args}):
 
 @_overload({fname}, inline="always")
 def _ol_{fname}({args}):
-    return {impl}
+    return {impl}.__wrapped__
 """
         else:
             code = f"""

diff --git a/tests/test_norm.py b/tests/test_norm.py
@@ -65,3 +65,12 @@ def test(x):
     y = test(x)
 
     assert_allclose(y, fn(x, 0, 1))
+
+
+@pytest.mark.filterwarnings("error")
+def test_rvs_njit():
+    @nb.njit
+    def test():
+        return norm.rvs(0.0, 1.0, 10, 1)
+
+    assert_allclose(test(), norm.rvs(0, 1, 10, 1))