python-graphblas · eriknw · Aug 3, 2026 · Aug 4, 2026
diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst
@@ -204,6 +204,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
@@ -112,6 +112,39 @@ def _bool_to_int8(dtype):
     return INT8 if dtype == BOOL else dtype
 
 
+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
@@ -270,6 +303,16 @@ def _summarize_numba_typing_error(exc):
                 return line
         return "Numba could not compile the function for these input types"
 
+    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.
 
@@ -1178,21 +1221,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.
@@ -1205,10 +1252,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
@@ -33,6 +33,7 @@
     TypedOpBase,
     _call_op,
     _hasop,
+    _validate_ret_dtype,
 )
 
 # Imported unconditionally (plain dict, no numba): ``_compile_udt`` consults it
@@ -47,7 +48,7 @@
         _compile_udf_for_udt,
         _finalize_udt_op,
         _get_udt_wrapper,
-        _resolve_udt_return_type,
+        _udt_ret_type,
     )
 
     try:
@@ -554,6 +555,7 @@ class BinaryOp(OpBase):
         "_numba_func",
         "_custom_dtype",
         "_defer_builds",
+        "_ret_dtype",
     )
     _module = binary
     _modname = "binary"
@@ -659,11 +661,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__", "<anonymous_binary>")
+        # 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
@@ -678,7 +686,14 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False):
         # interactively defined functions do not have, so ops registered by
         # users stay uncached.
         binary_udf = numba.njit(func, error_model="numpy", cache=cache)
-        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
@@ -751,7 +766,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
             )
@@ -790,7 +805,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.
@@ -819,19 +836,35 @@ 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
 
         """
         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.
 
@@ -867,6 +900,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):
@@ -890,6 +927,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] = (
@@ -899,14 +939,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)
@@ -1131,8 +1173,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