From 7f9f3f344a1b98d158c9bf22e3e49dcad85dc4dd Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 31 Jul 2026 12:00:07 -0500 Subject: [PATCH] Compile UDFs under Numba's numpy error model A UDF that divides by zero silently corrupted the result. Numba's default error model raises ZeroDivisionError inside the cfunc, where Numba prints the traceback and returns without writing, so GraphBLAS keeps whatever the output element already held. Measured at 6f1eb02: def _idiv(x, y): return x // y op = BinaryOp.register_anonymous(_idiv, "_probe_idiv") op(Vector.from_coo([0, 1], [10, 20]) & Vector.from_coo([0, 1], [2, 0])).new() -> [5, 5] # numpy gives [5, 0]; the 5 is the previous element The float case is the same shape: `[0.5, 0.5]` where numpy gives `[0.5, inf]`. The only user-visible signal is an "Exception ignored" traceback on stderr. Pass `error_model="numpy"` at every `numba.njit` and `numba.cfunc` site for user-defined ops. Setting it on the cfunc alone is not enough: `_build` calls `.compile(sig)` on the Dispatcher first, a Dispatcher keeps one compilation per signature, and whichever compile happens first fixes the model. Reverting only the Dispatcher argument, leaving it on the cfunc, reproduces `[5, 5]`. This changes what a dividing UDF returns, but the behaviour it replaces is an unwritten element rather than an error, so there is no correct result being taken away. --- graphblas/core/operator/base.py | 2 +- graphblas/core/operator/binary.py | 15 +++++++-- graphblas/core/operator/indexbinary.py | 12 +++++-- graphblas/core/operator/indexunary.py | 8 +++-- graphblas/core/operator/udt_utils.py | 2 +- graphblas/core/operator/unary.py | 10 ++++-- graphblas/tests/test_op.py | 46 ++++++++++++++++++++++++++ 7 files changed, 82 insertions(+), 13 deletions(-) diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index a893ba3ad..287e525d5 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -189,7 +189,7 @@ def _finalize_udt_op(parent_op, dtype, dtype2, ret_type, wrapper, wrapper_sig, t from ``typed_user_cls.opclass``. ``dtype2`` is ``None`` for unary ops; the rest pass both. Returns the cached ``TypedUser*Op``. """ - wrapper = numba.cfunc(wrapper_sig, nopython=True)(wrapper) + wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) c_typename = _GB_OBJ_C_TYPENAME[typed_user_cls.opclass] error_label = c_typename.removeprefix("GrB_").removeprefix("GxB_") gb_obj = ffi.new(f"{c_typename}*") diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 97e0a9e70..5329b4415 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -525,7 +525,14 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): if name is None: name = getattr(func, "__name__", "") success = False - binary_udf = numba.njit(func) + # The error model has to be set here, not only on the cfunc wrapper + # below: ``.compile(sig)`` further down builds the specialization the + # wrapper then reuses, and a Dispatcher keeps one compilation per + # signature. Whichever compile happens first fixes the model, so + # setting it only on the wrapper leaves ``x // 0`` raising + # ZeroDivisionError inside a cfunc, where Numba prints the traceback + # and returns, handing GraphBLAS an element it never wrote. + binary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=binary_udf) return_types = {} nt = numba.types @@ -581,7 +588,9 @@ def binary_wrapper(z, x, y): # pragma: no cover (numba) def binary_wrapper(z, x, y): # pragma: no cover (numba) z[0] = binary_udf(x[0], y[0]) - binary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(binary_wrapper) + binary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + binary_wrapper + ) new_binary = ffi_new("GrB_BinaryOp*") check_status_carg( lib.GrB_BinaryOp_new( @@ -981,7 +990,7 @@ def _initialize(cls): (binary.any, _second), ]: binop.orig_func = func - binop._numba_func = numba.njit(func) if _has_numba else None + binop._numba_func = numba.njit(func, error_model="numpy") if _has_numba else None binop._udt_types = {} binop._udt_ops = {} binary.any._numba_func = binary.second._numba_func diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index a19f5ffec..de2ca5f71 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -246,7 +246,9 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): if name is None: name = getattr(func, "__name__", "") success = False - indexbinary_udf = numba.njit(func) + # Set on the Dispatcher, not just the cfunc wrapper; see the note in + # ``BinaryOp._build``. + indexbinary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls( name, func, anonymous=anonymous, is_udt=is_udt, numba_func=indexbinary_udf ) @@ -333,7 +335,9 @@ def indexbinary_wrapper( ): # pragma: no cover (numba) z[0] = indexbinary_udf(x[0], ix, jx, y[0], iy, jy, theta[0]) - indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(indexbinary_wrapper) + indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + indexbinary_wrapper + ) new_idxbinop = ffi_new("GxB_IndexBinaryOp*") check_status_carg( lib.GxB_IndexBinaryOp_new( @@ -392,7 +396,9 @@ def _compile_udt(self, dtype, dtype2): numba_func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type ) - indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(indexbinary_wrapper) + indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + indexbinary_wrapper + ) new_idxbinop = ffi_new("GxB_IndexBinaryOp*") check_status_carg( lib.GxB_IndexBinaryOp_new( diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index c7af8ed78..5820e66c5 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -116,7 +116,9 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): if name is None: name = getattr(func, "__name__", "") success = False - indexunary_udf = numba.njit(func) + # Set on the Dispatcher, not just the cfunc wrapper; see the note in + # ``BinaryOp._build``. + indexunary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls( name, func, anonymous=anonymous, is_udt=is_udt, numba_func=indexunary_udf ) @@ -176,7 +178,9 @@ def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) z[0] = indexunary_udf(x[0], row, col, y[0]) - indexunary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(indexunary_wrapper) + indexunary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + indexunary_wrapper + ) new_indexunary = ffi_new("GrB_IndexUnaryOp*") check_status_carg( lib.GrB_IndexUnaryOp_new( diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index fc33ea8c0..3e2a60d45 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -558,7 +558,7 @@ def _make_record_func(leaf_paths, arity, py_op, *, x_is_scalar=False, y_is_scala func_name="_op", source_label=f"", ) - return numba.njit(op_func) + return numba.njit(op_func, error_model="numpy") def _make_array_wrapper( size, diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index 22822921a..64f68eb48 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -176,7 +176,9 @@ def _build(cls, name, func, *, anonymous=False, is_udt=False): if name is None: name = getattr(func, "__name__", "") success = False - unary_udf = numba.njit(func) + # Set on the Dispatcher, not just the cfunc wrapper; see the note in + # ``BinaryOp._build``. + unary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=unary_udf) return_types = {} nt = numba.types @@ -231,7 +233,9 @@ def unary_wrapper(z, x): def unary_wrapper(z, x): z[0] = unary_udf(x[0]) # pragma: no cover (numba) - unary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(unary_wrapper) + unary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + unary_wrapper + ) new_unary = ffi_new("GrB_UnaryOp*") check_status_carg( lib.GrB_UnaryOp_new( @@ -453,7 +457,7 @@ def _initialize(cls): ]: unop.orig_func = func if _has_numba: - unop._numba_func = numba.njit(func) + unop._numba_func = numba.njit(func, error_model="numpy") else: unop._numba_func = None unop._udt_types = {} diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index e439fc149..de1c01b88 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1359,6 +1359,52 @@ def _this_or_that(val, idx, _, thunk): # pragma: no cover (numba) assert result.isequal(w) +@pytest.mark.skipif("not supports_udfs") +def test_udf_division_by_zero_follows_numpy(): + """Dividing by zero in a UDF returns numpy's answer instead of losing the element. + + Under Numba's default error model the division raises ZeroDivisionError + inside the cfunc, where Numba prints the traceback and returns, so + GraphBLAS keeps whatever was in the output element (in practice the + previous element's value). ``error_model="numpy"`` fixes that, but only if + it is set on the ``njit`` Dispatcher: a Dispatcher holds one compilation + per signature, and ``_build`` calls ``.compile(sig)`` before the wrapper + exists, so setting it on the ``cfunc`` alone comes too late to matter. + """ + + def _idiv(x, y): # pragma: no cover (numba) + return x // y + + op = BinaryOp.register_anonymous(_idiv, "_udf_zero_idiv") + v = Vector.from_coo([0, 1], [10, 20], dtype=dtypes.INT64) + w = Vector.from_coo([0, 1], [2, 0], dtype=dtypes.INT64) + assert op(v & w).new().to_coo()[1].tolist() == [5, 0] + + def _tdiv(x, y): # pragma: no cover (numba) + return x / y + + op = BinaryOp.register_anonymous(_tdiv, "_udf_zero_tdiv") + v = Vector.from_coo([0, 1], [1.0, 1.0], dtype=dtypes.FP64) + w = Vector.from_coo([0, 1], [2.0, 0.0], dtype=dtypes.FP64) + assert op(v & w).new().to_coo()[1].tolist() == [0.5, float("inf")] + + # Same guarantee for a UDT UDF, which reaches the cfunc by another route. + udt = dtypes.register_anonymous( + np.dtype([("dz_a", np.int64), ("dz_b", np.int64)], align=True), "_UdfDivZeroRec" + ) + + def _rec_idiv(x, y): # pragma: no cover (numba) + return (x["dz_a"] // y["dz_a"], x["dz_b"]) + + op = BinaryOp.register_anonymous(_rec_idiv, "_udf_zero_rec", is_udt=True) + v = Vector(udt, size=1) + v[0] = (10, 5) + w = Vector(udt, size=1) + w[0] = (0, 1) + got = v.ewise_mult(w, op).new()[0].new().value + assert (got["dz_a"], got["dz_b"]) == (0, 5) + + @pytest.mark.skipif("not supports_udfs") def test_select_op_outlives_source_indexunary(): """A SelectOp keeps alive the IndexUnaryOp whose GraphBLAS handle it borrows.