Conversation
eriknw
marked this pull request as ready for review
August 4, 2026 16:07
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
2 times, most recently
from
August 5, 2026 00:06
3676d9c to
9ab0b14
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
from
August 5, 2026 03:18
9ab0b14 to
f5c4f3b
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
from
August 5, 2026 17:44
f5c4f3b to
5a10a1e
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
from
August 5, 2026 18:03
5a10a1e to
9fc6fa2
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
2 times, most recently
from
August 6, 2026 07:59
58c5cc1 to
bc96df1
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
from
August 6, 2026 15:39
bc96df1 to
b719ac2
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
2 times, most recently
from
August 6, 2026 20:36
cb1874e to
db8e82d
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
from
August 6, 2026 20:42
db8e82d to
1e2baf3
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
2 times, most recently
from
August 7, 2026 05:09
0ab8b52 to
4d91df6
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
from
August 26, 2026 17:30
4d91df6 to
f209a6f
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
from
September 18, 2026 03:09
f209a6f to
91c7fca
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
from
September 18, 2026 03:46
91c7fca to
e01f5f5
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
from
September 18, 2026 17:34
e01f5f5 to
4c49bfd
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
2 times, most recently
from
September 23, 2026 16:22
6f0e968 to
ddcf648
Compare
eriknw
removed this pull request from stack #627
September 23, 2026 18:18
eriknw
changed the base branch from
09-array-udt-udf-views
to
41-ci-dep-pools-refresh
September 23, 2026 18:18
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
from
September 23, 2026 18:19
ddcf648 to
6f0f862
Compare
eriknw
added this pull request to stack #634
September 23, 2026 18:19
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
from
September 24, 2026 20:06
6f0f862 to
519441d
Compare
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
3 times, most recently
from
September 26, 2026 05:23
9d92d84 to
a689494
Compare
A UDF over an array UDT whose return does not fit the element is not a type error. Numba's `Array` type carries `ndim` but not its extents, so returning `x[:2]` from a 9-element UDT types identically to returning `x`. The wrapper's `z[:] = ...` then raises inside the cfunc, where Numba prints the traceback and returns. GraphBLAS is handed no error, and the element keeps whatever was already in the buffer, so the caller gets a wrong answer and no exception. A record UDF that under-fills an array-typed leaf fails the same way and abandons the write part-way: every leaf after the short one, scalar leaves included, keeps what SuiteSparse had there. Extents can only be recovered by running the function, so `_get_udt_wrapper` and the IndexBinaryOp wrapper now call it once on stand-in operands and check that the result fits the UDT. Fitting is not the same as matching. Slice-assign broadcasts, so a `(1,)` return fills a `(6,)` element and a `(1, 3)` return fills both rows of a `(2, 3)` one; both work today and must keep working, and an equality check would reject them. `_fits_by_broadcast` applies numpy's assignment rule instead, and `test_udt_broadcast_matches_numba_slice_assign` pins it against Numba's own slice-assign in both directions, including the two ranks where broadcasting alone gives the wrong answer. The operands are ones rather than zeros for the same reason `dtypes._sample_values` avoids zeros: a UDF that divides by an operand would raise on a zero-filled probe, and a raising probe is treated as "cannot check" rather than as a failure. The check is best-effort, and this commit should not claim more. A UDF that raises on the probe values is not checked at all, so one that also returns a shape that does not fit still reaches the wrapper and still fails silently, exactly as it did before. The wrapper's slice-assign remains the only backstop that always runs. What the check buys is that the common case, a UDF that is simply wrong about the shape, becomes an error at registration instead of a wrong answer at apply time. Because the probe reports the shape for the values it was run on, the diagnostic says so. A UDF whose output shape depends on its operands' values rather than their types can otherwise be rejected on a shape it would never return for real data. This is a judgement call rather than a straight fix, and it is cheap to decline. Typing an op for a UDT now executes the user's function, where before it only compiled it. `OpBase.__contains__` is implemented as a typed lookup, so `udt in some_op` runs user code as a side effect of a membership test. Measured: registering an op runs no probe, since compilation stays lazy, but a single `udt in some_op` runs exactly one. That probe costs no extra compile, because it calls the function with an `ndarray` and the wrapper's `numba.carray` view needs that same specialization anyway. A/B with the check stubbed out: 93.4 ms median disabled against 92.1 ms enabled, two Numba specializations either way. Only array UDTs and records that actually have an array leaf are probed, so an ordinary record UDF is still typed without running user code. The array-UDT paragraph in `docs/user_guide/udt.rst`, added one commit earlier, described the rule as building "a new array of the same shape". It now states the rule the check enforces, and says the shape is learned by running the UDF. If this is declined, the array-UDT view change it sits on still stands. A return that does not fit stays memory-safe, because the slice-assign cannot overrun the element, and goes back to being silent.
test_udt_array_udf_shape_errors expected a flat FP64[3, 4] and a layered FP64[4] x 3 to be separate UDTs, so that a UDF building an array from the pair was ambiguous. "Flatten nested subarray dtypes so each UDT layout has one DataType" (#589) made them the same DataType, so nothing is ambiguous and the test failed with "DID NOT RAISE UdfParseError" on every CI job with Numba. Keep the same-type check that block ended with, now on udt9. The flattening itself is tested where it is introduced.
The shape check ran the UDF on sample values for every array return and every record return with an array leaf. An operand, or operand field, returned as-is keeps its extents in Numba's `NestedArray` type, so check those from the type and run the UDF only for an array it builds (a plain `Array`). Returning an operand no longer runs user code when the op is typed, and fields returned into the wrong slots are caught even if the UDF would raise on the sample values. The two checkers become one, `_check_udf_fills_output`. Also reject a record leaf the UDF can return as `None` (typed `Optional`, or always `None`), scalar leaves included. On that path the wrapper's assignment raises inside the cfunc and the element comes back holding another element's values. The probe docstring's reason for using ones was wrong: UDFs compile with `error_model="numpy"`, so dividing by zero does not raise. Drop two test cases that repeated #589's resolver tests, and add a test pinning when the UDF is run.
eriknw
force-pushed
the
10-udt-udf-shape-checks
branch
from
September 26, 2026 05:25
a689494 to
0793d8b
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 over an array UDT whose return does not fit the element is not a type
error. Numba's
Arraytype carriesndimbut not its extents, so returningx[:2]from a 9-element UDT types identically to returningx. Thewrapper's
z[:] = ...then raises inside the cfunc, where Numba prints thetraceback and returns. GraphBLAS is handed no error, and the element keeps
whatever was already in the buffer, so the caller gets a wrong answer and no
exception. A record UDF that under-fills an array-typed leaf fails the same
way and abandons the write part-way: every leaf after the short one, scalar
leaves included, keeps what SuiteSparse had there.
Extents can only be recovered by running the function, so
_get_udt_wrapperand the IndexBinaryOp wrapper now call it once on stand-in operands and check
that the result fits the UDT. Fitting is not the same as matching. Slice-assign
broadcasts, so a
(1,)return fills a(6,)element and a(1, 3)returnfills both rows of a
(2, 3)one; both work today and must keep working, andan equality check would reject them.
_fits_by_broadcastapplies numpy'sassignment rule instead, and
test_udt_broadcast_matches_numba_slice_assignpins it against Numba's own slice-assign in both directions, including the two
ranks where broadcasting alone gives the wrong answer. The operands are ones
rather than zeros for the same reason
dtypes._sample_valuesavoids zeros: aUDF that divides by an operand would raise on a zero-filled probe, and a
raising probe is treated as "cannot check" rather than as a failure.
The check is best-effort, and this commit should not claim more. A UDF that
raises on the probe values is not checked at all, so one that also returns a
shape that does not fit still reaches the wrapper and still fails silently,
exactly as it did before. The wrapper's slice-assign remains the only backstop
that always runs. What the check buys is that the common case, a UDF that is
simply wrong about the shape, becomes an error at registration instead of a
wrong answer at apply time.
Because the probe reports the shape for the values it was run on, the
diagnostic says so. A UDF whose output shape depends on its operands' values
rather than their types can otherwise be rejected on a shape it would never
return for real data.
This is a judgement call rather than a straight fix, and it is cheap to
decline. Typing an op for a UDT now executes the user's function, where before
it only compiled it.
OpBase.__contains__is implemented as a typed lookup, soudt in some_opruns user code as a side effect of a membership test.Measured: registering an op runs no probe, since compilation stays lazy, but a
single
udt in some_opruns exactly one. That probe costs no extra compile,because it calls the function with an
ndarrayand the wrapper'snumba.carrayview needs that same specialization anyway. A/B with the checkstubbed out: 93.4 ms median disabled against 92.1 ms enabled, two Numba
specializations either way.
Only array UDTs and records that actually have an array leaf are probed, so an
ordinary record UDF is still typed without running user code.
The array-UDT paragraph in
docs/user_guide/udt.rst, added one commit earlier,described the rule as building "a new array of the same shape". It now states
the rule the check enforces, and says the shape is learned by running the UDF.
If this is declined, the array-UDT view change it sits on still stands. A
return that does not fit stays memory-safe, because the slice-assign cannot
overrun the element, and goes back to being silent.
Stack created with GitHub Stacks CLI • Give Feedback 💬