Let UDF registration name its return dtype with ret_dtype - #613
Open
eriknw wants to merge 2 commits into
Open
Conversation
eriknw
marked this pull request as ready for review
August 4, 2026 16:07
eriknw
force-pushed
the
33-udf-ret-dtype
branch
from
August 4, 2026 16:12
2a84393 to
7dc89b9
Compare
eriknw
force-pushed
the
33-udf-ret-dtype
branch
from
August 5, 2026 00:06
7dc89b9 to
4784268
Compare
eriknw
force-pushed
the
33-udf-ret-dtype
branch
from
August 5, 2026 03:18
4784268 to
9f5a7a7
Compare
eriknw
force-pushed
the
33-udf-ret-dtype
branch
from
August 5, 2026 17:44
9f5a7a7 to
80e9dfa
Compare
eriknw
force-pushed
the
33-udf-ret-dtype
branch
from
August 5, 2026 18:03
80e9dfa to
38d1c10
Compare
eriknw
force-pushed
the
33-udf-ret-dtype
branch
2 times, most recently
from
August 6, 2026 07:59
51bb3a8 to
c9c6a51
Compare
eriknw
force-pushed
the
33-udf-ret-dtype
branch
from
August 6, 2026 15:39
c9c6a51 to
48bbd23
Compare
eriknw
force-pushed
the
33-udf-ret-dtype
branch
from
August 6, 2026 15:41
48bbd23 to
c9791e6
Compare
eriknw
force-pushed
the
33-udf-ret-dtype
branch
from
August 6, 2026 20:36
c9791e6 to
896bc49
Compare
eriknw
force-pushed
the
33-udf-ret-dtype
branch
2 times, most recently
from
August 7, 2026 02:49
067de67 to
dd7690a
Compare
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.
eriknw
force-pushed
the
33-udf-ret-dtype
branch
from
August 7, 2026 05:09
dd7690a to
a9b0b47
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
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.
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.
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.
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.
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.
Stack created with GitHub Stacks CLI • Give Feedback 💬