From f75adabbbce34a1da5f865c01439db428982064f Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 3 Aug 2026 14:58:39 -0500 Subject: [PATCH 01/20] Let UDF registration name its return dtype with ret_dtype A UDF's output type is inferred from what the function returns, matched back to one of the input dtypes by base element type and rank. Inference can only name a type the operator already has in hand, so an output UDT that is not an operand is unreachable. A rank-reducing unary op such as FP64[9] -> FP64[3] cannot be expressed at all: the inferred type is the input's, and the shape check then rejects the shorter array the UDF returns. The workaround of passing the desired type in as an extra operand does not rescue that case either, since two float64 rank-1 UDTs are indistinguishable to the matcher ("matches more than one input array UDT"). Nothing structural was in the way. GrB_UnaryOp_new and friends take ztype separately, and _finalize_udt_op already passes ret_type._carg independently of the operand types. Only the inference step was short-circuiting the choice. Add a ret_dtype= keyword to register_anonymous and register_new on UnaryOp, BinaryOp, IndexUnaryOp, and IndexBinaryOp. It is validated through lookup_dtype, stored on the parent op, and consulted by _udt_ret_type in place of _resolve_udt_return_type at the four call sites. The registration-time shape probe still runs. It lives inside _get_udt_wrapper, which receives the resolved return type, so declaring ret_dtype redirects the fit check rather than skipping it: a UDF whose result cannot fill the declared element is rejected at typing time, the same as in the inferred case. This is worth more here than it is for inference. An inferred type is derived from what the UDF returned and so can hardly contradict it, whereas a declared type is an independent claim the UDF can get wrong. test_udt_ret_dtype_still_shape_checked pins all three behaviors: a bad fit rejected, a broadcast-compatible return accepted, and the record-leaf half of the check. Every choice below is a judgement call, not a forced move. None is load-bearing for the feature, and each is listed with what reversing it would cost, because this is public API and the commitment is the maintainer's to make, not mine. 1. SelectOp gets no ret_dtype at all. For: GraphBLAS fixes a select operator's return type at BOOL, and SelectOp._compile_udt already hardcodes it rather than calling the resolver, so the parameter's only legal value would be its default. Against: the signature is then inconsistent with the other four classes, and a user who does not know the BOOL rule gets a bare TypeError from Python's argument binding rather than an explanation. To decline: add ret_dtype=None to the two SelectOp register methods and raise unless it resolves to BOOL. Purely additive, no caller changes; the only cost is that the error test changes shape. 2. ret_dtype requires is_udt=True. For: the builtin path derives a return type per input dtype by compiling the function against each sample value and then downcasting it toward the input type. One fixed dtype cannot describe that, and forcing one would silently recast results across every builtin type. Against: the downcast heuristic already carries the comment "There should be a way for users to be explicit", so the builtin path is arguably where an explicit return type is most wanted. To decline: honoring it there means overriding ret_type inside the per-sample-value loop in each _build. That is a few lines, but the wrapper signature is built from the return type, so a declared type that Numba will not store into changes a registration-time error into a wrong number. Widening later is compatible; narrowing later is not, which is the argument for starting narrow. 3. The dtype is fixed for the operator, not per input dtype. For: it matches how ret_dtype reads, and it keeps the stored state to a single attribute consulted at each compile. Against: an operator whose output type genuinely varies with its inputs now needs one registration per output type. To decline: accept a callable and call it with the operand dtypes at compile time. Backward compatible, since a DataType and a callable are distinguishable. Deliberately not built now. 4. ret_dtype is rejected with parameterized=True. For: a parameterized operator builds its function when called, and the inner registration graphblas performs accepts no return dtype today, so accepting the keyword at the outer call would promise a path that does not exist until the Parameterized* wrappers forward it. The error says the combination is unsupported rather than pointing at a register call the user cannot reach. Against: a user reasonably expects a keyword to work wherever is_udt does. To decline: thread it through the four Parameterized* classes (slot, __init__ kwarg, and the register_anonymous call in _call). Mechanical and additive; skipped here to keep the diff narrow while binary.py's register internals are being reworked in the preceding commit. 5. No output UDT is invented from the probe. For: the probe is best-effort and runs the UDF once on stand-in values, so a UDF whose output shape depends on its input values would mint the wrong type silently. Against: for the common value-independent UDF the type could have been inferred with no user input at all. To decline: this one should stay declined. A silently wrong dtype is the worst failure mode available here, and the explicit keyword costs the user one argument. Pickling preserves the declared type. The shared __reduce__ predated ret_dtype, so a cross-process round trip re-registered the op without it and the reconstruction failed its own shape probe; in-process pickling hid this, because _deserialize_udf finds the already-registered object and returns it. The reduce tuple now carries ret_dtype, passed on only when set, so SelectOp (which shares the path and has no such slot) and pickles written before this keyword keep working. The test crosses a real process boundary for exactly that reason. ret_dtype is also validated eagerly under lazy=True, so an invalid combination fails at the registration site instead of at first attribute touch of the delayed op. --- docs/user_guide/udt.rst | 37 ++++ graphblas/core/operator/base.py | 57 +++++- graphblas/core/operator/binary.py | 60 +++++- graphblas/core/operator/indexbinary.py | 61 ++++-- graphblas/core/operator/indexunary.py | 56 ++++-- graphblas/core/operator/udf.py | 26 +++ graphblas/core/operator/unary.py | 55 +++++- graphblas/tests/test_op.py | 263 +++++++++++++++++++++++++ 8 files changed, 567 insertions(+), 48 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 88b77086a..7f247388d 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -371,6 +371,43 @@ If your UDF references a field that doesn't exist, or returns the wrong arity, you'll get a ``UdfParseError`` with the actionable diagnostic line surfaced from Numba's typing pass instead of a 200-line traceback. +Naming the output type +~~~~~~~~~~~~~~~~~~~~~~ + +By default the return type is worked out from what the UDF returns, matched +against the input dtypes. That can only name a type the operator already has +in hand, so an output UDT that is not one of the operands is out of reach. +The ``x[:2]`` rejection above is a case of this: the output really is a +2-element UDT, but nothing says so. + +Pass ``ret_dtype`` to say it outright. It takes anything ``lookup_dtype`` +accepts and requires ``is_udt=True``:: + + nine = gb.dtypes.register_anonymous(np.dtype((np.float64, (9,))), "Nine") + three = gb.dtypes.register_anonymous(np.dtype((np.float64, (3,))), "Three") + + head = gb.core.operator.UnaryOp.register_anonymous( + lambda x: x[:3], "head", is_udt=True, ret_dtype=three + ) + head[nine].return_type # Three + +``ret_dtype`` is available on ``register_anonymous`` and ``register_new`` for +``UnaryOp``, ``BinaryOp``, ``IndexUnaryOp``, and ``IndexBinaryOp``. It is a +property of the operator, not of a particular input dtype: the same output +type applies to every dtype the operator is typed for. An operator whose +output type should vary with its inputs still needs one registration per +output type. + +Declaring the type does not switch off the shape check described above; it +points the check at the declared type instead. A UDF whose result cannot +fill a ``ret_dtype`` element is still rejected when the op is typed, so a +wrong ``ret_dtype`` is a registration error rather than a silently +mis-typed result. + +It does not apply to ``SelectOp``, whose return type GraphBLAS fixes at +``BOOL``, nor to builtin (non-UDT) dtypes, where the return type comes from +compiling the function against each input type in turn. + .. _udt_jit_introspection: JIT and introspection diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index f895c7a87..1ea1d5474 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -93,6 +93,39 @@ def _hasop(module, name): ) +def _validate_ret_dtype(ret_dtype, opclass, *, is_udt, parameterized): + """Normalize a user-supplied ``ret_dtype`` to a DataType, or raise. + + ``ret_dtype`` names the operator's output type outright instead of letting + it be inferred from what the UDF returns. Inference can only name a type it + can see, which is why this is limited to the UDT path: the builtin path + derives its output from Numba's typing of each sample input, and forcing a + single type across all of them would silently recast results. + """ + if ret_dtype is None: + return None + if not is_udt: + raise ValueError( + f"{opclass}: ret_dtype requires is_udt=True. The return type for builtin " + f"dtypes comes from compiling the function for each input type, so a single " + f"fixed type cannot describe it." + ) + if parameterized: + raise ValueError( + f"{opclass}: ret_dtype does not work with parameterized=True. " + f"A parameterized operator builds and registers its function when called, " + f"and that inner registration does not accept a return dtype; register the " + f"built function without parameterized=True to declare one." + ) + try: + return lookup_dtype(ret_dtype) + except (ValueError, TypeError) as exc: + raise ValueError( + f"{opclass}: ret_dtype={ret_dtype!r} is not a recognized dtype. " + f"Pass a DataType, a numpy dtype, or a name such as 'FP64'." + ) from exc + + class OpPath: def __init__(self, parent, name): self._parent = parent @@ -576,21 +609,25 @@ def _deserialize(cls, name, *args): return cls.register_new(name, *args) @classmethod - def _deserialize_udf(cls, name, orig_func, is_udt): + def _deserialize_udf(cls, name, orig_func, is_udt, ret_dtype=None): """Re-register a named UDF on unpickle, or reuse if already present. Shared by the five UDF-capable subclasses (UnaryOp, BinaryOp, IndexUnaryOp, SelectOp, IndexBinaryOp), all of which use the - default ``__reduce__`` below. + default ``__reduce__`` below. ``ret_dtype`` is passed only when set: + SelectOp shares this path and takes no ret_dtype, and pickles written + before ret_dtype existed carry a 3-tuple. """ if (rv := cls._find(name)) is not None: return rv - return cls.register_new(name, orig_func, is_udt=is_udt) + kwargs = {} if ret_dtype is None else {"ret_dtype": ret_dtype} + return cls.register_new(name, orig_func, is_udt=is_udt, **kwargs) @classmethod - def _deserialize_anon_udf(cls, func, name, is_udt): + def _deserialize_anon_udf(cls, func, name, is_udt, ret_dtype=None): """Re-register an anonymous UDF on unpickle.""" - return cls.register_anonymous(func, name, is_udt=is_udt) + kwargs = {} if ret_dtype is None else {"ret_dtype": ret_dtype} + return cls.register_anonymous(func, name, is_udt=is_udt, **kwargs) def __reduce__(self): """Default ``__reduce__`` for UDF-capable subclasses. @@ -603,10 +640,16 @@ def __reduce__(self): if self._anonymous: if hasattr(self.orig_func, "_parameterized_info"): return (_deserialize_parameterized, self.orig_func._parameterized_info) - return (type(self)._deserialize_anon_udf, (self.orig_func, self.name, self._is_udt)) + return ( + type(self)._deserialize_anon_udf, + (self.orig_func, self.name, self._is_udt, getattr(self, "_ret_dtype", None)), + ) if (name := f"{self._modname}.{self.name}") in _STANDARD_OPERATOR_NAMES: return name - return (type(self)._deserialize_udf, (self.name, self.orig_func, self._is_udt)) + return ( + type(self)._deserialize_udf, + (self.name, self.orig_func, self._is_udt, getattr(self, "_ret_dtype", None)), + ) @classmethod def _check_supports_udf(cls, method_name): diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 517ba72e1..57535f5ef 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -37,6 +37,7 @@ TypedOpBase, _call_op, _hasop, + _validate_ret_dtype, ) if _has_numba: @@ -48,7 +49,7 @@ _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper, - _resolve_udt_return_type, + _udt_ret_type, ) try: @@ -608,6 +609,7 @@ class BinaryOp(OpBase): "_numba_func", "_custom_dtype", "_defer_builds", + "_ret_dtype", ) _module = binary _modname = "binary" @@ -713,11 +715,17 @@ class BinaryOp(OpBase): } @classmethod - def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False): + def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False, ret_dtype=None): if not isinstance(func, FunctionType): raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") + # This rejects ret_dtype unless is_udt, which keeps it disjoint from the + # deferred builtin path below: that one only runs under ``not is_udt``, + # so an op can never be both deferred and carrying a declared return + # type. ``_build_deferred`` reads its ret_type back from ``.types`` and + # never consults ``_ret_dtype``. + ret_dtype = _validate_ret_dtype(ret_dtype, "binary", is_udt=is_udt, parameterized=False) success = False # The error model has to be set here as well as on the cfunc wrapper in # ``_finalize_typed_binaryop``. A Dispatcher keeps one compilation per @@ -741,7 +749,14 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False): binary_udf = numba.njit(func, error_model="numpy") else: binary_udf = numba.njit(func, error_model="numpy") - new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=binary_udf) + new_type_obj = cls( + name, + func, + anonymous=anonymous, + is_udt=is_udt, + numba_func=binary_udf, + ret_dtype=ret_dtype, + ) return_types = {} if not is_udt: # ``cache=True`` marks the module-level built-in UDFs (floordiv and @@ -827,7 +842,7 @@ def _compile_udt(self, dtype, dtype2): numba_func, sig, op_kind="binary", op_name=self.name, dtypes=(dtype, dtype2) ) numba_ret_type = numba_func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, dtype, dtype2) + ret_type = _udt_ret_type(self, numba_ret_type, dtype, dtype2) binary_wrapper, wrapper_sig = _get_udt_wrapper( numba_func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type ) @@ -886,7 +901,9 @@ def _compile_udt(self, dtype, dtype2): return op @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): + def register_anonymous( + cls, func, name=None, *, parameterized=False, is_udt=False, ret_dtype=None + ): """Register a BinaryOp without registering it in the ``graphblas.binary`` namespace. Because it is not registered in the namespace, the name is optional. @@ -915,6 +932,13 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals Setting ``is_udt=True`` is also helpful when the left and right dtypes need to be different. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + Without it the return type is inferred from what the function + returns, which can only name a type that is already an input, so + an output UDT that is not an operand needs this. The dtype is + fixed for the operator: it is the same for every input dtype. + Returns ------- BinaryOp or ParameterizedBinaryOp @@ -922,12 +946,21 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals """ cls._check_supports_udf("register_anonymous") if parameterized: + _validate_ret_dtype(ret_dtype, "binary", is_udt=is_udt, parameterized=True) return ParameterizedBinaryOp(name, func, anonymous=True, is_udt=is_udt) - return cls._build(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt, ret_dtype=ret_dtype) @classmethod def register_new( - cls, name, func, *, parameterized=False, is_udt=False, lazy=False, _cache=False + cls, + name, + func, + *, + parameterized=False, + is_udt=False, + lazy=False, + ret_dtype=None, + _cache=False, ): """Register a new BinaryOp and save it to ``graphblas.binary`` namespace. @@ -963,6 +996,10 @@ def register_new( delay compilation and only compile when the operator is used, which is done by setting ``lazy=True``. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + See :meth:`register_anonymous` for details. + Examples -------- >>> def max_zero(x, y): @@ -991,6 +1028,9 @@ def register_new( """ cls._check_supports_udf("register_new") + # Validate eagerly even for lazy=True, so a bad combination fails at + # the registration site rather than at first attribute touch. + _validate_ret_dtype(ret_dtype, "binary", is_udt=is_udt, parameterized=parameterized) module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( @@ -1000,14 +1040,16 @@ def register_new( "func": func, "parameterized": parameterized, "is_udt": is_udt, + "ret_dtype": ret_dtype, "_cache": _cache, }, ) elif parameterized: + _validate_ret_dtype(ret_dtype, "binary", is_udt=is_udt, parameterized=True) binary_op = ParameterizedBinaryOp(name, func, is_udt=is_udt) setattr(module, funcname, binary_op) else: - binary_op = cls._build(name, func, is_udt=is_udt, cache=_cache) + binary_op = cls._build(name, func, is_udt=is_udt, cache=_cache, ret_dtype=ret_dtype) setattr(module, funcname, binary_op) # Also save it to `graphblas.op` if not yet defined opmodule, funcname = cls._remove_nesting(name, module=op, modname="op", strict=False) @@ -1232,8 +1274,10 @@ def __init__( is_positional=False, is_udt=False, numba_func=None, + ret_dtype=None, ): super().__init__(name, anonymous=anonymous) + self._ret_dtype = ret_dtype self._monoid = None self._commutes_to = None self._semiring_commutes_to = None diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index 0e0eff04d..60cef2b4e 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -9,7 +9,7 @@ from ...exceptions import EmptyObject, UdfParseError, check_status_carg from .. import _has_numba, ffi, lib from ..dtypes import _sample_values, unify -from .base import OpBase, ParameterizedUdf, TypedOpBase +from .base import OpBase, ParameterizedUdf, TypedOpBase, _validate_ret_dtype from .binary import BinaryOp _has_idxbinop = hasattr(lib, "GxB_IndexBinaryOp_new") @@ -22,7 +22,7 @@ _compile_for_dtype, _compile_udf_for_udt, _get_udt_wrapper_indexbinary, - _resolve_udt_return_type, + _udt_ret_type, ) ffi_new = ffi.new @@ -330,7 +330,7 @@ class IndexBinaryOp(OpBase): no built-ins; all IndexBinaryOps are user-defined. """ - __slots__ = "orig_func", "_is_udt", "_numba_func" + __slots__ = "orig_func", "_is_udt", "_numba_func", "_ret_dtype" _module = indexbinary _modname = "indexbinary" _custom_dtype = None @@ -349,7 +349,7 @@ class IndexBinaryOp(OpBase): } @classmethod - def _build(cls, name, func, *, is_udt=False, anonymous=False): + def _build(cls, name, func, *, is_udt=False, anonymous=False, ret_dtype=None): if not _has_idxbinop: raise RuntimeError( "IndexBinaryOp requires SuiteSparse:GraphBLAS 9.4+ " @@ -359,12 +359,20 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") + ret_dtype = _validate_ret_dtype( + ret_dtype, "indexbinary", is_udt=is_udt, parameterized=False + ) success = False # 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 + name, + func, + anonymous=anonymous, + is_udt=is_udt, + numba_func=indexbinary_udf, + ret_dtype=ret_dtype, ) return_types = {} nt = numba.types @@ -509,7 +517,7 @@ def _compile_udt(self, dtype, dtype2, theta_type=None): numba_func, sig, op_kind="indexbinary", op_name=self.name, dtypes=dtypes ) numba_ret_type = numba_func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, *dtypes) + ret_type = _udt_ret_type(self, numba_ret_type, *dtypes) indexbinary_wrapper, wrapper_sig = _get_udt_wrapper_indexbinary( numba_func, ret_type, dtype, dtype2, theta_type, numba_ret_type=numba_ret_type ) @@ -572,7 +580,9 @@ def _add(self, op, *, is_jit=False): self._udt_ops[dtypes] = op @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): + def register_anonymous( + cls, func, name=None, *, parameterized=False, is_udt=False, ret_dtype=None + ): """Register an IndexBinaryOp without adding it to the ``indexbinary`` namespace. Because it is not registered in the namespace, the name is optional. @@ -593,6 +603,13 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals is_udt : bool, default False Whether the operator is intended to operate on user-defined types. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + Without it the return type is inferred from what the function + returns, which can only name a type that is already an input, so + an output UDT that is not an operand needs this. The dtype is + fixed for the operator: it is the same for every input dtype. + Returns ------- IndexBinaryOp or ParameterizedIndexBinaryOp @@ -600,11 +617,14 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals """ cls._check_supports_udf("register_anonymous") if parameterized: + _validate_ret_dtype(ret_dtype, "indexbinary", is_udt=is_udt, parameterized=True) return ParameterizedIndexBinaryOp(name, func, anonymous=True, is_udt=is_udt) - return cls._build(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt, ret_dtype=ret_dtype) @classmethod - def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + def register_new( + cls, name, func, *, parameterized=False, is_udt=False, lazy=False, ret_dtype=None + ): """Register a new IndexBinaryOp under the ``graphblas.indexbinary`` namespace. Parameters @@ -624,23 +644,37 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal lazy : bool, default False When True, defer compilation until the operator is first used. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + See :meth:`register_anonymous` for details. + Examples -------- >>> gb.indexbinary.register_new("index_dist", lambda x, ix, jx, y, iy, jy, t: abs(ix - iy)) """ cls._check_supports_udf("register_new") + # Validate eagerly even for lazy=True, so a bad combination fails at + # the registration site rather than at first attribute touch. + _validate_ret_dtype(ret_dtype, "indexbinary", is_udt=is_udt, parameterized=parameterized) module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( cls.register_new, - {"name": name, "func": func, "parameterized": parameterized, "is_udt": is_udt}, + { + "name": name, + "func": func, + "parameterized": parameterized, + "is_udt": is_udt, + "ret_dtype": ret_dtype, + }, ) elif parameterized: + _validate_ret_dtype(ret_dtype, "indexbinary", is_udt=is_udt, parameterized=True) idxbinop = ParameterizedIndexBinaryOp(name, func, is_udt=is_udt) setattr(module, funcname, idxbinop) else: - idxbinop = cls._build(name, func, is_udt=is_udt) + idxbinop = cls._build(name, func, is_udt=is_udt, ret_dtype=ret_dtype) setattr(module, funcname, idxbinop) if not cls._initialized: # pragma: no cover (safety) @@ -656,8 +690,11 @@ def _initialize(cls): # No built-in IndexBinaryOps to register. cls._initialized = True - def __init__(self, name, func=None, *, anonymous=False, is_udt=False, numba_func=None): + def __init__( + self, name, func=None, *, anonymous=False, is_udt=False, numba_func=None, ret_dtype=None + ): super().__init__(name, anonymous=anonymous) + self._ret_dtype = ret_dtype self.orig_func = func self._numba_func = numba_func self._is_udt = is_udt diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index 7d187bfe7..f3c6b57f6 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -7,7 +7,7 @@ from ...exceptions import UdfParseError, check_status_carg from .. import _has_numba, ffi, lib from ..dtypes import _sample_values -from .base import OpBase, ParameterizedUdf, TypedOpBase, _call_op +from .base import OpBase, ParameterizedUdf, TypedOpBase, _call_op, _validate_ret_dtype if _has_numba: import numba @@ -18,7 +18,7 @@ _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper, - _resolve_udt_return_type, + _udt_ret_type, ) ffi_new = ffi.new @@ -86,7 +86,7 @@ class IndexUnaryOp(OpBase): Built-in and registered IndexUnaryOps are located in the ``graphblas.indexunary`` namespace. """ - __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func" + __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func", "_ret_dtype" _module = indexunary _modname = "indexunary" _custom_dtype = None @@ -111,17 +111,23 @@ class IndexUnaryOp(OpBase): "rowindex", "colindex"} # fmt: skip @classmethod - def _build(cls, name, func, *, is_udt=False, anonymous=False): + def _build(cls, name, func, *, is_udt=False, anonymous=False, ret_dtype=None): if not isinstance(func, FunctionType): raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") + ret_dtype = _validate_ret_dtype(ret_dtype, "indexunary", is_udt=is_udt, parameterized=False) success = False # 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 + name, + func, + anonymous=anonymous, + is_udt=is_udt, + numba_func=indexunary_udf, + ret_dtype=ret_dtype, ) return_types = {} nt = numba.types @@ -218,7 +224,7 @@ def _compile_udt(self, dtype, dtype2): numba_func, sig, op_kind="indexunary", op_name=self.name, dtypes=(dtype, dtype2) ) numba_ret_type = numba_func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, dtype, dtype2) + ret_type = _udt_ret_type(self, numba_ret_type, dtype, dtype2) indexunary_wrapper, wrapper_sig = _get_udt_wrapper( numba_func, ret_type, dtype, dtype2, include_indexes=True, numba_ret_type=numba_ret_type ) @@ -227,7 +233,9 @@ def _compile_udt(self, dtype, dtype2): ) @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): + def register_anonymous( + cls, func, name=None, *, parameterized=False, is_udt=False, ret_dtype=None + ): """Register a IndexUnary without registering it in the ``graphblas.indexunary`` namespace. Because it is not registered in the namespace, the name is optional. @@ -261,6 +269,13 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals Setting ``is_udt=True`` is also helpful when the left and right dtypes need to be different. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + Without it the return type is inferred from what the function + returns, which can only name a type that is already an input, so + an output UDT that is not an operand needs this. The dtype is + fixed for the operator: it is the same for every input dtype. + Returns ------- return IndexUnaryOp or ParameterizedIndexUnaryOp @@ -268,11 +283,14 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals """ cls._check_supports_udf("register_anonymous") if parameterized: + _validate_ret_dtype(ret_dtype, "indexunary", is_udt=is_udt, parameterized=True) return ParameterizedIndexUnaryOp(name, func, anonymous=True, is_udt=is_udt) - return cls._build(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt, ret_dtype=ret_dtype) @classmethod - def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + def register_new( + cls, name, func, *, parameterized=False, is_udt=False, lazy=False, ret_dtype=None + ): """Register a new IndexUnaryOp and save it to ``graphblas.indexunary`` namespace. If the return type is Boolean, the function will also be registered as a SelectOp @@ -315,6 +333,10 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal delay compilation and only compile when the operator is used, which is done by setting ``lazy=True``. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + See :meth:`register_anonymous` for details. + Examples -------- >>> gb.indexunary.register_new("row_mod", lambda x, i, j, thunk: i % max(thunk, 2)) @@ -323,17 +345,27 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal """ cls._check_supports_udf("register_new") + # Validate eagerly even for lazy=True, so a bad combination fails at + # the registration site rather than at first attribute touch. + _validate_ret_dtype(ret_dtype, "indexunary", is_udt=is_udt, parameterized=parameterized) module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( cls.register_new, - {"name": name, "func": func, "parameterized": parameterized, "is_udt": is_udt}, + { + "name": name, + "func": func, + "parameterized": parameterized, + "is_udt": is_udt, + "ret_dtype": ret_dtype, + }, ) elif parameterized: + _validate_ret_dtype(ret_dtype, "indexunary", is_udt=is_udt, parameterized=True) indexunary_op = ParameterizedIndexUnaryOp(name, func, is_udt=is_udt) setattr(module, funcname, indexunary_op) else: - indexunary_op = cls._build(name, func, is_udt=is_udt) + indexunary_op = cls._build(name, func, is_udt=is_udt, ret_dtype=ret_dtype) setattr(module, funcname, indexunary_op) # If return type is BOOL, register additionally as a SelectOp if all(x == BOOL for x in indexunary_op.types.values()): @@ -397,8 +429,10 @@ def __init__( is_positional=False, is_udt=False, numba_func=None, + ret_dtype=None, ): super().__init__(name, anonymous=anonymous) + self._ret_dtype = ret_dtype self.orig_func = func self._numba_func = numba_func self.is_positional = is_positional diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index ab5f0b8a2..2a840eb41 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -190,6 +190,32 @@ def _summarize_numba_typing_error(exc): return "Numba could not compile the function for these input types" +def _array_udt_operands(dtypes): + """Return the distinct array-UDT operands in ``dtypes``, order preserved. + + ``x`` and ``y`` are usually the same UDT, hence the dedup. It is by + identity, not by shape: a flat ``FP64[2, 3]`` and a layered + ``FP64[3][2]`` are separate DataTypes with separate GrB_Type handles, + yet Numba collapses both to the same ``nestedarray(float64, (2, 3))``. + """ + operands = [] + for d in dtypes: + if d._is_udt and d.np_type.subdtype is not None and not any(d is seen for seen in operands): + operands.append(d) + return operands + + +def _udt_ret_type(parent_op, numba_ret_type, *dtypes): + """Return the operator's declared ``ret_dtype``, else infer one from the UDF. + + Inference can only name a type that is already an operand, so an output + UDT that appears nowhere in the inputs is unreachable without this. + """ + if (ret_dtype := parent_op._ret_dtype) is not None: + return ret_dtype + return _resolve_udt_return_type(numba_ret_type, *dtypes) + + def _resolve_udt_return_type(numba_ret_type, *dtypes): """Resolve a Numba return type to a DataType, matching Tuple returns to an input UDT. diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index b6e615ca5..a262c4dc9 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -28,6 +28,7 @@ ParameterizedUdf, TypedOpBase, _hasop, + _validate_ret_dtype, ) if _supports_complex: @@ -41,7 +42,7 @@ _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper, - _resolve_udt_return_type, + _udt_ret_type, ) ffi_new = ffi.new @@ -138,7 +139,7 @@ class UnaryOp(OpBase): as well as in the ``graphblas.op`` combined namespace. """ - __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func" + __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func", "_ret_dtype" _custom_dtype = None _module = unary _modname = "unary" @@ -172,16 +173,24 @@ class UnaryOp(OpBase): _positional = {"positioni", "positioni1", "positionj", "positionj1"} @classmethod - def _build(cls, name, func, *, anonymous=False, is_udt=False): + def _build(cls, name, func, *, anonymous=False, is_udt=False, ret_dtype=None): if type(func) is not FunctionType: raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") + ret_dtype = _validate_ret_dtype(ret_dtype, "unary", is_udt=is_udt, parameterized=False) success = False # 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) + new_type_obj = cls( + name, + func, + anonymous=anonymous, + is_udt=is_udt, + numba_func=unary_udf, + ret_dtype=ret_dtype, + ) return_types = {} nt = numba.types if not is_udt: @@ -288,7 +297,7 @@ def _compile_udt(self, dtype, dtype2): sig = (dtype.numba_type,) _compile_udf_for_udt(numba_func, sig, op_kind="unary", op_name=self.name, dtypes=(dtype,)) numba_ret_type = numba_func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, dtype) + ret_type = _udt_ret_type(self, numba_ret_type, dtype) unary_wrapper, wrapper_sig = _get_udt_wrapper( numba_func, ret_type, dtype, numba_ret_type=numba_ret_type ) @@ -297,7 +306,9 @@ def _compile_udt(self, dtype, dtype2): ) @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): + def register_anonymous( + cls, func, name=None, *, parameterized=False, is_udt=False, ret_dtype=None + ): """Register a UnaryOp without registering it in the ``graphblas.unary`` namespace. Because it is not registered in the namespace, the name is optional. @@ -324,6 +335,12 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals Whether the operator is intended to operate on user-defined types. If True, then the function will not be automatically compiled for builtin types, and it will be compiled "just in time" when used. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + Without it the return type is inferred from what the function + returns, which can only name a type that is already an input, so + an output UDT that is not an operand needs this. The dtype is + fixed for the operator: it is the same for every input dtype. Returns ------- @@ -332,11 +349,14 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals """ cls._check_supports_udf("register_anonymous") if parameterized: + _validate_ret_dtype(ret_dtype, "unary", is_udt=is_udt, parameterized=True) return ParameterizedUnaryOp(name, func, anonymous=True, is_udt=is_udt) - return cls._build(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt, ret_dtype=ret_dtype) @classmethod - def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + def register_new( + cls, name, func, *, parameterized=False, is_udt=False, lazy=False, ret_dtype=None + ): """Register a new UnaryOp and save it to ``graphblas.unary`` namespace. Parameters @@ -369,6 +389,9 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal Compiling functions can be slow, however, so you may want to delay compilation and only compile when the operator is used, which is done by setting ``lazy=True``. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + See :meth:`register_anonymous` for details. Examples -------- @@ -378,17 +401,27 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal """ cls._check_supports_udf("register_new") + # Validate eagerly even for lazy=True, so a bad combination fails at + # the registration site rather than at first attribute touch. + _validate_ret_dtype(ret_dtype, "unary", is_udt=is_udt, parameterized=parameterized) module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( cls.register_new, - {"name": name, "func": func, "parameterized": parameterized, "is_udt": is_udt}, + { + "name": name, + "func": func, + "parameterized": parameterized, + "is_udt": is_udt, + "ret_dtype": ret_dtype, + }, ) elif parameterized: + _validate_ret_dtype(ret_dtype, "unary", is_udt=is_udt, parameterized=True) unary_op = ParameterizedUnaryOp(name, func, is_udt=is_udt) setattr(module, funcname, unary_op) else: - unary_op = cls._build(name, func, is_udt=is_udt) + unary_op = cls._build(name, func, is_udt=is_udt, ret_dtype=ret_dtype) setattr(module, funcname, unary_op) # Also save it to `graphblas.op` if not yet defined opmodule, funcname = cls._remove_nesting(name, module=op, modname="op", strict=False) @@ -485,12 +518,14 @@ def __init__( is_positional=False, is_udt=False, numba_func=None, + ret_dtype=None, ): super().__init__(name, anonymous=anonymous) self.orig_func = func self._numba_func = numba_func self.is_positional = is_positional self._is_udt = is_udt + self._ret_dtype = ret_dtype if is_udt: self._udt_types = {} # {dtype: DataType} self._udt_ops = {} # {dtype: TypedUserUnaryOp} diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index ea39ac5ac..8ed3070d5 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -7273,3 +7273,266 @@ def test_deferred_commutes_to(): assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report assert lines.get("floordiv_ok") == "True", report assert lines.get("absfirst_ok") == "True", report + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_ret_dtype_names_an_output_udt(): + """``ret_dtype`` names an output UDT that is not one of the operands. + + Without it the return type is inferred from what the UDF builds, and the + only names in scope are the input dtypes. That makes a rank-reducing op + such as FP64[9] -> FP64[3] unreachable: the inferred type is the input's, + and the shape check then rejects the shorter array the UDF returns. + """ + # Shapes unique to this test; see the note in test_udt_array_udf_shape_errors. + nine = dtypes.register_anonymous(np.dtype((np.float64, (9,))), "_RetD9") + three = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_RetD3") + + def _first_three(x): # pragma: no cover (numba) + return x[:3] + + without = UnaryOp.register_anonymous(_first_three, "_ret_dtype_without", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(3,\) when run on sample values"): + without[nine] + + op_ = UnaryOp.register_anonymous(_first_three, "_ret_dtype_with", is_udt=True, ret_dtype=three) + assert op_[nine].return_type is three + + v = Vector(nine, size=2) + v[0] = np.arange(9.0) + w = op_(v).new() + assert w.dtype is three + np.testing.assert_array_equal(w[0].new().value, np.arange(3.0)) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_ret_dtype_binary_and_record(): + """``ret_dtype`` on a binary op, including a type that is neither operand.""" + a4 = dtypes.register_anonymous(np.dtype((np.float64, (4,))), "_RetDBinA4") + b6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_RetDBinB6") + out2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_RetDBinOut2") + + def _head_two(x, y): # pragma: no cover (numba) + return x[:2] + y[:2] + + op_ = BinaryOp.register_anonymous(_head_two, "_ret_dtype_bin", is_udt=True, ret_dtype=out2) + assert op_[a4, b6].return_type is out2 + + # ret_dtype equal to an operand's dtype agrees with what inference picks. + def _plus(x, y): # pragma: no cover (numba) + return x + y + + same = BinaryOp.register_anonymous(_plus, "_ret_dtype_same", is_udt=True, ret_dtype=a4) + inferred = BinaryOp.register_anonymous(_plus, "_ret_dtype_inferred", is_udt=True) + assert same[a4, a4].return_type is inferred[a4, a4].return_type is a4 + + # A record UDT output that appears in no operand. + rec = dtypes.register_anonymous( + np.dtype([("lo", np.float64), ("hi", np.float64)], align=True), "_RetDRec" + ) + + def _bounds(x, y): # pragma: no cover (numba) + return (min(x, y), max(x, y)) + + recop = BinaryOp.register_anonymous(_bounds, "_ret_dtype_rec", is_udt=True, ret_dtype=rec) + assert recop[FP64, FP64].return_type is rec + + w = Vector(FP64, size=2) + w[0] = 3.0 + u = Vector(FP64, size=2) + u[0] = 1.0 + res = recop(w & u).new() + assert res.dtype is rec + assert tuple(res[0].new().value.tolist()) == (1.0, 3.0) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_ret_dtype_errors(): + """``ret_dtype`` is rejected outside the UDT path and for unrecognized dtypes.""" + with pytest.raises(ValueError, match="not a recognized dtype"): + UnaryOp.register_anonymous(lambda x: x, "_ret_dtype_junk", is_udt=True, ret_dtype="NOPE") + with pytest.raises(ValueError, match="not a recognized dtype"): + UnaryOp.register_anonymous(lambda x: x, "_ret_dtype_junk2", is_udt=True, ret_dtype=object()) + + # The builtin path derives a return type per input dtype, so one fixed + # dtype cannot describe it. + with pytest.raises(ValueError, match="ret_dtype requires is_udt=True"): + UnaryOp.register_anonymous(lambda x: x, "_ret_dtype_builtin", ret_dtype=FP64) + with pytest.raises(ValueError, match="ret_dtype requires is_udt=True"): + BinaryOp.register_anonymous(lambda x, y: x + y, "_ret_dtype_builtin2", ret_dtype=FP64) + + # lazy=True must not defer the validation: a bad combination fails at the + # registration site, not at first attribute touch of the delayed op. + with pytest.raises(ValueError, match="parameterized=True"): + UnaryOp.register_new( + "_ret_dtype_lazy_param", + lambda x: x, + parameterized=True, + is_udt=True, + lazy=True, + ret_dtype=FP64, + ) + + # A parameterized operator builds its function when called; ret_dtype + # belongs to the register call for that function. + with pytest.raises(ValueError, match="does not work with parameterized=True"): + UnaryOp.register_anonymous( + lambda t=1: (lambda x: x + t), + "_ret_dtype_param", + parameterized=True, + is_udt=True, + ret_dtype=FP64, + ) + + # SelectOp takes no ret_dtype: GraphBLAS fixes its return type to BOOL. + with pytest.raises(TypeError, match="ret_dtype"): + SelectOp.register_anonymous( + lambda x, i, j, t: x > t, "_ret_dtype_select", is_udt=True, ret_dtype=BOOL + ) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_ret_dtype_index_ops(): + """``ret_dtype`` reaches the IndexUnaryOp and IndexBinaryOp compile paths too.""" + in5 = dtypes.register_anonymous(np.dtype((np.float64, (5,))), "_RetDIdx5") + out2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_RetDIdx2") + + def _head_plus_row(x, i, j, t): # pragma: no cover (numba) + return x[:2] + i + + iu = IndexUnaryOp.register_anonymous( + _head_plus_row, "_ret_dtype_indexunary", is_udt=True, ret_dtype=out2 + ) + assert iu[in5, INT64].return_type is out2 + + if lib.__dict__.get("GxB_IndexBinaryOp_new") is not None: + from graphblas.core.operator import IndexBinaryOp + + def _head_sum(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x[:2] + y[:2] + + ib = IndexBinaryOp.register_anonymous( + _head_sum, "_ret_dtype_indexbinary", is_udt=True, ret_dtype=out2 + ) + assert ib[in5, in5].return_type is out2 + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_ret_dtype_still_shape_checked(): + """The registration-time shape probe checks against the declared ``ret_dtype``. + + Short-circuiting the return-type inference must not take the probe with + it. Declaring the output type makes the check worth more, not less: the + inferred case can only ever compare a UDF against a type derived from + what it returned, while a declared type is an independent claim the UDF + can contradict. + """ + in5 = dtypes.register_anonymous(np.dtype((np.float64, (5,))), "_RetDPrb5") + out2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_RetDPrb2") + + def _three(x): # pragma: no cover (numba) + return x[:3] + + # Three elements cannot fill a declared two-element output. + op_ = UnaryOp.register_anonymous(_three, "_ret_dtype_prb_bad", is_udt=True, ret_dtype=out2) + with pytest.raises(UdfParseError, match=r"shape \(3,\).*_RetDPrb2 elements are \(2,\)"): + op_[in5] + + # The check is fit-by-broadcast, not equality: a one-element return + # legitimately fills every slot of the declared element. + def _one(x): # pragma: no cover (numba) + return x[:1] + + ok = UnaryOp.register_anonymous(_one, "_ret_dtype_prb_bcast", is_udt=True, ret_dtype=out2) + assert ok[in5].return_type is out2 + + # The record half of the probe checks the declared type's array leaves. + rec = dtypes.register_anonymous( + np.dtype([("v", np.float64, (4,)), ("n", np.int64)], align=True), "_RetDPrbRec" + ) + + def _short_leaf(x, y): # pragma: no cover (numba) + return (x[:2], 1) + + recop = BinaryOp.register_anonymous( + _short_leaf, "_ret_dtype_prb_rec", is_udt=True, ret_dtype=rec + ) + with pytest.raises(UdfParseError, match=r"shape \(2,\) for field \['v'\] of _RetDPrbRec"): + recop[in5, in5] + + +def _ret_dtype_pickle_udf(x): # pragma: no cover (numba) + return x[:5] + + +_RET_DTYPE_PICKLE_WRITER = """ +import pickle +import sys + +import numpy as np + +import graphblas as gb +from graphblas.core.operator.unary import UnaryOp +from graphblas.tests.test_op import _ret_dtype_pickle_udf + +print("package: " + gb.__file__) +five = gb.dtypes.register_anonymous(np.dtype((np.float32, (5,))), "_RetPickle5") +UnaryOp.register_new("_ret_dtype_pickled", _ret_dtype_pickle_udf, is_udt=True, ret_dtype=five) +anon = UnaryOp.register_anonymous( + _ret_dtype_pickle_udf, "_ret_dtype_pickled_anon", is_udt=True, ret_dtype=five +) +with open(sys.argv[1], "wb") as f: + pickle.dump((gb.unary._ret_dtype_pickled, anon), f) +print("wrote: ok") +""" + + +_RET_DTYPE_PICKLE_READER = """ +import pickle +import sys + +import numpy as np + +import graphblas as gb + +print("package: " + gb.__file__) +with open(sys.argv[1], "rb") as f: + named, anon = pickle.load(f) +ten = gb.dtypes.register_anonymous(np.dtype((np.float32, (10,))), "_RetPickle10") +print("named shape: " + str(named[ten].return_type.np_type.subdtype[1])) +print("anon shape: " + str(anon[ten].return_type.np_type.subdtype[1])) +""" + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_ret_dtype_survives_pickle(tmp_path): + # In-process unpickling takes the _find shortcut and returns the very + # same object, so only a cross-process round trip exercises what pickle + # exists for: the reduce tuple must carry ret_dtype, or the reconstructed + # op re-registers without it and fails its own shape probe. Both sides + # run in subprocesses so the named registration never lands in this + # process's operator namespace (test_operator_types enumerates it). + payload = tmp_path / "ops.pkl" + repo_root = Path(__file__).resolve().parents[2] + env = dict(os.environ, PYTHONPATH=str(repo_root)) + for probe, checks in ( + (_RET_DTYPE_PICKLE_WRITER, {"wrote": "ok"}), + (_RET_DTYPE_PICKLE_READER, {"named shape": "(5,)", "anon shape": "(5,)"}), + ): + result = subprocess.run( + [sys.executable, "-c", probe, str(payload)], + capture_output=True, + text=True, + check=False, + env=env, + ) + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + lines = dict(line.split(": ", 1) for line in result.stdout.splitlines() if ": " in line) + assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report + for key, want in checks.items(): + assert lines.get(key) == want, report From 14495452ad280c9d5ca448858d7dc945fc766c3e Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 4 Aug 2026 16:57:31 -0700 Subject: [PATCH 02/20] Use unique array shapes for the ret_dtype test UDTs --- graphblas/tests/test_op.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 8ed3070d5..4bd281740 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -7397,7 +7397,10 @@ def test_udt_ret_dtype_errors(): @pytest.mark.slow def test_udt_ret_dtype_index_ops(): """``ret_dtype`` reaches the IndexUnaryOp and IndexBinaryOp compile paths too.""" - in5 = dtypes.register_anonymous(np.dtype((np.float64, (5,))), "_RetDIdx5") + # Shape (11,) is used nowhere else: anonymous UDTs share one DataType per + # np.dtype, so a shape reused across tests would inherit whichever JIT C + # state was frozen first under random test ordering. + in11 = dtypes.register_anonymous(np.dtype((np.float64, (11,))), "_RetDIdx11") out2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_RetDIdx2") def _head_plus_row(x, i, j, t): # pragma: no cover (numba) @@ -7406,7 +7409,7 @@ def _head_plus_row(x, i, j, t): # pragma: no cover (numba) iu = IndexUnaryOp.register_anonymous( _head_plus_row, "_ret_dtype_indexunary", is_udt=True, ret_dtype=out2 ) - assert iu[in5, INT64].return_type is out2 + assert iu[in11, INT64].return_type is out2 if lib.__dict__.get("GxB_IndexBinaryOp_new") is not None: from graphblas.core.operator import IndexBinaryOp @@ -7417,7 +7420,7 @@ def _head_sum(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) ib = IndexBinaryOp.register_anonymous( _head_sum, "_ret_dtype_indexbinary", is_udt=True, ret_dtype=out2 ) - assert ib[in5, in5].return_type is out2 + assert ib[in11, in11].return_type is out2 @pytest.mark.skipif("not supports_udfs") @@ -7431,7 +7434,10 @@ def test_udt_ret_dtype_still_shape_checked(): what it returned, while a declared type is an independent claim the UDF can contradict. """ - in5 = dtypes.register_anonymous(np.dtype((np.float64, (5,))), "_RetDPrb5") + # Shape (14,) is used nowhere else: anonymous UDTs share one DataType per + # np.dtype, so a shape reused across tests would inherit whichever JIT C + # state was frozen first under random test ordering. + in14 = dtypes.register_anonymous(np.dtype((np.float64, (14,))), "_RetDPrb14") out2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_RetDPrb2") def _three(x): # pragma: no cover (numba) @@ -7440,7 +7446,7 @@ def _three(x): # pragma: no cover (numba) # Three elements cannot fill a declared two-element output. op_ = UnaryOp.register_anonymous(_three, "_ret_dtype_prb_bad", is_udt=True, ret_dtype=out2) with pytest.raises(UdfParseError, match=r"shape \(3,\).*_RetDPrb2 elements are \(2,\)"): - op_[in5] + op_[in14] # The check is fit-by-broadcast, not equality: a one-element return # legitimately fills every slot of the declared element. @@ -7448,7 +7454,7 @@ def _one(x): # pragma: no cover (numba) return x[:1] ok = UnaryOp.register_anonymous(_one, "_ret_dtype_prb_bcast", is_udt=True, ret_dtype=out2) - assert ok[in5].return_type is out2 + assert ok[in14].return_type is out2 # The record half of the probe checks the declared type's array leaves. rec = dtypes.register_anonymous( @@ -7462,7 +7468,7 @@ def _short_leaf(x, y): # pragma: no cover (numba) _short_leaf, "_ret_dtype_prb_rec", is_udt=True, ret_dtype=rec ) with pytest.raises(UdfParseError, match=r"shape \(2,\) for field \['v'\] of _RetDPrbRec"): - recop[in5, in5] + recop[in14, in14] def _ret_dtype_pickle_udf(x): # pragma: no cover (numba) From 05ea2554014b0aeb20da54be98e6261718fdc60f Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 23 Sep 2026 12:37:32 -0500 Subject: [PATCH 03/20] Drop the stale separate-DataTypes note from _array_udt_operands The docstring said a flat FP64[2, 3] and a layered FP64[3][2] are separate DataTypes with separate GrB_Type handles that Numba collapses to one type. Since "Flatten nested subarray dtypes so each UDT layout has one DataType" (#589), registration flattens the layered form, so the two are one DataType. The dedup is unchanged; only its explanation was out of date. --- graphblas/core/operator/udf.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index 2a840eb41..992f7d933 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -193,10 +193,7 @@ def _summarize_numba_typing_error(exc): def _array_udt_operands(dtypes): """Return the distinct array-UDT operands in ``dtypes``, order preserved. - ``x`` and ``y`` are usually the same UDT, hence the dedup. It is by - identity, not by shape: a flat ``FP64[2, 3]`` and a layered - ``FP64[3][2]`` are separate DataTypes with separate GrB_Type handles, - yet Numba collapses both to the same ``nestedarray(float64, (2, 3))``. + ``x`` and ``y`` are usually the same UDT, hence the dedup. """ operands = [] for d in dtypes: From a884f18a8e14c5d6e7eef175f524f1b26842b2a2 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 24 Sep 2026 23:02:39 -0500 Subject: [PATCH 04/20] Check a declared ret_dtype against what the UDF returns A declared ret_dtype skipped inference, and nothing checked it in its place. Some mismatches were silent wrong answers: a record given a longer tuple dropped the extra values, a 2-element array given a 3-tuple was left uninitialized, and FP64 given an array kept its first element. The rest failed with a multi-line Numba error that never mentioned ret_dtype. _check_declared_ret_type accepts what inference would: one value per field (or a value of the record type) for a record, an array or a scalar for an array UDT, and a scalar otherwise. The field count it reports for a nested record is the flattened one, and the message says so. Also, register_new no longer adds an IndexUnaryOp declared to return something other than BOOL to the select namespace, and an op without ret_dtype pickles as the 3-tuple it always did, so older versions can still read it. Remove the unused _array_udt_operands and a repeated validation call, give the output-UDT test an input shape no other test uses, cover lazy=True with ret_dtype and an operand returned as-is into a declared shape, and have the user guide use gb.unary.register_anonymous and describe the new check. --- docs/user_guide/udt.rst | 17 ++-- graphblas/core/operator/base.py | 16 +-- graphblas/core/operator/binary.py | 1 - graphblas/core/operator/indexbinary.py | 1 - graphblas/core/operator/indexunary.py | 10 +- graphblas/core/operator/udf.py | 69 ++++++++++--- graphblas/core/operator/unary.py | 1 - graphblas/tests/test_op.py | 136 +++++++++++++++++++++++-- 8 files changed, 207 insertions(+), 44 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 7f247388d..4831b3ecd 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -381,14 +381,14 @@ The ``x[:2]`` rejection above is a case of this: the output really is a 2-element UDT, but nothing says so. Pass ``ret_dtype`` to say it outright. It takes anything ``lookup_dtype`` -accepts and requires ``is_udt=True``:: +accepts and requires ``is_udt=True``: + +.. code-block:: python nine = gb.dtypes.register_anonymous(np.dtype((np.float64, (9,))), "Nine") three = gb.dtypes.register_anonymous(np.dtype((np.float64, (3,))), "Three") - head = gb.core.operator.UnaryOp.register_anonymous( - lambda x: x[:3], "head", is_udt=True, ret_dtype=three - ) + head = gb.unary.register_anonymous(lambda x: x[:3], "head", is_udt=True, ret_dtype=three) head[nine].return_type # Three ``ret_dtype`` is available on ``register_anonymous`` and ``register_new`` for @@ -399,10 +399,11 @@ output type should vary with its inputs still needs one registration per output type. Declaring the type does not switch off the shape check described above; it -points the check at the declared type instead. A UDF whose result cannot -fill a ``ret_dtype`` element is still rejected when the op is typed, so a -wrong ``ret_dtype`` is a registration error rather than a silently -mis-typed result. +points the check at the declared type instead. The result must also be +something the declared type can hold: one value per field for a record, an +array or a scalar for an array UDT, and a scalar for a scalar type. A UDF that +breaks either rule is rejected when the op is typed, so a wrong ``ret_dtype`` +is a registration error rather than a silently mis-typed result. It does not apply to ``SelectOp``, whose return type GraphBLAS fixes at ``BOOL``, nor to builtin (non-UDT) dtypes, where the return type comes from diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index 1ea1d5474..0b539744d 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -614,9 +614,9 @@ def _deserialize_udf(cls, name, orig_func, is_udt, ret_dtype=None): Shared by the five UDF-capable subclasses (UnaryOp, BinaryOp, IndexUnaryOp, SelectOp, IndexBinaryOp), all of which use the - default ``__reduce__`` below. ``ret_dtype`` is passed only when set: - SelectOp shares this path and takes no ret_dtype, and pickles written - before ret_dtype existed carry a 3-tuple. + default ``__reduce__`` below. ``ret_dtype`` is passed on only when set: + SelectOp shares this path and takes no ret_dtype, and an op without one + pickles as the 3-tuple it always did. """ if (rv := cls._find(name)) is not None: return rv @@ -637,19 +637,19 @@ def __reduce__(self): ``Aggregator`` define their own ``__reduce__`` because their pickle shape differs (they hold a binary op + identity, etc.). """ + # ``ret_dtype`` rides along only when set, so an op without one pickles + # exactly as before and older versions can still read it. + ret = () if (ret_dtype := getattr(self, "_ret_dtype", None)) is None else (ret_dtype,) if self._anonymous: if hasattr(self.orig_func, "_parameterized_info"): return (_deserialize_parameterized, self.orig_func._parameterized_info) return ( type(self)._deserialize_anon_udf, - (self.orig_func, self.name, self._is_udt, getattr(self, "_ret_dtype", None)), + (self.orig_func, self.name, self._is_udt, *ret), ) if (name := f"{self._modname}.{self.name}") in _STANDARD_OPERATOR_NAMES: return name - return ( - type(self)._deserialize_udf, - (self.name, self.orig_func, self._is_udt, getattr(self, "_ret_dtype", None)), - ) + return (type(self)._deserialize_udf, (self.name, self.orig_func, self._is_udt, *ret)) @classmethod def _check_supports_udf(cls, method_name): diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 57535f5ef..2fb86e26d 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -1045,7 +1045,6 @@ def register_new( }, ) elif parameterized: - _validate_ret_dtype(ret_dtype, "binary", is_udt=is_udt, parameterized=True) binary_op = ParameterizedBinaryOp(name, func, is_udt=is_udt) setattr(module, funcname, binary_op) else: diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index 60cef2b4e..ed6146350 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -670,7 +670,6 @@ def register_new( }, ) elif parameterized: - _validate_ret_dtype(ret_dtype, "indexbinary", is_udt=is_udt, parameterized=True) idxbinop = ParameterizedIndexBinaryOp(name, func, is_udt=is_udt) setattr(module, funcname, idxbinop) else: diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index f3c6b57f6..77cc80678 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -361,14 +361,18 @@ def register_new( }, ) elif parameterized: - _validate_ret_dtype(ret_dtype, "indexunary", is_udt=is_udt, parameterized=True) indexunary_op = ParameterizedIndexUnaryOp(name, func, is_udt=is_udt) setattr(module, funcname, indexunary_op) else: indexunary_op = cls._build(name, func, is_udt=is_udt, ret_dtype=ret_dtype) setattr(module, funcname, indexunary_op) - # If return type is BOOL, register additionally as a SelectOp - if all(x == BOOL for x in indexunary_op.types.values()): + # If return type is BOOL, register additionally as a SelectOp. A UDT + # op has no types until it is compiled, so only a declared + # ``ret_dtype`` can rule it out here. + declared = indexunary_op._ret_dtype + if (declared is None or declared == BOOL) and all( + x == BOOL for x in indexunary_op.types.values() + ): from .select import SelectOp select_module, funcname = SelectOp._remove_nesting(name, strict=False) diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index 992f7d933..c7550cf64 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -190,29 +190,72 @@ def _summarize_numba_typing_error(exc): return "Numba could not compile the function for these input types" -def _array_udt_operands(dtypes): - """Return the distinct array-UDT operands in ``dtypes``, order preserved. - - ``x`` and ``y`` are usually the same UDT, hence the dedup. - """ - operands = [] - for d in dtypes: - if d._is_udt and d.np_type.subdtype is not None and not any(d is seen for seen in operands): - operands.append(d) - return operands - - def _udt_ret_type(parent_op, numba_ret_type, *dtypes): """Return the operator's declared ``ret_dtype``, else infer one from the UDF. Inference can only name a type that is already an operand, so an output - UDT that appears nowhere in the inputs is unreachable without this. + UDT that appears nowhere in the inputs is unreachable without this. A + declared type is checked against the UDF's return first. """ if (ret_dtype := parent_op._ret_dtype) is not None: + _check_declared_ret_type(ret_dtype, numba_ret_type) return ret_dtype return _resolve_udt_return_type(numba_ret_type, *dtypes) +def _check_declared_ret_type(ret_dtype, numba_ret_type): + """Reject a UDF whose return cannot be written as the declared ``ret_dtype``. + + Inference only names a type the return can be written as; a declared + type skips it. Left to the wrapper, a mismatch is a raw Numba error when + the wrapper compiles, or a silent wrong answer when it runs: a longer + tuple loses its extra values, an array written to a scalar keeps its + first element, and a tuple written to an array element leaves it + uninitialized. Shapes, and record leaves that can be ``None``, are + checked afterwards by :func:`_check_udf_fills_output`. + """ + nt = numba.core.types + np_type = ret_dtype.np_type + if ret_dtype._is_udt and np_type.names is not None: + nleaves = len(list(_iter_record_leaves(np_type))) + if isinstance(numba_ret_type, nt.BaseTuple): + if len(numba_ret_type.types) == nleaves: + return + # The tuple is flat over the leaves, so a nested record counts + # the fields of its inner records, not the inner records. + nested = "" if nleaves == len(np_type.names) else " once nested records are flattened" + raise UdfParseError( + f"UDT UDF returned a tuple of length {len(numba_ret_type.types)}, but " + f"ret_dtype={ret_dtype} has {nleaves} fields{nested}. " + f"Return one value per field." + ) + if numba_ret_type == ret_dtype.numba_type: + return # a ``ret_dtype`` value, such as an operand of that type + wanted = f"a tuple of its {nleaves} fields, or a {ret_dtype} value" + elif ret_dtype._is_udt and np_type.subdtype is not None: + if isinstance(numba_ret_type, (nt.Array, nt.Number, nt.Boolean)): + return # a scalar fills every slot; the shape is checked later + wanted = "an array, or a scalar to fill every slot" + else: + if isinstance(numba_ret_type, (nt.Number, nt.Boolean)): + return + wanted = "a scalar" + if isinstance(numba_ret_type, nt.BaseTuple): + returned = f"a tuple of length {len(numba_ret_type.types)}" + elif isinstance(numba_ret_type, nt.Array): + returned = "an array" + elif isinstance(numba_ret_type, nt.Record): + returned = f"a record with fields {list(numba_ret_type.fields)}" + elif isinstance(numba_ret_type, (nt.Optional, nt.NoneType)): + returned = "a value that can be None" + else: + returned = f"{numba_ret_type}" + raise UdfParseError( + f"UDT UDF returned {returned}, which cannot be written as " + f"ret_dtype={ret_dtype}. Return {wanted}." + ) + + def _resolve_udt_return_type(numba_ret_type, *dtypes): """Resolve a Numba return type to a DataType, matching Tuple returns to an input UDT. diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index a262c4dc9..29ffb02e8 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -417,7 +417,6 @@ def register_new( }, ) elif parameterized: - _validate_ret_dtype(ret_dtype, "unary", is_udt=is_udt, parameterized=True) unary_op = ParameterizedUnaryOp(name, func, is_udt=is_udt) setattr(module, funcname, unary_op) else: diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 4bd281740..1c781e423 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -7281,12 +7281,14 @@ def test_udt_ret_dtype_names_an_output_udt(): """``ret_dtype`` names an output UDT that is not one of the operands. Without it the return type is inferred from what the UDF builds, and the - only names in scope are the input dtypes. That makes a rank-reducing op - such as FP64[9] -> FP64[3] unreachable: the inferred type is the input's, - and the shape check then rejects the shorter array the UDF returns. + only names in scope are the input dtypes. That makes an op that shortens + its element, such as FP64[10] -> FP64[3], unreachable: the inferred type + is the input's, and the shape check then rejects the shorter array the + UDF returns. """ - # Shapes unique to this test; see the note in test_udt_array_udf_shape_errors. - nine = dtypes.register_anonymous(np.dtype((np.float64, (9,))), "_RetD9") + # The input shape is unique to this test; see the note in + # test_udt_array_udf_shape_errors. + ten = dtypes.register_anonymous(np.dtype((np.float64, (10,))), "_RetD10") three = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_RetD3") def _first_three(x): # pragma: no cover (numba) @@ -7294,13 +7296,25 @@ def _first_three(x): # pragma: no cover (numba) without = UnaryOp.register_anonymous(_first_three, "_ret_dtype_without", is_udt=True) with pytest.raises(UdfParseError, match=r"shape \(3,\) when run on sample values"): - without[nine] + without[ten] op_ = UnaryOp.register_anonymous(_first_three, "_ret_dtype_with", is_udt=True, ret_dtype=three) - assert op_[nine].return_type is three + assert op_[ten].return_type is three - v = Vector(nine, size=2) - v[0] = np.arange(9.0) + # lazy=True must hand ret_dtype to the registration it defers, or the op + # that fires on first attribute access is built without it. + UnaryOp.register_new("_ret_dtype_lazy", _first_three, is_udt=True, lazy=True, ret_dtype=three) + try: + assert unary._ret_dtype_lazy[ten].return_type is three + finally: + # Same cleanup as test_udt_lazy_registration: the op lands in ``unary`` + # and ``op``, and both keep a ``_delayed`` entry until it fires. + for module in (unary, op): + vars(module).pop("_ret_dtype_lazy", None) + module._delayed.pop("_ret_dtype_lazy", None) + + v = Vector(ten, size=2) + v[0] = np.arange(10.0) w = op_(v).new() assert w.dtype is three np.testing.assert_array_equal(w[0].new().value, np.arange(3.0)) @@ -7422,6 +7436,16 @@ def _head_sum(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) ) assert ib[in11, in11].return_type is out2 + # register_new also adds an IndexUnaryOp to ``select`` when it returns BOOL. + # A UDT op has no types until it is compiled, so a declared non-BOOL + # ret_dtype is what keeps this one out. + IndexUnaryOp.register_new("_ret_dtype_iu_named", _head_plus_row, is_udt=True, ret_dtype=out2) + try: + assert "_ret_dtype_iu_named" not in vars(select) + finally: + vars(indexunary).pop("_ret_dtype_iu_named", None) + vars(select).pop("_ret_dtype_iu_named", None) + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow @@ -7456,6 +7480,15 @@ def _one(x): # pragma: no cover (numba) ok = UnaryOp.register_anonymous(_one, "_ret_dtype_prb_bcast", is_udt=True, ret_dtype=out2) assert ok[in14].return_type is out2 + # An operand returned as-is keeps its extents in its Numba type, so the + # declared shape is checked without running the UDF (no "sample values"). + def _as_is(x): # pragma: no cover (numba) + return x + + as_is = UnaryOp.register_anonymous(_as_is, "_ret_dtype_prb_as_is", is_udt=True, ret_dtype=out2) + with pytest.raises(UdfParseError, match=r"shape \(14,\), but _RetDPrb2 elements are \(2,\)"): + as_is[in14] + # The record half of the probe checks the declared type's array leaves. rec = dtypes.register_anonymous( np.dtype([("v", np.float64, (4,)), ("n", np.int64)], align=True), "_RetDPrbRec" @@ -7471,6 +7504,86 @@ def _short_leaf(x, y): # pragma: no cover (numba) recop[in14, in14] +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_ret_dtype_rejects_a_return_it_cannot_hold(): + """A declared ``ret_dtype`` is checked against what the UDF returns. + + Inference only names a type the return can be written as, and a declared + type skips it. Unchecked, a mismatch reached the wrapper, which either + failed to compile with a raw Numba error or, worse, ran and gave a wrong + answer with no error: a longer tuple lost its extra values, an array + written to a scalar kept its first element, and a tuple written to an + array element left it uninitialized. + """ + # Names and shapes unique to this test; see the note in + # test_udt_array_udf_shape_errors. + rec2 = dtypes.register_anonymous( + np.dtype([("rdm_lo", np.float64), ("rdm_hi", np.float64)], align=True), "_RetDMisRec2" + ) + rec3 = dtypes.register_anonymous( + np.dtype([("rdm_a", np.int64), ("rdm_b", np.int64), ("rdm_c", np.int64)], align=True), + "_RetDMisRec3", + ) + arr2 = dtypes.register_anonymous(np.dtype((np.float32, (2,))), "_RetDMisArr2") + arr3 = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_RetDMisArr3") + + def binary_rejects(func, ret_dtype, match): + op_ = BinaryOp.register_anonymous(func, is_udt=True, ret_dtype=ret_dtype) + with pytest.raises(UdfParseError, match=match): + op_[FP64, FP64] + + binary_rejects( + lambda x, y: (x, y, x + y), rec2, "tuple of length 3, but ret_dtype=_RetDMisRec2 has 2" + ) + binary_rejects(lambda x, y: (x,), rec2, "tuple of length 1, but ret_dtype=_RetDMisRec2 has 2") + binary_rejects( + lambda x, y: x + y, rec2, "returned float64, which cannot be written as ret_dtype" + ) + + # A nested record takes a flat tuple, one value per leaf, and the count in + # the message says so when it differs from the number of top-level fields. + nested = dtypes.register_anonymous( + np.dtype( + [("rdm_in", [("rdm_p", np.float64), ("rdm_q", np.float64)]), ("rdm_r", np.float64)], + align=True, + ), + "_RetDMisNested", + ) + binary_rejects( + lambda x, y: ((x, y), x + y), + nested, + "tuple of length 2, but ret_dtype=_RetDMisNested has 3 fields once nested records", + ) + flat = BinaryOp.register_anonymous(lambda x, y: (x, y, x + y), is_udt=True, ret_dtype=nested) + assert flat[FP64, FP64].return_type is nested + binary_rejects(lambda x, y: (x, y, x + y), arr2, "tuple of length 3, which cannot be written") + binary_rejects(lambda x, y: (x, y), FP64, "tuple of length 2, which cannot be written") + + def unary_rejects(func, ret_dtype, operand, match): + op_ = UnaryOp.register_anonymous(func, is_udt=True, ret_dtype=ret_dtype) + with pytest.raises(UdfParseError, match=match): + op_[operand] + + unary_rejects(lambda x: x * 1.0, FP64, arr3, "returned an array, which cannot be written") + unary_rejects(lambda x: x, rec3, rec2, r"record with fields \['rdm_lo', 'rdm_hi'\]") + unary_rejects(lambda x: x, arr3, rec2, "record with fields") + unary_rejects(lambda x: x, rec2, arr3, "returned an array, which cannot be written") + + # What a declared type can hold still works: one value per field, a + # scalar that fills an array element, and a scalar cast to the declared type. + ok = BinaryOp.register_anonymous(lambda x, y: (x, y), is_udt=True, ret_dtype=rec2) + fill = BinaryOp.register_anonymous(lambda x, y: x + y, is_udt=True, ret_dtype=arr2) + cast = BinaryOp.register_anonymous(lambda x, y: x + y, is_udt=True, ret_dtype=INT64) + w = Vector(FP64, size=1) + w[0] = 3.0 + u = Vector(FP64, size=1) + u[0] = 1.0 + assert tuple(w.ewise_mult(u, ok).new()[0].new().value.tolist()) == (3.0, 1.0) + np.testing.assert_array_equal(w.ewise_mult(u, fill).new()[0].new().value, [4.0, 4.0]) + assert w.ewise_mult(u, cast).new()[0].new().value == 4 + + def _ret_dtype_pickle_udf(x): # pragma: no cover (numba) return x[:5] @@ -7542,3 +7655,8 @@ def test_udt_ret_dtype_survives_pickle(tmp_path): assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report for key, want in checks.items(): assert lines.get(key) == want, report + + # An op without ret_dtype pickles as the 3-tuple it always did, which + # versions without ret_dtype can still read. + plain = UnaryOp.register_anonymous(_ret_dtype_pickle_udf, "_ret_dtype_plain", is_udt=True) + assert len(plain.__reduce__()[1]) == 3 From 6b70607cffa6a7837803cac8e801c015791979a5 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Sun, 27 Sep 2026 10:03:43 -0500 Subject: [PATCH 05/20] Ignore the UDT repr size warning in the ret_dtype mismatch test _RetDMisNested's numpy repr is 141 characters, over the 128 a type name holds below SuiteSparse 9, so its registration warns there, and the test failed on psg 8 draws that ran slow tests on the suitesparse backend. It is about the ret_dtype check, not serialization. --- graphblas/tests/test_op.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 1c781e423..c1158e0ed 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -7506,6 +7506,8 @@ def _short_leaf(x, y): # pragma: no cover (numba) @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow +# _RetDMisNested's numpy repr is 141 chars; see test_udt_eq_nested_record_with_nan_leaf. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") def test_udt_ret_dtype_rejects_a_return_it_cannot_hold(): """A declared ``ret_dtype`` is checked against what the UDF returns. From f4cd16c507775ab3ee3e2b5f16b0f32a3c667224 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 07:54:48 -0500 Subject: [PATCH 06/20] Cover a declared record ret_dtype returned as is, and a return that can be None _check_declared_ret_type accepts a value of the declared record type, such as the operand itself, and names a return that can be None in its error, but no test reached either. Return the record operand from a unary op with that record as ret_dtype and check the result, and check the error for x if x > y else None under ret_dtype=FP64. --- graphblas/tests/test_op.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index c1158e0ed..8ab2e4eec 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -7561,6 +7561,9 @@ def binary_rejects(func, ret_dtype, match): assert flat[FP64, FP64].return_type is nested binary_rejects(lambda x, y: (x, y, x + y), arr2, "tuple of length 3, which cannot be written") binary_rejects(lambda x, y: (x, y), FP64, "tuple of length 2, which cannot be written") + binary_rejects( + lambda x, y: x if x > y else None, FP64, "returned a value that can be None, which cannot" + ) def unary_rejects(func, ret_dtype, operand, match): op_ = UnaryOp.register_anonymous(func, is_udt=True, ret_dtype=ret_dtype) @@ -7572,9 +7575,11 @@ def unary_rejects(func, ret_dtype, operand, match): unary_rejects(lambda x: x, arr3, rec2, "record with fields") unary_rejects(lambda x: x, rec2, arr3, "returned an array, which cannot be written") - # What a declared type can hold still works: one value per field, a - # scalar that fills an array element, and a scalar cast to the declared type. + # What a declared type can hold still works: one value per field, a value of + # the declared record type, a scalar that fills an array element, and a scalar + # cast to the declared type. ok = BinaryOp.register_anonymous(lambda x, y: (x, y), is_udt=True, ret_dtype=rec2) + same = UnaryOp.register_anonymous(lambda x: x, is_udt=True, ret_dtype=rec2) fill = BinaryOp.register_anonymous(lambda x, y: x + y, is_udt=True, ret_dtype=arr2) cast = BinaryOp.register_anonymous(lambda x, y: x + y, is_udt=True, ret_dtype=INT64) w = Vector(FP64, size=1) @@ -7582,6 +7587,9 @@ def unary_rejects(func, ret_dtype, operand, match): u = Vector(FP64, size=1) u[0] = 1.0 assert tuple(w.ewise_mult(u, ok).new()[0].new().value.tolist()) == (3.0, 1.0) + r = Vector(rec2, size=1) + r[0] = (1.0, 2.0) + assert tuple(r.apply(same).new()[0].new().value.tolist()) == (1.0, 2.0) np.testing.assert_array_equal(w.ewise_mult(u, fill).new()[0].new().value, [4.0, 4.0]) assert w.ewise_mult(u, cast).new()[0].new().value == 4 From 71e8d007b3c1b995c76d6f9277c9c351bfa19e2f Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 07:54:51 -0500 Subject: [PATCH 07/20] Test that unpickling a named UDF re-registers it with its ret_dtype _deserialize_udf registers a named op again when the name is missing, which only happened in the cross-process pickle test's subprocess, so coverage never saw it. Register a named op with a ret_dtype, check that unpickling returns the registered op, remove the name, and check that unpickling registers a new op that still returns the declared type. The name is removed again afterwards, as the cross-process tests in test_pickle.py do. --- graphblas/tests/test_op.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 8ab2e4eec..09fd07f7a 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1,5 +1,6 @@ import itertools import os +import pickle import subprocess import sys from pathlib import Path @@ -7670,3 +7671,26 @@ def test_udt_ret_dtype_survives_pickle(tmp_path): # versions without ret_dtype can still read. plain = UnaryOp.register_anonymous(_ret_dtype_pickle_udf, "_ret_dtype_plain", is_udt=True) assert len(plain.__reduce__()[1]) == 3 + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_ret_dtype_named_op_reregisters_on_unpickle(): + # A named op unpickled where its name is not registered is registered again, + # with its ret_dtype. The name is removed when done, since test_operator_types + # enumerates the operator namespaces. + five = dtypes.register_anonymous(np.dtype((np.float32, (5,))), "_RetPickle5") + twelve = dtypes.register_anonymous(np.dtype((np.float32, (12,))), "_RetPickle12") + name = "_ret_dtype_reregistered" + try: + named = UnaryOp.register_new(name, _ret_dtype_pickle_udf, is_udt=True, ret_dtype=five) + payload = pickle.dumps(named) + assert pickle.loads(payload) is named + delattr(unary, name) + vars(op).pop(name) + new = pickle.loads(payload) + assert new is not named + assert getattr(unary, name) is new + assert new[twelve].return_type is five + finally: + vars(unary).pop(name, None) + vars(op).pop(name, None) From 5c0ceaf035afd540abb7328aa6898ff8b7adc37c Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 13:58:17 -0800 Subject: [PATCH 08/20] Refuse a UDT UDF that raises for its operand types when the op is typed The shape probe ran the UDF on stand-in operands and gave up on any exception, so an op whose UDF cannot run on its operand types typed, and its cfunc raised for every element, where Numba prints and swallows the error: each element kept what SuiteSparse had in the buffer. iu = IndexUnaryOp.register_anonymous(lambda x, i, j, t: x[:2] * t, is_udt=True, ret_dtype=R2) v.apply(iu, 3.0) (v of a 6-element array UDT R6; the literal is an R6) before: [[0, 0], [0, 0], [0, 0]], and "Exception ignored ... ValueError: unable to broadcast argument 1 to output array" on stderr after: UdfParseError: UDT UDF raised ValueError when run on sample values of (R6, UINT64, UINT64, R6): unable to broadcast ... The same for a BinaryOp with right=3.0. A UDF that raises on the probe's ones is now run again on twos, since what it raises may depend on the values, and only an exception on both is refused; an ArithmeticError still leaves the op unchecked. test_udt_udf_shape_check_runs_udf_only_for_built_arrays keeps its UDF that raises for ones only, which now types from the twos. On #612 the same UDF without ret_dtype gave the same zeros. --- graphblas/core/operator/udf.py | 50 +++++++++++++++++++++------------- graphblas/tests/test_op.py | 35 +++++++++++++++++++++++- 2 files changed, 65 insertions(+), 20 deletions(-) diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index c7550cf64..e68da5343 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -539,8 +539,8 @@ def _fits_by_broadcast(shape, expected): return False -def _udf_probe_value(dtype): - """Build a stand-in operand of ``dtype`` for :func:`_run_udf_probe`. +def _udf_probe_value(dtype, fill=1): + """Build a stand-in operand of ``dtype`` for :func:`_run_udf_probe`, ``fill`` in every leaf. Only the shape of the UDF's result is read, so any values it accepts will do. @@ -548,20 +548,25 @@ def _udf_probe_value(dtype): if dtype._is_udt: np_type = dtype.np_type if np_type.subdtype is None: - return np.ones(1, dtype=np_type)[0] + return np.full(1, fill, dtype=np_type)[0] base_np_type, shape = np_type.subdtype - return np.ones(shape, dtype=base_np_type) + return np.full(shape, fill, dtype=base_np_type) return _sample_values[dtype] def _run_udf_probe(numba_func, arg_dtypes): - """Call the UDF on stand-in operands. Returns ``(result,)``, or ``None`` if it raised. - - A UDF that raises on the stand-in values is not checked, so one that - also returns a shape that does not fit is left to the wrapper, which - fails without telling the caller (see :func:`_compose_wrapper_body`). - Only the call is guarded, so a mistake in building the operands still - surfaces. + """Call the UDF on stand-in operands. Returns ``(result,)``, or ``None`` if it gave up. + + A UDF that raises on the stand-ins of ones is run again on twos, since + what it raises may depend on the values. One that raises an + ``ArithmeticError`` is then not checked, and if it also returns a shape + that does not fit, that is left to the wrapper, which fails without + telling the caller (see :func:`_compose_wrapper_body`). Any other + exception raised on both comes from the operands' types and shapes (a + broadcast, an index past the end): the cfunc would raise it for every + element, where Numba prints and swallows it and the element is never + written, so it raises ``UdfParseError`` here instead. Only the call is + guarded, so a mistake in building the operands still surfaces. This is where typing an op for a UDT runs the user's function rather than only compiling it, and ``OpBase.__contains__`` is a typed lookup, @@ -570,14 +575,21 @@ def _run_udf_probe(numba_func, arg_dtypes): array, a record, a scalar), so the call compiles the specialization the wrapper needs anyway. """ - args = [_udf_probe_value(d) for d in arg_dtypes] - try: - return (numba_func(*args),) - except NumbaError as exc: # pragma: no cover (safety) - # The wrapper's operands have these types, so it would fail to compile too. - raise UdfParseError(_summarize_numba_typing_error(exc)) from exc - except Exception: - return None + for fill in (1, 2): + args = [_udf_probe_value(d, fill) for d in arg_dtypes] + try: + return (numba_func(*args),) + except NumbaError as exc: # pragma: no cover (safety) + # The wrapper's operands have these types, so it would fail to compile too. + raise UdfParseError(_summarize_numba_typing_error(exc)) from exc + except ArithmeticError: + return None + except Exception as exc: + error = exc + names = ", ".join(str(dtype) for dtype in arg_dtypes) + raise UdfParseError( + f"UDT UDF raised {type(error).__name__} when run on sample values of ({names}): {error}" + ) from error def _raise_unless_fits(shape, expected, return_type, path, *, probed): diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 09fd07f7a..6edf518b2 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -3191,7 +3191,7 @@ def _sum(x, y): # pragma: no cover (numba) assert op[udt].return_type is udt assert probed == ["_sum"] - # A UDF that raises on the probe's values is not checked, and still types. + # A UDF that raises on the probe's ones is run again on twos, and types. def _no_ones(x, y): # pragma: no cover (numba) if x[0] == 1: raise ValueError("the probe's values") @@ -3202,6 +3202,39 @@ def _no_ones(x, y): # pragma: no cover (numba) assert probed == ["_sum", "_no_ones"] +@pytest.mark.skipif("not supports_udfs") +def test_udt_udf_raising_for_its_operand_types_raises_when_typed(): + """A UDF that raises on any values of its operand types is refused when the op is typed. + + The probe gave up on any exception, so the op typed, and its cfunc raised + for every element, where Numba prints and swallows the error: the result + held zeros. Here ``x[:2] * t`` cannot broadcast when ``t`` is a 6-element + array, as it is for an ``(R6, R6)`` op. + """ + r6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_RaiseTyped6") + r2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_RaiseTyped2") + iu = IndexUnaryOp.register_anonymous( + lambda x, i, j, t: x[:2] * t, "_raise_typed_iu", is_udt=True, ret_dtype=r2 + ) + bop = BinaryOp.register_anonymous( + lambda x, y: x[:2] * y, "_raise_typed_b", is_udt=True, ret_dtype=r2 + ) + for parent, types in [ + (iu, "_RaiseTyped6, UINT64, UINT64, _RaiseTyped6"), + (bop, "_RaiseTyped6, _RaiseTyped6"), + ]: + with pytest.raises( + UdfParseError, match=rf"raised ValueError when run on sample values of \({types}\)" + ): + parent[r6, r6] + # With a scalar operand the op types and runs + v = Vector(r6, size=2) + v[0] = np.arange(6.0) + three = gb.Scalar.from_value(3.0) + assert v.apply(iu, three).new()[0].new().value.tolist() == [0.0, 3.0] + assert v.apply(bop, right=three).new()[0].new().value.tolist() == [0.0, 3.0] + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_record_array_field_roundtrip(): From adc32fa7bd3010a52ba6f38e76afc32a13133d9d Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 14:00:37 -0800 Subject: [PATCH 09/20] Type a number as itself beside a user's op that cannot take it as an element A literal beside a user's is_udt op (an untyped BinaryOp, IndexUnaryOp or SelectOp) becomes an element of the operand's UDT. A narrowing op with ret_dtype that uses the number, as x[:2] * t does, cannot take a 6-element t, so a literal failed where a Scalar of it works: iu = IndexUnaryOp.register_anonymous(lambda x, i, j, t: x[:2] * t, is_udt=True, ret_dtype=R2) v.apply(iu, 3.0), v.apply(iu, np.float64(3.0)), and a BinaryOp's right=3.0 or left=3.0 (v of a 6-element array UDT R6) on #613: [[0, 0], [0, 0], [0, 0]], silently on the previous commit: UdfParseError (the UDF raises on R6, R6) after: [[0, 3], [30, 33], [60, 63]], as with Scalar.from_value(3.0) _own_scalar_type_for_user_op types a Python or numpy number as itself (FP64 for 3.0, INT64 for 3, FP32 for np.float32) when the op cannot be typed with an element and can be typed with the number; otherwise it stays an element, so an op that takes an element (x + t on INT8[3], typed INT8[3] as before), and the fit check (0.5 for INT8[3] raises), are unchanged. A record op that adds a number to each field, which raised before, now works. Built-in ops, typed ops and Monoids are untouched, as is an int beyond INT64. udt.rst's literal row says so. --- docs/user_guide/udt.rst | 5 +++- graphblas/core/literals.py | 36 ++++++++++++++++++++++++++ graphblas/tests/test_op.py | 53 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 93 insertions(+), 1 deletion(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 4831b3ecd..25aea9935 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -209,7 +209,10 @@ dtype, so it never narrows. Below, ``v`` is a Vector of an ``INT8[3]`` UDT. but converts exactly, such as ``0`` for a bool field or ``2.0`` for an int field, still converts with a ``DeprecationWarning`` and will raise in a future version; any other raises ``ValueError``, and a float too - large for a float field raises ``OverflowError``. + large for a float field raises ``OverflowError``. A Python or numpy + number beside a user-defined op that cannot be typed with an element + (``x[:2] * t`` with a 6-element ``t``) is typed as itself instead, as a + ``Scalar`` of it would be. - ``v.ewise_union(w, binary.plus, 0.5, 0)`` raises A Monoid used element-wise (``apply``, ``ewise_mult``, ``ewise_union``) uses diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index d43c97aa8..87a36ab42 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -127,10 +127,46 @@ def _literal_operand(dtype, value, op, *, left=False): literal_type = _literal_type(dtype, value, exact=exact) if literal_type is not None: return value, literal_type + if (own := _own_scalar_type_for_user_op(dtype, value, op, left)) is not None: + return value, own _check_literal_fits(dtype, value) return value, dtype +def _own_scalar_type_for_user_op(dtype, value, op, left): + """Return a number's own type, for a user's op that cannot take it as an element, or None. + + Beside a user's ``is_udt`` op (an untyped BinaryOp, IndexUnaryOp or + SelectOp), a literal becomes an element of the UDT. A Python or numpy + number for an op whose function cannot be typed with that element (it + fails to compile, or raises on sample values: ``x[:2] * t`` with a + 6-element ``t``) is typed as itself instead, as a ``Scalar`` of it would + be, when the op can be typed so. Otherwise it stays an element, and the + op reports why it cannot take it. + """ + from ..exceptions import UdfParseError + + if ( + not getattr(op, "_is_udt", False) + or isinstance(op, (TypedOpBase, Monoid)) + or not ( + type(value) in {bool, int, float, complex} or isinstance(value, (np.number, np.bool_)) + ) + or (type(value) is int and not -(2**63) <= value < 2**63) # no INT64 holds it + or (own := lookup_dtype(type(value), value))._is_udt + ): + return None + + def takes(dtype2): + try: + op[(dtype2, dtype) if left else (dtype, dtype2)] + except UdfParseError: + return False + return True + + return own if not takes(dtype) and takes(own) else None + + # Built-in ops that compare their operands, and so keep no result type. _COMPARISONS = {"eq", "ne", "lt", "le", "gt", "ge"} _COMPARISONS.update([f"value{name}" for name in list(_COMPARISONS)]) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 6edf518b2..2a1df5273 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -3235,6 +3235,59 @@ def test_udt_udf_raising_for_its_operand_types_raises_when_typed(): assert v.apply(bop, right=three).new()[0].new().value.tolist() == [0.0, 3.0] +@pytest.mark.skipif("not supports_udfs") +def test_udt_user_op_takes_a_number_as_itself_when_it_cannot_take_an_element(): + """A number beside a user's op is typed as itself when the op cannot take it as an element. + + A literal beside a user's ``is_udt`` op becomes an element of the operand's + UDT, so ``x[:2] * t`` met a 6-element ``t``: its cfunc raised for every + element and the result held zeros, where a Scalar thunk worked. When the op + cannot be typed with an element, a Python or numpy number now takes its own + type, as a Scalar of it would. + """ + r6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_OwnNum6") + r2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_OwnNum2") + iu = IndexUnaryOp.register_anonymous( + lambda x, i, j, t: x[:2] * t, "_own_num_iu", is_udt=True, ret_dtype=r2 + ) + bop = BinaryOp.register_anonymous( + lambda x, y: x[:2] * y, "_own_num_b", is_udt=True, ret_dtype=r2 + ) + left = BinaryOp.register_anonymous( + lambda x, y: y[:2] * x, "_own_num_left", is_udt=True, ret_dtype=r2 + ) + v = Vector(r6, size=2) + v[0] = np.arange(6.0) + v[1] = np.arange(6.0) + 10 + expected = [[0.0, 3.0], [30.0, 33.0]] + for literal, own in [(3.0, FP64), (np.float64(3.0), FP64), (3, INT64), (np.float32(3), FP32)]: + for expr in [ + v.apply(iu, literal), + v.apply(bop, right=literal), + v.apply(left, left=literal), + ]: + assert own in {expr.op.type, expr.op._type2} + assert expr.new().to_coo()[1].tolist() == expected + # An op that can take an element still gets one, as udt.rst says + i3 = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_OwnNumI3") + add = IndexUnaryOp.register_anonymous(lambda x, i, j, t: x + t, "_own_num_add", is_udt=True) + w = Vector(i3, size=1) + w[0] = [1, 2, 3] + expr = w.apply(add, 1) + assert expr.op._type2 is i3 + assert expr.new()[0].new().value.tolist() == [2, 3, 4] + with pytest.raises(ValueError, match="does not fit _OwnNumI3"): + w.apply(add, 0.5) + # A record op that adds a number to each field takes it as itself + rec = dtypes.register_anonymous(np.dtype([("on_a", np.int32), ("on_b", np.float64)])) + each = BinaryOp.register_anonymous( + lambda x, y: (x.on_a + y, x.on_b + y), "_own_num_each", is_udt=True + ) + r = Vector(rec, size=1) + r[0] = (1, 2.0) + assert r.apply(each, right=2).new()[0].new().value.tolist() == (3, 4.0) + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_record_array_field_roundtrip(): From ba2a00a107f983f0ff0b83c5966998f2ee163cb5 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 14:08:12 -0800 Subject: [PATCH 10/20] Let ret_dtype name the operand's own date, duration or text UDT _check_declared_ret_type's scalar branch took only numbers and bools, while inference names a date, duration or bytes UDT for an operand returned as is, so declaring the type inference picks broke an op that works: dt = register_anonymous(np.dtype("M8[s]")) (also m8[s], S5) UnaryOp.register_anonymous(lambda x: x, is_udt=True, ret_dtype=dt)[dt] before: UdfParseError: UDT UDF returned datetime64[s], which cannot be written as ret_dtype=... Return a scalar. after: typed, with return type dt A value of the declared UDT's own Numba type is accepted, as the record branch accepts a record of the declared type. --- graphblas/core/operator/udf.py | 2 ++ graphblas/tests/test_op.py | 13 +++++++++++++ 2 files changed, 15 insertions(+) diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index e68da5343..c256ca2fc 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -239,6 +239,8 @@ def _check_declared_ret_type(ret_dtype, numba_ret_type): else: if isinstance(numba_ret_type, (nt.Number, nt.Boolean)): return + if ret_dtype._is_udt and numba_ret_type == ret_dtype.numba_type: + return # a value of a date, duration or text UDT, such as an operand returned as is wanted = "a scalar" if isinstance(numba_ret_type, nt.BaseTuple): returned = f"a tuple of length {len(numba_ret_type.types)}" diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 2a1df5273..12a7071f1 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -7679,6 +7679,19 @@ def unary_rejects(func, ret_dtype, operand, match): assert tuple(r.apply(same).new()[0].new().value.tolist()) == (1.0, 2.0) np.testing.assert_array_equal(w.ewise_mult(u, fill).new()[0].new().value, [4.0, 4.0]) assert w.ewise_mult(u, cast).new()[0].new().value == 4 + # So does an operand of a date, duration or text UDT returned as its own type, + # which said "Return a scalar" + for np_type, value in [ + ("M8[s]", np.datetime64(5, "s")), + ("m8[s]", np.timedelta64(5, "s")), + ("S5", b"abc"), + ]: + udt = dtypes.register_anonymous(np.dtype(np_type)) + own = UnaryOp.register_anonymous(lambda x: x, is_udt=True, ret_dtype=udt) + assert own[udt].return_type is udt + v = Vector(udt, size=1) + v[0] = value + assert v.apply(own).new()[0].new().value == value def _ret_dtype_pickle_udf(x): # pragma: no cover (numba) From a294ebda014729e082e47fa2010a59bccb64f88b Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 14:09:01 -0800 Subject: [PATCH 11/20] Say that a complex return cannot be written as a real ret_dtype Numba has no cast from complex to real, so the wrapper failed to compile and _summarize_numba_typing_error kept only Numba's first line: BinaryOp.register_anonymous(lambda x, y: x * 1j + y, is_udt=True, ret_dtype=FP64)[FP64, FP64] before: UdfParseError: binary. does not work with (FP64, FP64): Failed in nopython mode pipeline (step: native lowering) after: UdfParseError: UDT UDF returned complex128, which cannot be written as ret_dtype=FP64. Return a real value, or declare a complex ret_dtype. The same for a complex scalar or array into a real array UDT, which said "No implementation of function ... setitem". --- graphblas/core/operator/udf.py | 17 +++++++++++++++++ graphblas/tests/test_op.py | 3 +++ 2 files changed, 20 insertions(+) diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index c256ca2fc..78e71d379 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -233,10 +233,12 @@ def _check_declared_ret_type(ret_dtype, numba_ret_type): return # a ``ret_dtype`` value, such as an operand of that type wanted = f"a tuple of its {nleaves} fields, or a {ret_dtype} value" elif ret_dtype._is_udt and np_type.subdtype is not None: + _refuse_complex_into_real(ret_dtype, np_type.subdtype[0], numba_ret_type) if isinstance(numba_ret_type, (nt.Array, nt.Number, nt.Boolean)): return # a scalar fills every slot; the shape is checked later wanted = "an array, or a scalar to fill every slot" else: + _refuse_complex_into_real(ret_dtype, np_type, numba_ret_type) if isinstance(numba_ret_type, (nt.Number, nt.Boolean)): return if ret_dtype._is_udt and numba_ret_type == ret_dtype.numba_type: @@ -258,6 +260,21 @@ def _check_declared_ret_type(ret_dtype, numba_ret_type): ) +def _refuse_complex_into_real(ret_dtype, element, numba_ret_type): + """Raise unless a complex return goes into complex ``element``s of ``ret_dtype``. + + Numba has no cast from complex to real, so the wrapper fails to compile, + with a message that loses why (``Failed in nopython mode pipeline``). + """ + nt = numba.core.types + returned = numba_ret_type.dtype if isinstance(numba_ret_type, nt.Array) else numba_ret_type + if isinstance(returned, nt.Complex) and element.kind != "c": + raise UdfParseError( + f"UDT UDF returned {numba_ret_type}, which cannot be written as " + f"ret_dtype={ret_dtype}. Return a real value, or declare a complex ret_dtype." + ) + + def _resolve_udt_return_type(numba_ret_type, *dtypes): """Resolve a Numba return type to a DataType, matching Tuple returns to an input UDT. diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 12a7071f1..8e716b980 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -7648,6 +7648,9 @@ def binary_rejects(func, ret_dtype, match): assert flat[FP64, FP64].return_type is nested binary_rejects(lambda x, y: (x, y, x + y), arr2, "tuple of length 3, which cannot be written") binary_rejects(lambda x, y: (x, y), FP64, "tuple of length 2, which cannot be written") + # A complex return into real values; Numba's own message lost why + for ret_dtype in [FP64, arr2]: + binary_rejects(lambda x, y: x * 1j + y, ret_dtype, "complex128, .* Return a real value") binary_rejects( lambda x, y: x if x > y else None, FP64, "returned a value that can be None, which cannot" ) From dceefdfbe54a67000c3ed1d5a15d0f32760460c1 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 14:09:24 -0800 Subject: [PATCH 12/20] Say when a wrong ret_dtype raises, and what it and inference cover udt.rst called a wrong ret_dtype "a registration error"; registering succeeds, and the error comes when the op is typed (op[dtype], or its first use), as the sentence before says. It also said ret_dtype does not apply to "builtin (non-UDT) dtypes", but it does apply to built-in input dtypes and accepts a built-in return type; what it excludes is an op registered without is_udt=True. The four register_anonymous docstrings (and _udt_ret_type's) said inference "can only name a type that is already an input"; it also names a built-in dtype, such as FP64 from a record. --- docs/user_guide/udt.rst | 9 +++++---- graphblas/core/operator/binary.py | 4 ++-- graphblas/core/operator/indexbinary.py | 4 ++-- graphblas/core/operator/indexunary.py | 4 ++-- graphblas/core/operator/udf.py | 4 ++-- graphblas/core/operator/unary.py | 4 ++-- 6 files changed, 15 insertions(+), 14 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 25aea9935..0ef6ea576 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -405,12 +405,13 @@ Declaring the type does not switch off the shape check described above; it points the check at the declared type instead. The result must also be something the declared type can hold: one value per field for a record, an array or a scalar for an array UDT, and a scalar for a scalar type. A UDF that -breaks either rule is rejected when the op is typed, so a wrong ``ret_dtype`` -is a registration error rather than a silently mis-typed result. +breaks either rule is rejected when the op is typed (``op[dtype]``, or its +first use), so a wrong ``ret_dtype`` is an error rather than a silently +mis-typed result. It does not apply to ``SelectOp``, whose return type GraphBLAS fixes at -``BOOL``, nor to builtin (non-UDT) dtypes, where the return type comes from -compiling the function against each input type in turn. +``BOOL``, nor to an op registered without ``is_udt=True``, whose return type +comes from compiling the function against each input type in turn. .. _udt_jit_introspection: diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 2fb86e26d..9ba0818ef 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -935,8 +935,8 @@ def register_anonymous( ret_dtype : dtype, optional The dtype the operator returns. Requires ``is_udt=True``. Without it the return type is inferred from what the function - returns, which can only name a type that is already an input, so - an output UDT that is not an operand needs this. The dtype is + returns, which names a built-in dtype or one of the input UDTs, + so an output UDT that is not an operand needs this. The dtype is fixed for the operator: it is the same for every input dtype. Returns diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index ed6146350..91b31a812 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -606,8 +606,8 @@ def register_anonymous( ret_dtype : dtype, optional The dtype the operator returns. Requires ``is_udt=True``. Without it the return type is inferred from what the function - returns, which can only name a type that is already an input, so - an output UDT that is not an operand needs this. The dtype is + returns, which names a built-in dtype or one of the input UDTs, + so an output UDT that is not an operand needs this. The dtype is fixed for the operator: it is the same for every input dtype. Returns diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index 77cc80678..2bab291bd 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -272,8 +272,8 @@ def register_anonymous( ret_dtype : dtype, optional The dtype the operator returns. Requires ``is_udt=True``. Without it the return type is inferred from what the function - returns, which can only name a type that is already an input, so - an output UDT that is not an operand needs this. The dtype is + returns, which names a built-in dtype or one of the input UDTs, + so an output UDT that is not an operand needs this. The dtype is fixed for the operator: it is the same for every input dtype. Returns diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index 78e71d379..7ae1ed459 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -193,8 +193,8 @@ def _summarize_numba_typing_error(exc): def _udt_ret_type(parent_op, numba_ret_type, *dtypes): """Return the operator's declared ``ret_dtype``, else infer one from the UDF. - Inference can only name a type that is already an operand, so an output - UDT that appears nowhere in the inputs is unreachable without this. A + Inference names a built-in dtype or an operand's type, so an output UDT + that appears nowhere in the inputs is unreachable without this. A declared type is checked against the UDF's return first. """ if (ret_dtype := parent_op._ret_dtype) is not None: diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index 29ffb02e8..de4424f6e 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -338,8 +338,8 @@ def register_anonymous( ret_dtype : dtype, optional The dtype the operator returns. Requires ``is_udt=True``. Without it the return type is inferred from what the function - returns, which can only name a type that is already an input, so - an output UDT that is not an operand needs this. The dtype is + returns, which names a built-in dtype or one of the input UDTs, + so an output UDT that is not an operand needs this. The dtype is fixed for the operator: it is the same for every input dtype. Returns From c04bd832a02f651daa25cccb332e7a9d5684ed18 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 14:10:53 -0800 Subject: [PATCH 13/20] Test ret_dtype with lazy and parameterized registration for every op class Only UnaryOp was tested, so dropping ret_dtype from BinaryOp's or IndexUnaryOp's deferred (lazy=True) registration, or BinaryOp's parameterized check, left every test passing. The new test runs both over UnaryOp, BinaryOp, IndexUnaryOp and, where SuiteSparse has it, IndexBinaryOp; each of the three mutations now fails it. test_udt_ret_dtype_index_ops also runs its IndexUnaryOp and IndexBinaryOp on a Vector and checks the values; it only checked return_type. --- graphblas/tests/test_op.py | 65 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 8e716b980..f376c43ad 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -7494,6 +7494,62 @@ def test_udt_ret_dtype_errors(): ) +def _ret_dtype_class_cases(): + """(class, namespace, UDF giving the first 3 of x, typed lookup) per class with ret_dtype.""" + cases = [ + (UnaryOp, "unary", lambda x: x[:3], lambda o, t: o[t]), + (BinaryOp, "binary", lambda x, y: x[:3] + y[:3], lambda o, t: o[t, t]), + (IndexUnaryOp, "indexunary", lambda x, i, j, th: x[:3] + i, lambda o, t: o[t, INT64]), + ] + if lib.__dict__.get("GxB_IndexBinaryOp_new") is not None: + from graphblas.core.operator import IndexBinaryOp + + cases.append( + ( + IndexBinaryOp, + "indexbinary", + lambda x, ix, jx, y, iy, jy, th: x[:3] + y[:3], + lambda o, t: o[t, t], + ) + ) + return cases + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_ret_dtype_lazy_and_parameterized_for_each_op_class(): + """``lazy=True`` keeps ``ret_dtype``, and ``parameterized=True`` refuses it, for every class. + + Only UnaryOp was tested, so dropping ``ret_dtype`` from BinaryOp's or + IndexUnaryOp's deferred registration, or BinaryOp's parameterized check, + left every test passing. + """ + twelve = dtypes.register_anonymous(np.dtype((np.float64, (12,))), "_RetDCls12") + three = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_RetDCls3") + for cls, modname, func, typed in _ret_dtype_class_cases(): + module = getattr(gb, modname) + name = f"_ret_dtype_lazy_{modname}" + cls.register_new(name, func, is_udt=True, lazy=True, ret_dtype=three) + try: + assert typed(getattr(module, name), twelve).return_type is three + finally: + for ns in (module, op): + vars(ns).pop(name, None) + ns._delayed.pop(name, None) + with pytest.raises(ValueError, match="does not work with parameterized=True"): + cls.register_anonymous( + lambda t=1, func=func: func, is_udt=True, parameterized=True, ret_dtype=three + ) + with pytest.raises(ValueError, match="parameterized=True"): + cls.register_new( + f"{name}_param", + lambda t=1, func=func: func, + is_udt=True, + parameterized=True, + lazy=True, + ret_dtype=three, + ) + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_ret_dtype_index_ops(): @@ -7511,6 +7567,12 @@ def _head_plus_row(x, i, j, t): # pragma: no cover (numba) _head_plus_row, "_ret_dtype_indexunary", is_udt=True, ret_dtype=out2 ) assert iu[in11, INT64].return_type is out2 + v = Vector(in11, size=3) + v[1] = np.arange(11.0) + v[2] = np.arange(11.0) + 10 + result = v.apply(iu, gb.Scalar.from_value(0)).new() + assert result.dtype is out2 + assert result.to_coo()[1].tolist() == [[1.0, 2.0], [12.0, 13.0]] if lib.__dict__.get("GxB_IndexBinaryOp_new") is not None: from graphblas.core.operator import IndexBinaryOp @@ -7522,6 +7584,9 @@ def _head_sum(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) _head_sum, "_ret_dtype_indexbinary", is_udt=True, ret_dtype=out2 ) assert ib[in11, in11].return_type is out2 + result = v.ewise_mult(v, ib(0.0)).new() + assert result.dtype is out2 + assert result.to_coo()[1].tolist() == [[0.0, 2.0], [20.0, 22.0]] # register_new also adds an IndexUnaryOp to ``select`` when it returns BOOL. # A UDT op has no types until it is compiled, so a declared non-BOOL From 30020cb812e8fc051213b093c1234b17f9fac8c7 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 15:10:34 -0800 Subject: [PATCH 14/20] Type a UDF that raises on stand-ins as before; only the literal rule reads it "Refuse a UDT UDF that raises for its operand types when the op is typed" refused a UDF that raised a non-arithmetic exception on stand-in operands of ones and of twos. Such stand-ins are degenerate, and the index and thunk stand-ins do not vary, so valid UDFs that work on real data were refused (r612): UnaryOp(np.linalg.inv) on FP64[2, 2] (a matrix of ones is singular) UnaryOp: if x[0] >= x[1]: raise ...; x * 2 (equal fields) IndexUnaryOp: x.reshape((int(t), -1)).sum(0), and if t < 1.0: raise ...; x[:2] * t (the thunk stand-in is 0.5) on 542713c1: UdfParseError when typed after: typed and run, as at #613's head _run_udf_probe now only reports: it returns (ran, error), where error is the exception raised on both stand-ins, and typing never raises for it. _finalize_udt_op records such an op in _RAISES_ON_SAMPLES. The literal rule reads that record: beside an op whose UDF raised on every stand-in element, a number is typed as itself if the op types with the number's type (its stand-in is a fixed sample, not the number, so it need not pass the probe). So v.apply(iu, 3.0) still gives [0, 3], and so does the thunk check above with a literal 3.0. test_udt_udf_raising_for_its_operand_types_raises_when_typed becomes test_udt_udf_raising_on_stand_ins_still_types, with r612's four UDFs. --- graphblas/core/literals.py | 24 +++++++---- graphblas/core/operator/udf.py | 57 ++++++++++++++++--------- graphblas/tests/test_op.py | 77 +++++++++++++++++++++------------- 3 files changed, 100 insertions(+), 58 deletions(-) diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index 87a36ab42..b9255093e 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -139,10 +139,10 @@ def _own_scalar_type_for_user_op(dtype, value, op, left): Beside a user's ``is_udt`` op (an untyped BinaryOp, IndexUnaryOp or SelectOp), a literal becomes an element of the UDT. A Python or numpy number for an op whose function cannot be typed with that element (it - fails to compile, or raises on sample values: ``x[:2] * t`` with a - 6-element ``t``) is typed as itself instead, as a ``Scalar`` of it would - be, when the op can be typed so. Otherwise it stays an element, and the - op reports why it cannot take it. + fails to compile, or raises on every set of sample values: ``x[:2] * t`` + with a 6-element ``t``) is typed as itself instead, as a ``Scalar`` of it + would be, when the op can be typed so. Otherwise it stays an element, + and the op reports why it cannot take it. """ from ..exceptions import UdfParseError @@ -157,14 +157,20 @@ def _own_scalar_type_for_user_op(dtype, value, op, left): ): return None - def takes(dtype2): + from .operator.udf import _RAISES_ON_SAMPLES + + def types(dtype2): try: - op[(dtype2, dtype) if left else (dtype, dtype2)] + return op[(dtype2, dtype) if left else (dtype, dtype2)] except UdfParseError: - return False - return True + return None - return own if not takes(dtype) and takes(own) else None + # A UDF that raised on every stand-in element types, but cannot take an element. + # The number's own type need only type: its stand-in is a fixed sample, not the number. + element = types(dtype) + if element is not None and element not in _RAISES_ON_SAMPLES or types(own) is None: + return None + return own # Built-in ops that compare their operands, and so keep no result type. diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index 7ae1ed459..fa32a6122 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -19,6 +19,7 @@ import re import warnings +import weakref import numba import numpy as np @@ -100,6 +101,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``. """ + raises_on_samples = getattr(wrapper, "_raises_on_samples", None) try: wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) except NumbaError as exc: @@ -130,6 +132,8 @@ def _finalize_udt_op(parent_op, dtype, dtype2, ret_type, wrapper, wrapper_sig, t key = (dtype, dtype2) parent_op._udt_types[key] = ret_type parent_op._udt_ops[key] = op + if raises_on_samples is not None: + _RAISES_ON_SAMPLES[op] = raises_on_samples return op @@ -574,18 +578,17 @@ def _udf_probe_value(dtype, fill=1): def _run_udf_probe(numba_func, arg_dtypes): - """Call the UDF on stand-in operands. Returns ``(result,)``, or ``None`` if it gave up. - - A UDF that raises on the stand-ins of ones is run again on twos, since - what it raises may depend on the values. One that raises an - ``ArithmeticError`` is then not checked, and if it also returns a shape - that does not fit, that is left to the wrapper, which fails without - telling the caller (see :func:`_compose_wrapper_body`). Any other - exception raised on both comes from the operands' types and shapes (a - broadcast, an index past the end): the cfunc would raise it for every - element, where Numba prints and swallows it and the element is never - written, so it raises ``UdfParseError`` here instead. Only the call is - guarded, so a mistake in building the operands still surfaces. + """Call the UDF on stand-in operands. Returns ``(ran, error)``. + + ``ran`` is ``(result,)``, or ``None`` if the UDF raised on the stand-ins + of ones: then its result is not checked, and if it also returns a shape + that does not fit, that is left to the wrapper, whose output then keeps + its zeros (see :func:`_compose_wrapper_body`). Such a UDF may raise only + for some values, so it still types (``np.linalg.inv`` raises on a matrix + of ones). ``error`` is the exception when it raised on twos as well, and + it was no ``ArithmeticError``; only the literal rule reads it (see + :data:`_RAISES_ON_SAMPLES`). Only the call is guarded, so a mistake in + building the operands still surfaces. This is where typing an op for a UDT runs the user's function rather than only compiling it, and ``OpBase.__contains__`` is a typed lookup, @@ -594,21 +597,28 @@ def _run_udf_probe(numba_func, arg_dtypes): array, a record, a scalar), so the call compiles the specialization the wrapper needs anyway. """ + error = None for fill in (1, 2): args = [_udf_probe_value(d, fill) for d in arg_dtypes] try: - return (numba_func(*args),) + result = numba_func(*args) except NumbaError as exc: # pragma: no cover (safety) # The wrapper's operands have these types, so it would fail to compile too. raise UdfParseError(_summarize_numba_typing_error(exc)) from exc except ArithmeticError: - return None + return None, None except Exception as exc: error = exc - names = ", ".join(str(dtype) for dtype in arg_dtypes) - raise UdfParseError( - f"UDT UDF raised {type(error).__name__} when run on sample values of ({names}): {error}" - ) from error + continue + return ((result,) if fill == 1 else None), None + return None, error + + +# The UDT ops whose UDF raised, other than an ArithmeticError, on both sets of stand-in +# operands, with the exception. Such an op still types, since the UDF may raise only +# for some values, but beside it a number is typed as itself rather than as an element +# (``literals._own_scalar_type_for_user_op``), as a Scalar of it would be. +_RAISES_ON_SAMPLES = weakref.WeakKeyDictionary() def _raise_unless_fits(shape, expected, return_type, path, *, probed): @@ -686,6 +696,8 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ length exactly (:func:`_raise_unless_sequence_fits`); a tuple's length is in its type, and a list's is learned the same way as a built array's. A UDF that returns no built array or list is never run. + Returns the exception the UDF raised on every stand-in (see + :func:`_run_udf_probe`), or None. """ if zkind == "array_elements": # (record field path, Numba type, slot shape, index into the tuple) @@ -731,7 +743,10 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ ) elif isinstance(nb_type, numba.core.types.List): unknown.append((path, expected, i)) # a list's length is not in its type - if unknown and (ran := _run_udf_probe(numba_func, arg_dtypes)) is not None: + if not unknown: + return None + ran, error = _run_udf_probe(numba_func, arg_dtypes) + if ran is not None: for path, expected, i in unknown: value = ran[0] if i is None else ran[0][i] if isinstance(value, list): @@ -740,6 +755,7 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ ) else: _raise_unless_fits(np.shape(value), expected, return_type, path, probed=True) + return error def _get_udt_wrapper( @@ -769,7 +785,7 @@ def _get_udt_wrapper( yarg = ", y_ptr" yderef_expr = f", {yderef}" wrapper_sig = nt.void(*wrapper_args) - _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes) + error = _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes) rcidx = ", row, col" if include_indexes else "" signature_line = f"def wrapper(z_ptr, x_ptr{rcidx}{yarg}):" @@ -784,6 +800,7 @@ def _get_udt_wrapper( source_label=f"", extra_ns={"numba_func": numba_func}, ) + wrapper._raises_on_samples = error # for _finalize_udt_op return wrapper, wrapper_sig diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index f376c43ad..ab85d2b7c 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -3191,7 +3191,7 @@ def _sum(x, y): # pragma: no cover (numba) assert op[udt].return_type is udt assert probed == ["_sum"] - # A UDF that raises on the probe's ones is run again on twos, and types. + # A UDF that raises on the probe's values is not checked, and still types. def _no_ones(x, y): # pragma: no cover (numba) if x[0] == 1: raise ValueError("the probe's values") @@ -3203,36 +3203,55 @@ def _no_ones(x, y): # pragma: no cover (numba) @pytest.mark.skipif("not supports_udfs") -def test_udt_udf_raising_for_its_operand_types_raises_when_typed(): - """A UDF that raises on any values of its operand types is refused when the op is typed. +def test_udt_udf_raising_on_stand_ins_still_types(): + """A UDF that raises on the probe's stand-in operands still types, and runs on real data. - The probe gave up on any exception, so the op typed, and its cfunc raised - for every element, where Numba prints and swallows the error: the result - held zeros. Here ``x[:2] * t`` cannot broadcast when ``t`` is a 6-element - array, as it is for an ``(R6, R6)`` op. + Stand-ins of ones and twos are degenerate: a matrix of ones is singular, + and an interval check fails on equal fields. The index and thunk + stand-ins do not vary at all. So no UDF is refused for raising on them; + "Refuse a UDT UDF that raises for its operand types when the op is typed" + refused each of these. Only the literal rule reads that a UDF raised on + every stand-in: beside such an op a number is typed as itself. """ - r6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_RaiseTyped6") - r2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_RaiseTyped2") - iu = IndexUnaryOp.register_anonymous( - lambda x, i, j, t: x[:2] * t, "_raise_typed_iu", is_udt=True, ret_dtype=r2 - ) - bop = BinaryOp.register_anonymous( - lambda x, y: x[:2] * y, "_raise_typed_b", is_udt=True, ret_dtype=r2 - ) - for parent, types in [ - (iu, "_RaiseTyped6, UINT64, UINT64, _RaiseTyped6"), - (bop, "_RaiseTyped6, _RaiseTyped6"), - ]: - with pytest.raises( - UdfParseError, match=rf"raised ValueError when run on sample values of \({types}\)" - ): - parent[r6, r6] - # With a scalar operand the op types and runs - v = Vector(r6, size=2) - v[0] = np.arange(6.0) - three = gb.Scalar.from_value(3.0) - assert v.apply(iu, three).new()[0].new().value.tolist() == [0.0, 3.0] - assert v.apply(bop, right=three).new()[0].new().value.tolist() == [0.0, 3.0] + m22 = dtypes.register_anonymous(np.dtype((np.float64, (2, 2))), "_StandIn22") + p2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_StandIn2") + p3 = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_StandIn3") + p6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_StandIn6") + + def inverse(x): # pragma: no cover (numba) + return np.linalg.inv(x) + + def widths(x): # pragma: no cover (numba) + if x[0] >= x[1]: + raise ValueError("empty interval") + return x * 2.0 + + def rows(x, i, j, t): # pragma: no cover (numba) + return x.reshape((int(t), -1)).sum(axis=0) + + def need_t(x, i, j, t): # pragma: no cover (numba) + if t < 1.0: + raise ValueError("t must be at least 1") + return x[:2] * t + + m = Vector(m22, size=1) + m[0] = [[2.0, 1.0], [1.0, 3.0]] + result = m.apply(UnaryOp.register_anonymous(inverse, is_udt=True)).new() + np.testing.assert_allclose(result[0].new().value, [[0.6, -0.2], [-0.2, 0.4]]) + p = Vector(p2, size=1) + p[0] = [1.0, 4.0] + result = p.apply(UnaryOp.register_anonymous(widths, is_udt=True)).new() + assert result[0].new().value.tolist() == [2.0, 8.0] + s = Vector(p6, size=1) + s[0] = np.arange(6.0) + by_rows = IndexUnaryOp.register_anonymous(rows, is_udt=True, ret_dtype=p3) + assert s.apply(by_rows, gb.Scalar.from_value(2.0)).new()[0].new().value.tolist() == [3, 5, 7] + checked = IndexUnaryOp.register_anonymous(need_t, is_udt=True, ret_dtype=p2) + for thunk in [gb.Scalar.from_value(3.0), 3.0]: + assert s.apply(checked, thunk).new()[0].new().value.tolist() == [0.0, 3.0] + # An op whose UDF raises for every element thunk types too, unchecked + iu = IndexUnaryOp.register_anonymous(lambda x, i, j, t: x[:2] * t, is_udt=True, ret_dtype=p2) + assert iu[p6, p6].return_type is p2 @pytest.mark.skipif("not supports_udfs") From 038f803e44fe51653f4c3e7ebfb9c96b5783196d Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 15:22:45 -0800 Subject: [PATCH 15/20] Write 0 (False) into a UDF's built-in output before calling it Numba prints and swallows an exception raised in a cfunc, so a UDF that raised left its output as it was. For a select that is SuiteSparse's "bool keep", an uninitialized local (GB_select_generic_bitmap.c), which then went into the bitmap and the entry count: r612 saw nvals from 6 to 256 in a Vector of size 2 when a user SelectOp's UDF raised. Every wrapper of a user's UDF with a built-in output (UnaryOp, BinaryOp, IndexUnaryOp and SelectOp, IndexBinaryOp; built-in operand types, and the UDT wrapper's built-in outputs) now reads its operands, writes 0, then calls the UDF. Operands are read first because the output can be one of them (z = f(z, y) in a reduce); writing 0 first broke a dozen reduce tests. Expected and now checked, for a UDF that raises for negative values: INT64 [-1, 5, 2, -3, 7], select x > 3 -> nvals 2, indices [1, 4], values [5, 7] the same with apply x * 10 -> [0, 50, 20, 0, 70] a record UDT select whose UDF raises for entries 0 and 3 -> nvals 1, index [1] Those held by chance here before (the stray byte was 0). A second test calls a UDT wrapper directly on an output of 0xAB bytes: it now leaves 0, for a bool and for an FP64; it left 0xAB. A UDT output is not zeroed: it can also be an operand, copying a UDT operand first would copy every element, and Numba cannot re-raise, so catching the exception would silence it. _compose_wrapper_body says so. --- graphblas/core/operator/binary.py | 20 ++++- graphblas/core/operator/indexbinary.py | 24 ++++-- graphblas/core/operator/indexunary.py | 20 ++++- graphblas/core/operator/udf.py | 28 +++++-- graphblas/core/operator/unary.py | 20 ++++- graphblas/tests/test_op.py | 111 +++++++++++++++++++++++++ 6 files changed, 195 insertions(+), 28 deletions(-) diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 9ba0818ef..cb8abde73 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -505,26 +505,38 @@ def _finalize_typed_binaryop(parent, binary_udf, name, type_, ret_type): nt.CPointer(input_type.numba_type), ) + # Each wrapper writes 0 (False) before it calls the UDF: if the UDF raises, Numba + # prints the error and returns, and GraphBLAS uses the output as it is (for a + # select, an unwritten bool decides whether to keep the entry). The operands are + # read first, since the output can be one of them (z = f(z, y) in a reduce). if type_ == BOOL: if ret_type == BOOL: def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = bool(binary_udf(bool(x[0]), bool(y[0]))) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = bool(binary_udf(bool(x0), bool(y0))) else: def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = binary_udf(bool(x[0]), bool(y[0])) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = binary_udf(bool(x0), bool(y0)) elif ret_type == BOOL: def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = bool(binary_udf(x[0], y[0])) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = bool(binary_udf(x0, y0)) else: def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = binary_udf(x[0], y[0]) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = binary_udf(x0, y0) binary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(binary_wrapper) new_binary = ffi_new("GrB_BinaryOp*") diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index 91b31a812..d9c427d65 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -420,16 +420,20 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False, ret_dtype=None): nt.CPointer(input_type.numba_type), ) + # Each wrapper writes 0 (False) before it calls the UDF: if the UDF raises, Numba + # prints the error and returns, and GraphBLAS uses the output as it is (for a + # select, an unwritten bool decides whether to keep the entry). The operands are + # read first, since the output can be one of them (z = f(z, y) in a reduce). if type_ == BOOL: if ret_type == BOOL: def indexbinary_wrapper( z, x, ix, jx, y, iy, jy, theta ): # pragma: no cover (numba) + x0, y0, theta0 = x[0], y[0], theta[0] + z[0] = 0 z[0] = bool( - indexbinary_udf( - bool(x[0]), ix, jx, bool(y[0]), iy, jy, bool(theta[0]) - ) + indexbinary_udf(bool(x0), ix, jx, bool(y0), iy, jy, bool(theta0)) ) else: @@ -437,23 +441,27 @@ def indexbinary_wrapper( def indexbinary_wrapper( z, x, ix, jx, y, iy, jy, theta ): # pragma: no cover (numba) - z[0] = indexbinary_udf( - bool(x[0]), ix, jx, bool(y[0]), iy, jy, bool(theta[0]) - ) + x0, y0, theta0 = x[0], y[0], theta[0] + z[0] = 0 + z[0] = indexbinary_udf(bool(x0), ix, jx, bool(y0), iy, jy, bool(theta0)) elif ret_type == BOOL: def indexbinary_wrapper( z, x, ix, jx, y, iy, jy, theta ): # pragma: no cover (numba) - z[0] = bool(indexbinary_udf(x[0], ix, jx, y[0], iy, jy, theta[0])) + x0, y0, theta0 = x[0], y[0], theta[0] + z[0] = 0 + z[0] = bool(indexbinary_udf(x0, ix, jx, y0, iy, jy, theta0)) else: def indexbinary_wrapper( z, x, ix, jx, y, iy, jy, theta ): # pragma: no cover (numba) - z[0] = indexbinary_udf(x[0], ix, jx, y[0], iy, jy, theta[0]) + x0, y0, theta0 = x[0], y[0], theta[0] + z[0] = 0 + z[0] = indexbinary_udf(x0, ix, jx, y0, iy, jy, theta0) indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( indexbinary_wrapper diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index 2bab291bd..7f1a2e310 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -162,26 +162,38 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False, ret_dtype=None): nt.CPointer(input_type.numba_type), ) + # Each wrapper writes 0 (False) before it calls the UDF: if the UDF raises, Numba + # prints the error and returns, and GraphBLAS uses the output as it is (for a + # select, an unwritten bool decides whether to keep the entry). The operands are + # read first, since the output can be one of them (z = f(z, y) in a reduce). if type_ == BOOL: if ret_type == BOOL: def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) - z[0] = bool(indexunary_udf(bool(x[0]), row, col, bool(y[0]))) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = bool(indexunary_udf(bool(x0), row, col, bool(y0))) else: def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) - z[0] = indexunary_udf(bool(x[0]), row, col, bool(y[0])) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = indexunary_udf(bool(x0), row, col, bool(y0)) elif ret_type == BOOL: def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) - z[0] = bool(indexunary_udf(x[0], row, col, y[0])) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = bool(indexunary_udf(x0, row, col, y0)) else: def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) - z[0] = indexunary_udf(x[0], row, col, y[0]) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = indexunary_udf(x0, row, col, y0) indexunary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( indexunary_wrapper diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index fa32a6122..bc8813469 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -442,12 +442,15 @@ def _input_operand(dtype, var): var, nt.CPointer(base_numba), ) + # A built-in operand is read before the wrapper zeros a built-in output, which + # GraphBLAS can pass as an operand too (z = f(z, y) in a reduce). + setup = f" {var}0 = {var}_ptr[0]\n" if dtype == BOOL: # Numba can't compile bool ptrs (numba/numba#5395); expose them # as int8 and cast on deref. # MAINT 2026-10-18: still hits on Numba 0.68; re-test periodically. - return "", f"bool({var}_ptr[0])", nt.CPointer(INT8.numba_type) - return "", f"{var}_ptr[0]", nt.CPointer(dtype.numba_type) + return setup, f"bool({var}0)", nt.CPointer(INT8.numba_type) + return setup, f"{var}0", nt.CPointer(dtype.numba_type) def _output_handler(return_type, numba_ret_type): @@ -510,13 +513,22 @@ def _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr): compile a nested-tuple assignment to an outer record field. Array returns and array-typed record leaves slice-assign into the caller's buffer, so a wrong-shape return cannot overrun it. The error does not reach - the caller either: the resulting ``ValueError`` is raised inside a - cfunc, which Numba prints and swallows, so the write stops there and - every leaf from that point on keeps whatever SuiteSparse had in the - buffer. :func:`_check_udf_fills_output` rejects such a return when the - op is typed, but it cannot always learn the shape, so this stays the - backstop. + the caller either: the resulting ``ValueError``, like any exception the + UDF raises, is raised inside a cfunc, which Numba prints and swallows, + and GraphBLAS takes the output as the wrapper left it. So a built-in + output (a select's bool, say) is set to 0 before the UDF is called: an + element the UDF raises for holds 0, or False, which a select reads as + "do not keep". A select's output is an uninitialized bool in SuiteSparse, + so it could keep the entry and corrupt the entry count. A UDT output is + not zeroed: it can be an operand too (z = f(z, y) in a reduce), and + copying a UDT operand first would cost a copy of every element. Numba + cannot re-raise, so catching the exception would silence it. + :func:`_check_udf_fills_output` rejects a return that cannot fill the + output when the op is typed, but it cannot always learn the shape, so + this stays the backstop. """ + if zkind == "direct" and zinfo[2] == "z_ptr[0]": + body_setup = f"{body_setup} z_ptr[0] = 0\n" if zkind == "array_elements": return ( f"{signature_line}\n" diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index de4424f6e..23438328a 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -221,26 +221,38 @@ def _build(cls, name, func, *, anonymous=False, is_udt=False, ret_dtype=None): nt.CPointer(input_type.numba_type), ) + # Each wrapper writes 0 (False) before it calls the UDF: if the UDF raises, Numba + # prints the error and returns, and GraphBLAS uses the output as it is (for a + # select, an unwritten bool decides whether to keep the entry). The operands are + # read first, since the output can be one of them (z = f(z, y) in a reduce). if type_ == BOOL: if ret_type == BOOL: def unary_wrapper(z, x): - z[0] = bool(unary_udf(bool(x[0]))) # pragma: no cover (numba) + x0 = x[0] # pragma: no cover (numba) + z[0] = 0 # pragma: no cover (numba) + z[0] = bool(unary_udf(bool(x0))) # pragma: no cover (numba) else: def unary_wrapper(z, x): - z[0] = unary_udf(bool(x[0])) # pragma: no cover (numba) + x0 = x[0] # pragma: no cover (numba) + z[0] = 0 # pragma: no cover (numba) + z[0] = unary_udf(bool(x0)) # pragma: no cover (numba) elif ret_type == BOOL: def unary_wrapper(z, x): - z[0] = bool(unary_udf(x[0])) # pragma: no cover (numba) + x0 = x[0] # pragma: no cover (numba) + z[0] = 0 # pragma: no cover (numba) + z[0] = bool(unary_udf(x0)) # pragma: no cover (numba) else: def unary_wrapper(z, x): - z[0] = unary_udf(x[0]) # pragma: no cover (numba) + x0 = x[0] # pragma: no cover (numba) + z[0] = 0 # pragma: no cover (numba) + z[0] = unary_udf(x0) # pragma: no cover (numba) unary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( unary_wrapper diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index ab85d2b7c..79247e800 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -3202,6 +3202,117 @@ def _no_ones(x, y): # pragma: no cover (numba) assert probed == ["_sum", "_no_ones"] +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning") +def test_udf_raising_for_some_elements_gives_zeros_there(): + """An element the UDF raises for holds 0, and a select does not keep it. + + Numba prints and swallows an exception raised in a cfunc, so the wrapper + wrote nothing there. A select then read SuiteSparse's ``bool keep``, an + uninitialized local, which could keep the entry and add a stray byte to + the entry count (r612 saw nvals from 6 to 256 in a Vector of size 2). The + wrappers now write 0, or False, into a built-in output before they call + the UDF. What this checks held by chance before (the stray byte was 0); + test_udt_wrapper_zeros_a_builtin_output_before_the_udf checks the wrapper + itself. + """ + + def keep_above(x, i, j, t): # pragma: no cover (numba) + if x < 0: + raise ValueError("negative") + return x > t + + def times_ten(x): # pragma: no cover (numba) + if x < 0: + raise ValueError("negative") + return x * 10 + + v = Vector.from_coo([0, 1, 2, 3, 4], [-1, 5, 2, -3, 7]) + # -1 and -3 raise, so they are not kept; of the others, 5 and 7 are above 3 + kept = v.select(SelectOp.register_anonymous(keep_above, "_raise_keep_above"), 3).new() + assert kept.nvals == 2 + assert kept.to_coo()[0].tolist() == [1, 4] + assert kept.to_coo()[1].tolist() == [5, 7] + # -1 and -3 raise, so they hold 0 + tens = v.apply(UnaryOp.register_anonymous(times_ten, "_raise_times_ten")).new() + assert tens.to_coo()[1].tolist() == [0, 50, 20, 0, 70] + + rec = dtypes.register_anonymous(np.dtype([("rz_a", np.int64), ("rz_b", np.float64)])) + + def keep_rec(x, i, j, t): # pragma: no cover (numba) + if x.rz_a < 0: + raise ValueError("negative") + return x.rz_b > t + + r = Vector(rec, size=4) + r[0] = (-1, 9.0) + r[1] = (1, 9.0) + r[2] = (1, 0.0) + r[3] = (-2, 9.0) + # Entries 0 and 3 raise; of the others, only 1 has rz_b above 1 + sel = SelectOp.register_anonymous(keep_rec, "_raise_keep_rec", is_udt=True) + kept = r.select(sel, gb.Scalar.from_value(1.0)).new() + assert kept.nvals == 1 + assert kept.to_coo()[0].tolist() == [1] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning") +def test_udt_wrapper_zeros_a_builtin_output_before_the_udf(): + """A UDT op's wrapper sets a built-in output to 0 (False) before it calls the UDF. + + Called on an output full of 0xAB with an input the UDF raises for, it + leaves 0: a select's bool, and an FP64. It left the 0xAB bytes, which a + select read as "keep". (A UDT output is not zeroed; see + ``udf._compose_wrapper_body``.) + """ + import ctypes + + import numba + + from graphblas.core.operator.udf import _get_udt_wrapper + + rec = dtypes.register_anonymous( + np.dtype([("zo_n", np.int8), ("zo_s", "S3"), ("zo_d", "M8[s]")], align=True) + ) + + def keep(x, i, j, t): # pragma: no cover (numba) + if x.zo_n < 0: + raise ValueError("negative") + return True + + def half(x): # pragma: no cover (numba) + if x.zo_n < 0: + raise ValueError("negative") + return 0.5 + + x = np.zeros(1, dtype=rec.np_type) + x["zo_n"] = -1 + thunk = np.zeros(1, np.float64) + index = np.zeros(1, np.uint64) + for func, out, kwargs in [ + (keep, dtypes.BOOL, {"include_indexes": True, "dtype2": FP64}), + (half, FP64, {}), + ]: + numba_func = numba.njit(func) + arg_types = (rec.numba_type,) + if kwargs: + arg_types += (numba.types.uint64, numba.types.uint64, numba.types.float64) + numba_func.compile(arg_types) + ret = numba_func.overloads[arg_types].signature.return_type + wrapper, sig = _get_udt_wrapper(numba_func, out, rec, numba_ret_type=ret, **kwargs) + cfunc = numba.cfunc(sig, nopython=True)(wrapper) + z = np.full( + np.dtype(np.int8 if out is dtypes.BOOL else out.np_type).itemsize, 0xAB, np.uint8 + ) + args = [ctypes.cast(z.ctypes.data, cfunc.ctypes.argtypes[0])] + args.append(ctypes.cast(x.ctypes.data, cfunc.ctypes.argtypes[1])) + if kwargs: + args += [index[0], index[0], ctypes.cast(thunk.ctypes.data, cfunc.ctypes.argtypes[4])] + cfunc.ctypes(*args) + assert not z.any(), out + + @pytest.mark.skipif("not supports_udfs") def test_udt_udf_raising_on_stand_ins_still_types(): """A UDF that raises on the probe's stand-in operands still types, and runs on real data. From 1bcc6bd65b1feffbb1f12b1a92bfa42ba2502c1b Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 15:25:30 -0800 Subject: [PATCH 16/20] Probe a select UDF that returns an array, so a literal thunk can type as itself A user's SelectOp UDF that uses its thunk as a number, such as x[0] > t, gets an element of the UDT for a literal thunk, so it returns three bools where the select needs one, and the wrapper raised for every entry: v.select(first_above, 0) (v of an INT8[3] UDT: [1, 2, 3], [-4, 5, 6]) before: nothing kept (with a stray byte deciding, before the previous commit) after: [0] kept, as with Scalar.from_value(0); also for np.int8(0) and 0.5 SelectOp._compile_udt now hands the UDF's Numba return type to the wrapper, and _check_udf_fills_output runs the UDF on stand-ins when it returns an array for a scalar output. One that raises, or returns more than one value, is recorded in _RAISES_ON_SAMPLES (not refused: the shape may depend on the values), so the literal rule types the number as itself. A select that takes an element ((x == t).all()) still gets one. --- graphblas/core/operator/select.py | 4 +++- graphblas/core/operator/udf.py | 10 +++++++++- graphblas/tests/test_op.py | 29 +++++++++++++++++++++++++++++ 3 files changed, 41 insertions(+), 2 deletions(-) diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 1c1138ae5..a78b391b1 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -272,8 +272,10 @@ def _compile_udt(self, dtype, dtype2): _compile_udf_for_udt( numba_func, sig, op_kind="select", op_name=self.name, dtypes=(dtype, dtype2) ) + # Its return type lets the wrapper see an array returned for the bool + numba_ret_type = numba_func.overloads[sig].signature.return_type select_wrapper, wrapper_sig = _get_udt_wrapper( - numba_func, BOOL, dtype, dtype2, include_indexes=True + numba_func, BOOL, dtype, dtype2, include_indexes=True, numba_ret_type=numba_ret_type ) return _finalize_udt_op( self, dtype, dtype2, BOOL, select_wrapper, wrapper_sig, TypedUserSelectOp diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index bc8813469..6a9acf0d4 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -725,8 +725,16 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ (path, numba_ret_type.types[i], leaf.shape, i) for i, (path, _c, leaf) in enumerate(leaves) ] + elif isinstance(numba_ret_type, numba.core.types.Array): + # An array returned for a scalar (a select's bool, say) fills it only with one + # value, which running the UDF shows; one that cannot is recorded, not refused, + # since its shape may depend on the values (``x[x > 0]``). + ran, error = _run_udf_probe(numba_func, arg_dtypes) + if ran is not None and np.size(ran[0]) != 1: + error = ValueError(f"returned an array of shape {np.shape(ran[0])} for a scalar") + return error else: - return + return None unknown = [] for path, nb_type, expected, i in slots: if isinstance(nb_type, (numba.core.types.Optional, numba.core.types.NoneType)): diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 79247e800..d2b6ea54a 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -3365,6 +3365,35 @@ def need_t(x, i, j, t): # pragma: no cover (numba) assert iu[p6, p6].return_type is p2 +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning") +def test_udt_select_takes_a_number_as_itself_when_it_cannot_take_an_element(): + """A literal thunk for a user's select that cannot take an element is typed as itself. + + ``x[0] > t`` with ``t`` an INT8[3] element returns three bools for the one + the select needs, so the wrapper raised for every entry and nothing was + kept (or, before the select's bool was written first, a stray byte decided). + The select's UDF is now run on stand-ins when it returns an array, and a + number beside it is typed as itself, as a Scalar of it is. + """ + i3 = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_SelOwnI3") + first_above = SelectOp.register_anonymous( + lambda x, i, j, t: x[0] > t, "_sel_own_first", is_udt=True + ) + v = Vector(i3, size=2) + v[0] = [1, 2, 3] + v[1] = [-4, 5, 6] + for thunk in [0, np.int8(0), 0.5, gb.Scalar.from_value(0)]: + kept = v.select(first_above, thunk).new() + assert kept.to_coo()[0].tolist() == [0] + # A select that takes an element still gets one + same = SelectOp.register_anonymous( + lambda x, i, j, t: (x == t).all(), "_sel_own_same", is_udt=True + ) + v[1] = [1, 1, 1] + assert v.select(same, 1).new().to_coo()[0].tolist() == [1] + + @pytest.mark.skipif("not supports_udfs") def test_udt_user_op_takes_a_number_as_itself_when_it_cannot_take_an_element(): """A number beside a user's op is typed as itself when the op cannot take it as an element. From 61a3899446f4f5961a9a5a36210e2d76a3d5ce39 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 15:25:53 -0800 Subject: [PATCH 17/20] Let a complex return be declared as ret_dtype=BOOL again "Say that a complex return cannot be written as a real ret_dtype" also refused BOOL, which worked: the wrapper writes bool(...) for a BOOL output. BinaryOp.register_anonymous(lambda x, y: x * 1j + y, is_udt=True, ret_dtype=BOOL), on 1.5 and 0.0 #613's head: True before: UdfParseError: ... cannot be written as ret_dtype=BOOL after: True --- graphblas/core/operator/udf.py | 7 ++++--- graphblas/tests/test_op.py | 6 +++++- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index 6a9acf0d4..73becd3c6 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -265,14 +265,15 @@ def _check_declared_ret_type(ret_dtype, numba_ret_type): def _refuse_complex_into_real(ret_dtype, element, numba_ret_type): - """Raise unless a complex return goes into complex ``element``s of ``ret_dtype``. + """Raise unless a complex return goes into complex, or bool, ``element``s of ``ret_dtype``. Numba has no cast from complex to real, so the wrapper fails to compile, - with a message that loses why (``Failed in nopython mode pipeline``). + with a message that loses why (``Failed in nopython mode pipeline``). A + bool output is written as ``bool(...)``, which takes a complex value. """ nt = numba.core.types returned = numba_ret_type.dtype if isinstance(numba_ret_type, nt.Array) else numba_ret_type - if isinstance(returned, nt.Complex) and element.kind != "c": + if isinstance(returned, nt.Complex) and element.kind not in "cb": raise UdfParseError( f"UDT UDF returned {numba_ret_type}, which cannot be written as " f"ret_dtype={ret_dtype}. Return a real value, or declare a complex ret_dtype." diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index d2b6ea54a..4498964ea 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -7872,9 +7872,13 @@ def binary_rejects(func, ret_dtype, match): assert flat[FP64, FP64].return_type is nested binary_rejects(lambda x, y: (x, y, x + y), arr2, "tuple of length 3, which cannot be written") binary_rejects(lambda x, y: (x, y), FP64, "tuple of length 2, which cannot be written") - # A complex return into real values; Numba's own message lost why + # A complex return into real values; Numba's own message lost why. Into BOOL it is + # written as bool(...), so it works for ret_dtype in [FP64, arr2]: binary_rejects(lambda x, y: x * 1j + y, ret_dtype, "complex128, .* Return a real value") + nonzero = BinaryOp.register_anonymous(lambda x, y: x * 1j + y, is_udt=True, ret_dtype=BOOL) + one, zero = Vector.from_coo([0], [1.5]), Vector.from_coo([0], [0.0]) + assert one.ewise_mult(zero, nonzero).new()[0].new().value is True binary_rejects( lambda x, y: x if x > y else None, FP64, "returned a value that can be None, which cannot" ) From dd30f6cfece56dd0fbfa695456b639471d581f90 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 15:26:05 -0800 Subject: [PATCH 18/20] Say in udt.rst what a UDF that raises gets, at typing and when it runs The shape check's sentence said a UDF that raises on the sample values is not checked, which holds again after "Type a UDF that raises on stand-ins as before; only the literal rule reads it"; it now says why (it may raise only for some values), and what an exception at run time leaves: the element is not written, and a built-in output (a select's bool) is 0 or False. --- docs/user_guide/udt.rst | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 0ef6ea576..9bb5e803a 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -368,7 +368,11 @@ array returned into an array-typed field of a record UDT. A tuple or list also fills such a field, but only a one-dimensional one and only at its exact length; unlike an array, it does not broadcast. Numba does not know the shape of an array the UDF builds, or the length of a list, so the check runs the UDF -once on sample values; a UDF that raises on those values is not checked. +on sample values; a UDF that raises on those values is not checked, since it +may raise only for some values. An exception the UDF raises when the op runs +is printed and goes no further (Numba cannot pass it to GraphBLAS), so that +element is not written; a built-in output, such as a select's bool, is 0 or +``False`` there. If your UDF references a field that doesn't exist, or returns the wrong arity, you'll get a ``UdfParseError`` with the actionable diagnostic line surfaced From 860cd1ccd507d5dfab1e873e627e31cb880da3b2 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 15:58:16 -0800 Subject: [PATCH 19/20] Keep a UDF's stand-in failure as text, so the op it records can be freed _RAISES_ON_SAMPLES kept the exception the UDF raised on its stand-ins, and its traceback holds the frames of _check_udf_fills_output, _get_udt_wrapper and _compile_udt, which hold the typed op and its parent, so the weakly keyed record kept its own key alive (r612): 5 IndexUnaryOps (x[:2] * t, ret_dtype FP64[2]) applied with a literal thunk, then dropped and gc.collect() before: 5 of 5 ops alive, with their GrB handles; the record holds 5 after: 0 of 5 alive; the record is empty The probe now returns the failure as "Type: message", which the wrapper and the record keep instead; only membership is read. A test drops three recorded ops and checks that weak references to each op and typed op die. --- graphblas/core/operator/udf.py | 17 +++++++++-------- graphblas/tests/test_op.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 8 deletions(-) diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index 73becd3c6..1fb644b73 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -598,10 +598,11 @@ def _run_udf_probe(numba_func, arg_dtypes): that does not fit, that is left to the wrapper, whose output then keeps its zeros (see :func:`_compose_wrapper_body`). Such a UDF may raise only for some values, so it still types (``np.linalg.inv`` raises on a matrix - of ones). ``error`` is the exception when it raised on twos as well, and - it was no ``ArithmeticError``; only the literal rule reads it (see - :data:`_RAISES_ON_SAMPLES`). Only the call is guarded, so a mistake in - building the operands still surfaces. + of ones). ``error`` names the exception (as text) when it raised on twos + as well, and it was no ``ArithmeticError``; only the literal rule reads it + (see :data:`_RAISES_ON_SAMPLES`). The exception itself is not kept: its + traceback holds the frames that hold the op. Only the call is guarded, so + a mistake in building the operands still surfaces. This is where typing an op for a UDT runs the user's function rather than only compiling it, and ``OpBase.__contains__`` is a typed lookup, @@ -621,14 +622,14 @@ def _run_udf_probe(numba_func, arg_dtypes): except ArithmeticError: return None, None except Exception as exc: - error = exc + error = f"{type(exc).__name__}: {exc}" continue return ((result,) if fill == 1 else None), None return None, error # The UDT ops whose UDF raised, other than an ArithmeticError, on both sets of stand-in -# operands, with the exception. Such an op still types, since the UDF may raise only +# operands, with the exception as text. Such an op still types, since the UDF may raise only # for some values, but beside it a number is typed as itself rather than as an element # (``literals._own_scalar_type_for_user_op``), as a Scalar of it would be. _RAISES_ON_SAMPLES = weakref.WeakKeyDictionary() @@ -709,7 +710,7 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ length exactly (:func:`_raise_unless_sequence_fits`); a tuple's length is in its type, and a list's is learned the same way as a built array's. A UDF that returns no built array or list is never run. - Returns the exception the UDF raised on every stand-in (see + Returns what the UDF raised on every stand-in, as text (see :func:`_run_udf_probe`), or None. """ if zkind == "array_elements": @@ -732,7 +733,7 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ # since its shape may depend on the values (``x[x > 0]``). ran, error = _run_udf_probe(numba_func, arg_dtypes) if ran is not None and np.size(ran[0]) != 1: - error = ValueError(f"returned an array of shape {np.shape(ran[0])} for a scalar") + error = f"returned an array of shape {np.shape(ran[0])} for a scalar" return error else: return None diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 4498964ea..e4c06848d 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -3365,6 +3365,37 @@ def need_t(x, i, j, t): # pragma: no cover (numba) assert iu[p6, p6].return_type is p2 +@pytest.mark.skipif("not supports_udfs") +def test_udt_op_whose_udf_raised_on_stand_ins_can_be_collected(): + """An op the probe records as raising on every stand-in is freed once dropped. + + The record kept the exception, whose traceback holds the frames that hold the + typed op, so the op, its parent and their GraphBLAS handles were never freed + (r612: 5 of 5 alive). It keeps the exception as text now. + """ + import gc + import weakref + + from graphblas.core.operator.udf import _RAISES_ON_SAMPLES + + r6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_FreedStandIn6") + r2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_FreedStandIn2") + v = Vector(r6, size=1) + v[0] = np.arange(6.0) + refs = [] + for _ in range(3): + iu = IndexUnaryOp.register_anonymous( + lambda x, i, j, t: x[:2] * t, is_udt=True, ret_dtype=r2 + ) + assert v.apply(iu, 3.0).new()[0].new().value.tolist() == [0.0, 3.0] + typed = iu[r6, r6] + assert typed in _RAISES_ON_SAMPLES + refs += [weakref.ref(iu), weakref.ref(typed)] + del iu, typed + gc.collect() + assert [ref() for ref in refs] == [None] * len(refs) + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning") def test_udt_select_takes_a_number_as_itself_when_it_cannot_take_an_element(): From 450d00e6f65f6fc0444be0201db8d08cc05ac120 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 15:59:08 -0800 Subject: [PATCH 20/20] Test that each built-in-type UDF wrapper writes 0 before calling the UDF Nothing pinned the zero writes in the 16 wrappers a user's UnaryOp, BinaryOp, IndexUnaryOp (and SelectOp) and IndexBinaryOp get for built-in types: deleting every one left the tests passing (r612), since a select whose UDF raises keeps nothing by chance as often as not. The test catches the cfuncs from numba.cfunc while each op registers a UDF that always raises, and calls each with an output of 0xAB bytes: the output must be 0. With the 16 writes deleted it fails on the first wrapper. --- graphblas/tests/test_op.py | 79 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index e4c06848d..32bf229c7 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -3365,6 +3365,85 @@ def need_t(x, i, j, t): # pragma: no cover (numba) assert iu[p6, p6].return_type is p2 +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning") +def test_builtin_type_udf_wrappers_write_zero_before_the_udf(monkeypatch): + """Each wrapper a user's op gets for a built-in type writes 0 before it calls the UDF. + + The wrappers are made while the op registers, so they are caught from + ``numba.cfunc`` then, and each is called with an output of 0xAB bytes and a + UDF that always raises: the output must be 0 (False). Without the zero + writes it kept the 0xAB bytes, which a select reads as "keep". + """ + import ctypes + + import numba + + from graphblas.core.operator import IndexBinaryOp + + caught = [] + cfunc = numba.cfunc + + def catching_cfunc(*args, **kwargs): + decorate = cfunc(*args, **kwargs) + + def catch(func): + caught.append(decorate(func)) + return caught[-1] + + return catch + + # Numba types a return only if the raise is conditional + def raises_unary(x): # pragma: no cover (numba) + if x == x: + raise ValueError("raises") + return x + + def raises_binary(x, y): # pragma: no cover (numba) + if x == x: + raise ValueError("raises") + return x + + def raises_indexunary(x, i, j, t): # pragma: no cover (numba) + if x == x: + raise ValueError("raises") + return x + + def raises_indexbinary(x, ix, jx, y, iy, jy, t): # pragma: no cover (numba) + if x == x: + raise ValueError("raises") + return x + + monkeypatch.setattr(numba, "cfunc", catching_cfunc) + counts = {} + for cls, func in [ + (UnaryOp, raises_unary), + (BinaryOp, raises_binary), + (IndexUnaryOp, raises_indexunary), + (IndexBinaryOp, raises_indexbinary), + ]: + if cls is IndexBinaryOp and lib.__dict__.get("GxB_IndexBinaryOp_new") is None: + continue # pragma: no cover (SuiteSparse < 9.4) + start = len(caught) + cls.register_anonymous(func) + counts[cls.__name__] = len(caught) - start + monkeypatch.undo() + assert all(counts.values()), counts # one wrapper per built-in type the UDF compiles for + for wrapper in caught: + args = [] + out = None + for argtype, ctype in zip(wrapper._sig.args, wrapper.ctypes.argtypes, strict=True): + if not isinstance(argtype, numba.types.CPointer): + args.append(1) # an index + continue + np_type = numba.np.numpy_support.as_dtype(argtype.dtype) + buf = np.full(np_type.itemsize, 0xAB, np.uint8) if out is None else np.ones(1, np_type) + out = buf if out is None else out + args.append(ctypes.cast(buf.ctypes.data, ctype)) + wrapper.ctypes(*args) + assert not out.any(), wrapper._sig + + @pytest.mark.skipif("not supports_udfs") def test_udt_op_whose_udf_raised_on_stand_ins_can_be_collected(): """An op the probe records as raising on every stand-in is freed once dropped.