From 997ce44c8cf375d2596cf68f3dcb554c523b4694 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 31 Jul 2026 12:48:36 -0500 Subject: [PATCH 01/54] Reject record UDT pairs whose leaf counts differ `_check_udt_pair` rejects two record operands that disagree on shape, but it compared only top-level field names. Two records can share those and still nest differently: `[("a", f8), ("b", f8)]` against `[("a", [("n1", f8), ("n2", f8)]), ("b", f8)]` both report fields `["a", "b"]` while contributing two leaves and three. The codegen walks one operand's leaf paths and applies them to both, so applying the scalar path `["a"]` to the nested side asks Numba to add a float to a record. Measured before this change, that pair gets through the checks and comes back from Numba's typing pass as: UdfParseError: binary.plus does not work with (_NestFlat, _NestDeep): No implementation of function Function() found for signature: >>> add(float64, Record(nst_n1[type=float64;offset=0], nst_n2[type=float64;offset=8];16;False)) That reports a compile failure for what is the same shape disagreement the three sibling checks in the same function already report as a KeyError. Comparing leaf counts makes it the fourth of those checks. The exception class changes, and that is visible in both directions: UdfParseError derives from GraphblasException, not from KeyError, so code catching KeyError to mean "no such op" now catches this pair, and code catching UdfParseError no longer does. `OpBase.__contains__` catches both, so `udt in binary.plus` is unaffected either way. --- graphblas/core/operator/udt_utils.py | 27 +++++++++++++++++------ graphblas/tests/test_op.py | 32 ++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 6 deletions(-) diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index ede25b467..e5349235a 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -325,12 +325,27 @@ def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): f"binary.{op_name} does not work with ({dtype}, {dtype2}): " f"cannot mix record and array UDTs in a single element-wise op." ) - if kind_x == "record" and detail_x != detail_y: - raise KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"record UDTs must share field names; got {list(detail_x)} vs " - f"{list(detail_y)}." - ) + if kind_x == "record": + if detail_x != detail_y: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"record UDTs must share field names; got {list(detail_x)} vs " + f"{list(detail_y)}." + ) + # Matching top-level names is not enough: the codegen pairs operands + # leaf by leaf, and a field that is a sub-record on one side and a + # scalar on the other contributes a different number of leaves. Left + # unchecked the pair reaches Numba, whose typing failure arrives as a + # UdfParseError, reporting a compile error for what is really the same + # shape disagreement the checks above report as a KeyError. + leaves_x = [c for _py, c, _d in _iter_record_leaves(dtype.np_type)] + leaves_y = [c for _py, c, _d in _iter_record_leaves(dtype2.np_type)] + if len(leaves_x) != len(leaves_y): + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"record UDTs must nest the same way, so that each has the same " + f"number of leaf fields; got {leaves_x} vs {leaves_y}." + ) if kind_x == "array" and detail_x != detail_y: raise KeyError( f"binary.{op_name} does not work with ({dtype}, {dtype2}): " diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 36152dfe0..82abeb1c0 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -3167,6 +3167,38 @@ def test_udt_eq_ne_rejects_incompatible_pairs(): binary.eq(v_uv & v_arr).new() +@pytest.mark.skipif("not supports_udfs") +def test_udt_record_nesting_mismatch_is_a_keyerror(): + """Records sharing field names but not nesting depth are rejected as a KeyError. + + ``_check_udt_pair`` matched on top-level names only, but the codegen pairs + operands leaf by leaf, and a field that is a sub-record on one side and a + scalar on the other contributes a different number of leaves. Without the + guard the pair reaches Numba, whose typing failure arrives as a + ``UdfParseError``: a compile error reported for what is really the same + shape disagreement its sibling checks raise ``KeyError`` for. + """ + flat = dtypes.register_anonymous( + np.dtype([("nst_a", np.float64), ("nst_b", np.float64)], align=True), "_NestFlat" + ) + nested = dtypes.register_anonymous( + np.dtype( + [ + ("nst_a", np.dtype([("nst_n1", np.float64), ("nst_n2", np.float64)])), + ("nst_b", np.float64), + ], + align=True, + ), + "_NestDeep", + ) + v = Vector(flat, size=1) + v[0] = (1.0, 2.0) + w = Vector(nested, size=1) + w[0] = ((3.0, 4.0), 5.0) + with pytest.raises(KeyError, match="same number of leaf fields"): + v.ewise_mult(w, binary.plus).new() + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_aggregators(): From 480f08f330a3ce69c93ded1e19cbd2df4352a710 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 4 Aug 2026 16:57:25 -0700 Subject: [PATCH 02/54] Ignore the UDT repr size warning in the nesting mismatch test --- graphblas/tests/test_op.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 82abeb1c0..c0b90d6bd 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -3168,6 +3168,10 @@ def test_udt_eq_ne_rejects_incompatible_pairs(): @pytest.mark.skipif("not supports_udfs") +# SS < 9 has no GrB_NAME setter, so registration falls back to storing the +# numpy repr in the type name and warns when it does not fit in 128 chars. +# _NestDeep's repr is 142; how it serializes is not what the test is about. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") def test_udt_record_nesting_mismatch_is_a_keyerror(): """Records sharing field names but not nesting depth are rejected as a KeyError. From 1a7db6cd9fbff0c8651297e7650739def775f7be Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Sun, 4 Oct 2026 01:20:10 -0500 Subject: [PATCH 03/54] Compare record leaf paths and field shapes, not leaf counts Equal leaf counts do not make two records pair up. The codegen reads both operands at the left one's leaf paths, so records nested under different fields, or with differently named sub-records, still reached Numba's typing: a UdfParseError for plus, a raw TypingError for eq. The check now requires the same leaf paths in the same order, as it already does for top-level names. That also rejects reordered sub-record fields. They work here, but they give wrong answers once #593 pairs leaf dtypes by position and #595 branches on them. Array fields whose shapes do not broadcast were worse: (3,) with (4,) raised inside the cfunc, unseen, and left garbage in the result. Those are now KeyErrors, as is a scalar field on the left of an array one, which never compiled. The arithmetic ops also need the broadcast to fit the left field. A (1,) field left of a (3,) one was caught only by the UDF return check, worded as a UDF the user did not write. eq and ne reduce each field to one bool and keep accepting that pair. Fields that broadcast, such as (3,) with () or (1, 3), still combine; across 64 field-shape pairs under plus, min and eq, everything that worked before still does. The user guide now states the pairing rules. --- docs/user_guide/udt.rst | 10 ++ graphblas/core/operator/udt_utils.py | 59 +++++++++--- graphblas/tests/test_op.py | 132 +++++++++++++++++++++------ 3 files changed, 162 insertions(+), 39 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index ae2cef9e2..61f5d0af4 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -141,6 +141,16 @@ UDTs with any complex leaf (no ordering, no integer modulus); use ``plus``, ``minus``, ``times``, or ``truediv`` for complex arithmetic, or register a custom binary op. +The lifted binary ops also take two different UDTs. Two record UDTs pair up +when they nest the same way, with the same field names in the same order at +every level. Array fields combine as numpy arrays do, by broadcasting, and the +arithmetic ops give each result field the shape of the left operand's field, so +a ``(3,)`` field pairs with a ``(1,)`` or scalar one on its right but not on +its left. ``eq`` and ``ne`` reduce each field to one bool, so for them an array +field on the left need only broadcast with the right one. Two array UDTs pair +up when they share a base dtype and flat size. Any other pair raises +``KeyError`` on the first lookup. + Composite aggregators (``agg.hypot``, ``agg.L1norm``, ``agg.Linfnorm``, ``agg.sum_of_squares``, ``agg.sum_of_inverses``) do *not* auto-lift to UDTs; they reference scalar-only binary ops that don't generalize trivially. diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index e5349235a..c06dbbee5 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -314,7 +314,8 @@ def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): ``info_x`` and ``info_y`` are ``_get_udt_info`` results. When either is ``None`` (one side is a scalar broadcast), no check fires. Without this gate the generated wrapper code raises a cryptic Numba ``TypingError`` - on the first field access. + on the first field access, or, for array fields that do not broadcast, + fails inside the cfunc without telling the caller. """ if info_x is None or info_y is None: return @@ -332,19 +333,53 @@ def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): f"record UDTs must share field names; got {list(detail_x)} vs " f"{list(detail_y)}." ) - # Matching top-level names is not enough: the codegen pairs operands - # leaf by leaf, and a field that is a sub-record on one side and a - # scalar on the other contributes a different number of leaves. Left - # unchecked the pair reaches Numba, whose typing failure arrives as a - # UdfParseError, reporting a compile error for what is really the same - # shape disagreement the checks above report as a KeyError. - leaves_x = [c for _py, c, _d in _iter_record_leaves(dtype.np_type)] - leaves_y = [c for _py, c, _d in _iter_record_leaves(dtype2.np_type)] - if len(leaves_x) != len(leaves_y): + # Matching top-level names is not enough. The codegen reads both + # operands at the left one's leaf paths, so a field that is a + # sub-record on one side and a scalar on the other, or sub-records + # with different names, fails Numba's typing: a compile error + # reported for what is really the same shape disagreement the check + # above reports as a KeyError. The paths must also come in the same + # order, as the top-level names must, so that leaf ``i`` of one + # operand is leaf ``i`` of the other. + leaves_x = list(_iter_record_leaves(dtype.np_type)) + leaves_y = list(_iter_record_leaves(dtype2.np_type)) + if [py for py, _c, _d in leaves_x] != [py for py, _c, _d in leaves_y]: raise KeyError( f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"record UDTs must nest the same way, so that each has the same " - f"number of leaf fields; got {leaves_x} vs {leaves_y}." + f"record UDTs must nest the same way, with the same field names at " + f"every level; got {[c for _py, c, _d in leaves_x]} vs " + f"{[c for _py, c, _d in leaves_y]}." + ) + # Array fields combine as numpy arrays do, by broadcasting. Two shapes + # that do not broadcast raise inside the cfunc, where the caller never + # sees it, and the result keeps whatever was in the buffer. Each field's + # code is generated from the left operand's leaf, so a scalar field on + # the left of an array one never compiles, and the arithmetic ops write + # the broadcast into a field shaped like the left one, so it must fit + # there (``_fits_by_broadcast``, the rule the wrapper applies to a UDF's + # return; left to the wrapper, the error names a UDF the user did not + # write). ``eq`` and ``ne`` reduce the field to one bool, so they only + # need the shapes to broadcast. + from .base import _fits_by_broadcast + + for (_py, c, leaf_x), (_py_y, _c_y, leaf_y) in zip(leaves_x, leaves_y, strict=True): + try: + combined = np.broadcast_shapes(leaf_x.shape, leaf_y.shape) + except ValueError: + rule = "must have fields whose shapes broadcast together" + else: + if not leaf_x.shape and leaf_y.shape: + rule = ( + "can pair an array field with a scalar field only when the array is " + "on the left" + ) + elif op_name not in ("eq", "ne") and not _fits_by_broadcast(combined, leaf_x.shape): + rule = "must have fields that broadcast to the left field's shape" + else: + continue + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"record UDTs {rule}; field {c} is {leaf_x.shape} vs {leaf_y.shape}." ) if kind_x == "array" and detail_x != detail_y: raise KeyError( diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index c0b90d6bd..981ad2036 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -3170,37 +3170,115 @@ def test_udt_eq_ne_rejects_incompatible_pairs(): @pytest.mark.skipif("not supports_udfs") # SS < 9 has no GrB_NAME setter, so registration falls back to storing the # numpy repr in the type name and warns when it does not fit in 128 chars. -# _NestDeep's repr is 142; how it serializes is not what the test is about. +# Each nested record's repr here is 141 or 142; how they serialize is not what +# the test is about. @pytest.mark.filterwarnings("ignore:UDT repr is too large") -def test_udt_record_nesting_mismatch_is_a_keyerror(): - """Records sharing field names but not nesting depth are rejected as a KeyError. - - ``_check_udt_pair`` matched on top-level names only, but the codegen pairs - operands leaf by leaf, and a field that is a sub-record on one side and a - scalar on the other contributes a different number of leaves. Without the - guard the pair reaches Numba, whose typing failure arrives as a - ``UdfParseError``: a compile error reported for what is really the same - shape disagreement its sibling checks raise ``KeyError`` for. +@pytest.mark.parametrize("op", ["plus", "eq"]) +@pytest.mark.parametrize("case", ["leaf_count", "nesting", "inner_names", "inner_order"]) +def test_udt_record_nesting_mismatch_is_a_keyerror(case, op): + """Records sharing top-level field names but nesting differently are a KeyError. + + ``_check_udt_pair`` matched on top-level names only, but the codegen reads + both operands at the left one's leaf paths. A field that is a sub-record on + one side and a scalar on the other, or sub-records with different names, + reached Numba, whose typing failure was a ``UdfParseError`` for ``plus`` + and a raw ``TypingError`` for ``eq``: compile errors for what is really the + shape disagreement its sibling checks raise ``KeyError`` for. Sub-record + fields in a different order are rejected as top-level ones are, so that + leaf ``i`` of one operand is leaf ``i`` of the other. """ - flat = dtypes.register_anonymous( - np.dtype([("nst_a", np.float64), ("nst_b", np.float64)], align=True), "_NestFlat" - ) - nested = dtypes.register_anonymous( - np.dtype( - [ - ("nst_a", np.dtype([("nst_n1", np.float64), ("nst_n2", np.float64)])), - ("nst_b", np.float64), - ], - align=True, + f8 = np.float64 + sub = [("nst_n1", f8), ("nst_n2", f8)] + left, right = { + # A scalar field on the left is a sub-record on the right. + "leaf_count": ([("nst_a", f8), ("nst_b", f8)], [("nst_a", sub), ("nst_b", f8)]), + # Three leaves each, but nested under different fields. + "nesting": ([("nst_a", sub), ("nst_b", f8)], [("nst_a", f8), ("nst_b", sub)]), + "inner_names": ( + [("nst_a", sub), ("nst_b", f8)], + [("nst_a", [("nst_m1", f8), ("nst_m2", f8)]), ("nst_b", f8)], ), - "_NestDeep", - ) - v = Vector(flat, size=1) - v[0] = (1.0, 2.0) - w = Vector(nested, size=1) - w[0] = ((3.0, 4.0), 5.0) - with pytest.raises(KeyError, match="same number of leaf fields"): + "inner_order": ([("nst_a", sub), ("nst_b", f8)], [("nst_a", sub[::-1]), ("nst_b", f8)]), + }[case] + v = Vector(dtypes.register_anonymous(np.dtype(left, align=True), "_NestLeft"), size=1) + w = Vector(dtypes.register_anonymous(np.dtype(right, align=True), "_NestRight"), size=1) + with pytest.raises(KeyError, match="must nest the same way, with the same field names"): + v.ewise_mult(w, getattr(binary, op)).new() + + +def _record_with_field_shape(shape, name): + """Register a record UDT whose ``fsh_a`` field has ``shape`` (a scalar for ``()``).""" + field = ("fsh_a", np.float64, shape) if shape else ("fsh_a", np.float64) + return dtypes.register_anonymous(np.dtype([field, ("fsh_b", np.float64)]), name) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.parametrize("op", ["plus", "eq"]) +@pytest.mark.parametrize( + ("left", "right", "match"), + [ + ((3,), (4,), "shapes broadcast together"), + ((2, 3), (3, 2), "shapes broadcast together"), + ((), (3,), "only when the array is on the left"), + ], +) +def test_udt_record_field_shape_mismatch_is_a_keyerror(left, right, match, op): + """Record array fields that cannot be combined are a KeyError. + + The codegen combines array fields as numpy does, by broadcasting, so a + ``(3,)`` field with a ``(4,)`` one raised inside the cfunc, where the + caller never saw it, and the result element kept whatever was in the + buffer (``eq`` read it as its bool). A scalar field on the left of an + array one failed Numba's typing instead. + """ + v = Vector(_record_with_field_shape(left, "_FshLeft"), size=1) + w = Vector(_record_with_field_shape(right, "_FshRight"), size=1) + with pytest.raises(KeyError, match=match): + v.ewise_mult(w, getattr(binary, op)).new() + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.parametrize(("left", "right"), [((1,), (3,)), ((3,), (3, 1)), ((3, 1), (3,))]) +def test_udt_record_field_broadcast_exceeding_left_is_a_keyerror(left, right): + """An arithmetic op writes the broadcast into a field shaped like the left one. + + ``_check_udf_fills_output`` caught these pairs when the wrapper was built, + but as "UDT UDF returned an array of shape (3,)" on sample values, naming a + UDF the user did not write. ``eq`` reduces the field to one bool, so it + takes the same pair. + """ + v = Vector(_record_with_field_shape(left, "_FshLeft"), size=1) + w = Vector(_record_with_field_shape(right, "_FshRight"), size=1) + with pytest.raises(KeyError, match="broadcast to the left field's shape"): v.ewise_mult(w, binary.plus).new() + v[0] = (np.ones(left), 1.0) + w[0] = (np.ones(right), 1.0) + assert binary.eq(v & w).new().isequal(Vector.from_coo([0], [True])) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.parametrize( + ("left", "right"), [((3,), ()), ((3,), (1,)), ((3,), (1, 3)), ((2, 3), (3,))] +) +def test_udt_record_field_shapes_that_broadcast(left, right): + """Record array fields that broadcast together still combine as numpy arrays do.""" + X = _record_with_field_shape(left, "_FshLeft") + xa = np.arange(1.0, 1.0 + np.prod(left)).reshape(left) + v = Vector(X, size=2) + v[0] = (xa, 2.0) + v[1] = (np.full(left, 10.0), 2.0) + w = Vector(_record_with_field_shape(right, "_FshRight"), size=2) + w[0] = (np.full(right, 10.0), 2.0) + w[1] = (np.full(right, 10.0), 2.0) + total = v.ewise_mult(w, binary.plus).new() + assert total.dtype == X + expected = np.empty(left) + expected[...] = xa + np.full(right, 10.0) + first = total[0].new().value + np.testing.assert_array_equal(first["fsh_a"], expected) + assert first["fsh_b"] == 4.0 + # Element 0 differs only in ``fsh_a``, so its False comes from the arrays. + assert binary.eq(v & w).new().isequal(Vector.from_coo([0, 1], [False, True])) @pytest.mark.skipif("not supports_udfs") From 4210418933f021f134807291342f775ff66efd7f Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 7 Oct 2026 00:07:47 -0500 Subject: [PATCH 04/54] Type UDT ops like built-in ops, cast UDT stores, broadcast array UDTs Lifted ops (plus, minus, times, truediv, floordiv, min, max, eq, ne, ainv, abs) type each UDT element or record leaf as the built-in op types that dtype, so results never narrow. A mixed pair gives an operand's type if it holds the result, else the promoted layout. Records pair on field names, order, nesting and field shapes; array UDTs broadcast as numpy arrays do. A Python number beside a UDT is weak per leaf (NEP 50), the same on numpy 1 and 2. A literal that must become a UDT element (tuple, array, user-op literal, ewise_union default) is converted only if no value changes. apply and ewise ops with a Monoid on a mixed pair use its BinaryOp. Storing into an object of another UDT casts each leaf as GraphBLAS casts built-ins; incompatible layouts and UDT/built-in pairs raise DomainMismatch. Assigning or setting an element from another UDT used to read raw bytes, past the end of a smaller element. ewise_add converts operands to the result type, which GraphBLAS cannot cast. Array codegen uses loops, in Numba and in C: typing an op on FP64[64,64] took ~112 s and now takes ~0.02 s. Also fixes several errors, eq on 2-D array fields, array-field truediv on small ints, a UDF registered under a lifted op's name, and packed vs aligned records. Docs and tests updated. --- docs/user_guide/udt.rst | 92 +- graphblas/core/base.py | 77 +- graphblas/core/infix.py | 15 +- graphblas/core/matrix.py | 54 +- graphblas/core/operator/base.py | 9 +- graphblas/core/operator/binary.py | 122 +-- graphblas/core/operator/udt_utils.py | 656 +++++++++++--- graphblas/core/operator/unary.py | 21 + graphblas/core/operator/utils.py | 16 + graphblas/core/scalar.py | 270 +++++- graphblas/core/vector.py | 56 +- graphblas/tests/test_op.py | 1217 ++++++++++++++++++++++++-- graphblas/tests/test_ssjit.py | 3 +- 13 files changed, 2327 insertions(+), 281 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 61f5d0af4..5c7a597e5 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -135,21 +135,83 @@ For example: total = v.reduce(agg.sum[edge_dtype]).new() # field-wise reduce The lift is field-by-field for record UDTs and element-by-element for array -UDTs. Field types that don't support the operation raise ``KeyError`` on the -first lookup. ``binary.min``, ``binary.max``, and ``binary.floordiv`` reject -UDTs with any complex leaf (no ordering, no integer modulus); use ``plus``, -``minus``, ``times``, or ``truediv`` for complex arithmetic, or register a -custom binary op. - -The lifted binary ops also take two different UDTs. Two record UDTs pair up -when they nest the same way, with the same field names in the same order at -every level. Array fields combine as numpy arrays do, by broadcasting, and the -arithmetic ops give each result field the shape of the left operand's field, so -a ``(3,)`` field pairs with a ``(1,)`` or scalar one on its right but not on -its left. ``eq`` and ``ne`` reduce each field to one bool, so for them an array -field on the left need only broadcast with the right one. Two array UDTs pair -up when they share a base dtype and flat size. Any other pair raises -``KeyError`` on the first lookup. +UDTs. Each field or element of the result has the dtype the op gives on the +built-in dtypes, so a lifted op does not narrow: ``truediv`` on an ``INT64[3]`` +UDT gives ``FP64[3]``, as ``truediv`` on two INT64 vectors gives FP64. Field +types that don't support the operation raise ``KeyError`` on the first lookup. +``binary.min``, ``binary.max``, and ``binary.floordiv`` reject UDTs with any +complex leaf (no ordering, no integer modulus); use ``plus``, ``minus``, +``times``, or ``truediv`` for complex arithmetic, or register a custom binary +op. + +A scalar operand combines with every field or element. A vector, ``Scalar``, +numpy scalar or 0-d array of a built-in dtype keeps its dtype, so with ``v`` of +an ``INT8[3]`` UDT, ``v + np.int64(1)`` is ``INT64[3]``. A Python number is +weak, as numpy 2 treats it (NEP 50): in each field or element it takes that +field's dtype when the field's kind can hold it, and its own default dtype +(INT64, FP64 or FC64) when not. So ``v + 1`` stays ``INT8[3]``, ``v * 2.5`` +is ``FP64[3]``, and ``v + 300`` raises ``OverflowError``. A record with an ``INT8`` field and an ``FP32`` +field, plus ``0.5``, has an ``FP64`` field and an ``FP32`` one. python-graphblas +applies this rule itself, so it does not depend on the installed numpy (vectors +of built-in dtypes still type a Python number through numpy, which differs +between numpy 1 and 2). ``eq`` and ``ne`` compare with the number as it is. + +A literal is converted to the UDT instead when it is a whole element (a tuple, +a list, or a numpy array that is not 0-d), when the op is user-defined or a +Monoid, and for the defaults of ``ewise_union``, which must have the op's input +types (a default given as a Scalar of another type is converted too). +``apply`` with a Monoid applies its BinaryOp, so there a literal is typed as +that BinaryOp types it. The conversion must keep every value: +``v + (1, 2, 3)`` works, but ``v + (0.5, 1.5, 2.5)`` raises ``ValueError`` +rather than adding ``(0, 1, 2)``, and ``v.ewise_union(w, binary.plus, 0.5, 0)`` +raises rather than using ``0`` for a missing ``v`` entry. A number may still +round to a narrower float field, as it does beside an ``FP32`` field above. + +The lifted binary ops also take two different UDTs whose shapes match. Two +record UDTs match when they nest the same way, with the same field names in the +same order at every level, and each array field has the same shape in both; +array fields do not broadcast against each other. Two array UDTs match when +their shapes broadcast together as numpy arrays do, and the result has the +broadcast shape: ``FP64[3]`` pairs with ``INT64[2, 3]`` to give ``FP64[2, 3]``, +and ``FP64[3, 1]`` with ``FP64[1, 4]`` gives ``FP64[3, 4]``. ``eq`` is True +when every pair of elements numpy would compare is equal. Any other pair, such +as ``FP64[2, 3]`` with ``FP64[3, 2]``, raises ``KeyError`` on the first lookup. +Field and element dtypes may differ, and each pair of them follows the +built-in dtype rules; a dtype that is not a built-in one, such as bytes, pairs +only with itself. A Monoid takes one type, so ``ewise_add``, ``ewise_mult`` and +``ewise_union`` with a Monoid on two different UDTs use its BinaryOp. + +The result is an operand's UDT when that type holds it, on either side. +Otherwise it is the promoted type: for array UDTs, the structural type such as +``FP64[3]``, and for records, the record with the same names and nesting and +the promoted field dtypes, so ``{"x": INT64, "y": FP64}`` plus +``{"x": FP64, "y": INT64}`` gives ``{"x": FP64, "y": FP64}``. That is the UDT +registered with that layout if there is one, else an anonymous UDT. ``eq`` and +``ne`` compare the values and return ``BOOL``. + +A result stored in an object of another UDT type is cast to that type, as +GraphBLAS casts built-in dtypes when it stores a result: the object's type is +the request. GraphBLAS cannot cast a UDT itself, so python-graphblas casts each +element or record leaf the way GraphBLAS casts that pair of built-in dtypes. +Integers wrap; a float stored as an integer is truncated toward zero, saturates +at the integer type's bounds, and is 0 if NaN; a complex number loses its +imaginary part; and anything nonzero is ``True``. So +``int_udt_vec << int_udt_vec + 0.5``, ``int_udt_vec += 0.5`` and +``(int_udt_vec + 0.5).new(dtype=int_udt)`` truncate, as they do for an ``INT8`` +vector, and so does assigning, accumulating or masking into an object of +another UDT type. The two types must correspond as the operands of a lifted op +must, above (same field names and nesting, same shapes). numpy also casts +records by position whatever their names, and arrays to another length, but +those raise ``DomainMismatch`` here, as does storing a UDT result in an object +of a built-in dtype or the reverse. The cast is an extra pass over the result, +which is first computed in its own type. ``ewise_add`` has a related limit: +GraphBLAS casts a value that has no partner to the op's output type, so it +takes UDT operands only when both have the result's type. With a lifted +arithmetic op, it converts each operand to the result type first, which changes +no value, and applies the op there. So ``v + w`` and +``v.ewise_add(w, binary.max)`` work for two UDTs of different types, and +``truediv`` works on an integer UDT. With any other op it raises +``DomainMismatch``; use ``ewise_union`` with a default for each side instead. Composite aggregators (``agg.hypot``, ``agg.L1norm``, ``agg.Linfnorm``, ``agg.sum_of_squares``, ``agg.sum_of_inverses``) do *not* auto-lift to UDTs; diff --git a/graphblas/core/base.py b/graphblas/core/base.py index 15f66bc2f..04fb0ae04 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -162,6 +162,59 @@ def _expect_op(self, op, values, *, within, **kwargs): raise TypeError(message) from None +# Whether a UDT result may be stored in an object of another UDT type, cast +# element by element as GraphBLAS casts built-in types. False makes every +# store that would need a UDT cast raise DomainMismatch. +_udt_store_casts = True + + +def _store_cast_op(result_dtype, output_dtype): + """Return the UnaryOp that converts a result for an object of ``output_dtype``, or ``None``. + + GraphBLAS casts between built-in types on store (C's casts, with a float + saturating into an integer type), but it cannot cast to or from a UDT; + left to GraphBLAS, that is a bare GrB_DOMAIN_MISMATCH or, for a single + element, bytes read as the wrong type. So when a UDT is involved, the + result is stored by applying the op returned here, which casts each + element or record leaf the way GraphBLAS casts built-in types, between + UDTs whose layouts correspond (``udt_utils._udt_cast_error``). Any other + pair raises ``DomainMismatch``. ``None`` means GraphBLAS stores the result + as it is. This is the one place that decides what may be stored where. + """ + if result_dtype is output_dtype or not (result_dtype._is_udt or output_dtype._is_udt): + return None + from . import _has_numba + + reason = "UDTs do not cast" + if _udt_store_casts and _has_numba and result_dtype._is_udt and output_dtype._is_udt: + from .operator.udt_utils import _udt_cast_error + + if (reason := _udt_cast_error(result_dtype, output_dtype)) is None: + from .operator.unary import _udt_cast_op + + return _udt_cast_op(result_dtype, output_dtype) + elif _udt_store_casts and _has_numba: + reason = "a UDT casts only to another UDT" + from ..exceptions import DomainMismatch + + raise DomainMismatch( + f"cannot store a result of type {result_dtype} in an object of type " + f"{output_dtype}: {reason}. Use .new() for an object of the result's type, or " + "compute the result in the type you want to store." + ) + + +def _scalar_accum_op(scalar, accum): + """Return the op to accumulate into ``scalar`` with, as ``scalar.ewise_add`` runs it. + + ``__call__`` types ``accum`` on the object's dtype. For a UDT Scalar, that + would convert a literal into the UDT (``s += 0.5``), where a Vector + computes ``s + 0.5`` in its own type and casts as it stores; the untyped + op does the same here. + """ + return accum.parent if scalar.dtype._is_udt else accum + + AmbiguousAssignOrExtract._expect_op = _expect_op AmbiguousAssignOrExtract._expect_type = _expect_type @@ -341,7 +394,7 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * if expr._is_scalar and self._is_scalar: # Extract element (s << v[1]) if accum is not None: - self(**opts) << self.ewise_add(expr, accum) + self(**opts) << self.ewise_add(expr, _scalar_accum_op(self, accum)) return expr.parent._extract_element( expr.resolved_indexes, @@ -362,9 +415,13 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * expr = expr._extract_delayed() elif type(expr) is type(self): # Simple assignment (w << v) + if (cast_op := _store_cast_op(expr.dtype, self.dtype)) is not None: + expr = expr.apply(cast_op) + self._update(expr, mask, accum, replace, input_mask, opts=opts) + return if self._is_scalar: if accum is not None: - self(**opts) << self.ewise_add(expr, accum) + self(**opts) << self.ewise_add(expr, _scalar_accum_op(self, accum)) return if opts: # Ignore opts for now @@ -386,7 +443,7 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * # s << (v @ v) expr = expr._to_expr() elif accum is not None: - self(**opts) << self.ewise_add(expr, accum) + self(**opts) << self.ewise_add(expr, _scalar_accum_op(self, accum)) return else: if opts: @@ -446,6 +503,20 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * if input_mask is not None: raise TypeError("`input_mask` argument may only be used for extract") + # With an accumulator too: it is typed on this object's dtype, so the + # result still has to be that type. Before the Aggregator branch, which + # hands GraphBLAS the output as it is: agg.first on another UDT read its + # bytes as this object's type. + if (cast_op := _store_cast_op(expr.dtype, self.dtype)) is not None: + # Compute in the result's type, then cast as the result is stored, + # so the mask, accumulator and replace apply to the cast values. + # Masking the temporary too only skips values the store drops. + if self._is_scalar: + result = expr.new(name="s_result", **opts) + else: + result = expr.new(mask=mask, name="result", **opts) + self._update(result.apply(cast_op), mask, accum, replace, opts=opts) + return if expr.op is not None and expr.op.opclass == "Aggregator": updater = self(mask=mask, accum=accum, replace=replace, **opts) expr.op._new(updater, expr) diff --git a/graphblas/core/infix.py b/graphblas/core/infix.py index 73b3b91b5..bc66836fb 100644 --- a/graphblas/core/infix.py +++ b/graphblas/core/infix.py @@ -1,6 +1,6 @@ from .. import backend, binary from ..dtypes import BOOL -from ..exceptions import DimensionMismatch +from ..exceptions import DimensionMismatch, DomainMismatch from ..monoid import land, lor from ..semiring import any_pair from . import automethods, recorder, utils @@ -492,9 +492,16 @@ def _dummy(obj, obj_type): def _mismatched(left, right, method, op): # Create dummy expression to raise on incompatible dimensions - getattr(_dummy(left) if isinstance(left, InfixExprBase) else left, method)( - _dummy(right) if isinstance(right, InfixExprBase) else right, op - ) + try: + getattr(_dummy(left) if isinstance(left, InfixExprBase) else left, method)( + _dummy(right) if isinstance(right, InfixExprBase) else right, op + ) + except DomainMismatch: + # GraphBLAS reports a UDT that does not cast before the shapes, but the + # shapes are what this is here to report. + raise DimensionMismatch( + f"Dimensions not compatible for {method}: {left.shape} and {right.shape}" + ) from None raise DimensionMismatch # pragma: no cover diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 4625e4258..4e0f7e895 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -26,6 +26,12 @@ ScalarExpression, ScalarIndexExpr, _as_scalar, + _cast_for_store, + _check_literal_fits, + _check_scalar_fits, + _ewise_add_converts_operands, + _ewise_add_in_result_type, + _literal_dtype, _scalar_index, ) from .utils import ( @@ -1921,14 +1927,29 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): # Per the spec, op may be a semiring, but this is weird, so don't. self._expect_op(op, ("BinaryOp", "Monoid"), within=method_name, argname="op") - if other.ndim == 1: + if other.ndim == 1 and self._ncols != other._size: # Broadcast rowwise from the right - if self._ncols != other._size: - raise DimensionMismatch( - "Dimensions not compatible for broadcasting Vector from the right " - f"to rows of Matrix in {method_name}. Matrix.ncols (={self._ncols}) " - f"must equal Vector.size (={other._size})." - ) + raise DimensionMismatch( + "Dimensions not compatible for broadcasting Vector from the right " + f"to rows of Matrix in {method_name}. Matrix.ncols (={self._ncols}) " + f"must equal Vector.size (={other._size})." + ) + if _ewise_add_converts_operands(op, self.dtype, other.dtype): + # See Vector._ewise_add: convert both operands to the result type, + # then apply the op in that one type, since GraphBLAS cannot cast a UDT. + expr = MatrixExpression( + method_name, + None, + [self, other, _ewise_add_in_result_type, (self, other, op)], + nrows=self._nrows, + ncols=self._ncols, + op=op, + ) + if other.ndim == 2 and self.shape != other.shape: + expr.new(name="") # incompatible shape; raise now + return expr + + if other.ndim == 1: return MatrixExpression( method_name, None, @@ -2102,6 +2123,8 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): left_dtype = temp_op.type dtype = left_dtype if left_dtype._is_udt else None if type(left_default) is not Scalar: + if dtype is not None: + _check_literal_fits(dtype, left_default) try: left = Scalar.from_value( left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -2116,10 +2139,13 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): op=op, ) else: + _check_scalar_fits(dtype, left_default) left = _as_scalar(left_default, dtype, is_cscalar=False) # pragma: is_grbscalar right_dtype = temp_op.type2 dtype = right_dtype if right_dtype._is_udt else None if type(right_default) is not Scalar: + if dtype is not None: + _check_literal_fits(dtype, right_default) try: right = Scalar.from_value( right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -2134,6 +2160,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): op=op, ) else: + _check_scalar_fits(dtype, right_default) right = _as_scalar(right_default, dtype, is_cscalar=False) # pragma: is_grbscalar if is_infix: @@ -2419,6 +2446,11 @@ def apply(self, op, right=None, *, left=None): if isinstance(op, str): op = op_from_string(op) op, opclass = find_opclass(op) + if opclass == "Monoid" and self.dtype._is_udt: + # Applying a Monoid applies its BinaryOp. Typing it as that op lets a + # literal promote as it would with the BinaryOp, where a Monoid takes + # one type and would need the literal converted into the UDT. + op, opclass = op.binaryop, "BinaryOp" if opclass in {"IndexUnaryOp", "SelectOp"}: # Provide default value for index unary if right is None: @@ -2440,7 +2472,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = None elif right is None: if type(left) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_dtype(self.dtype, left, op) try: left = Scalar.from_value(left, dtype, is_cscalar=None, name="") except TypeError: @@ -2476,7 +2508,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = "{1.name}.apply({op}, left={0._expr_name})" elif left is None: if type(right) is not Scalar: - dtype = self.dtype if (self.dtype._is_udt and not op.is_positional) else None + dtype = _literal_dtype(self.dtype, right, op) try: right = Scalar.from_value(right, dtype, is_cscalar=None, name="") except TypeError: @@ -3087,6 +3119,8 @@ def _prep_for_extract(self, resolved_indexes): def _assign_element(self, resolved_indexes, value): rowidx, colidx = resolved_indexes.indices + if type(value) is Scalar: + value = _cast_for_store(value, self.dtype) if type(value) is not Scalar: dtype = self.dtype if self.dtype._is_udt else None try: @@ -3115,6 +3149,8 @@ def _assign_element(self, resolved_indexes, value): def _prep_for_assign(self, resolved_indexes, value, mask, is_submask, replace, opts): method_name = "__setitem__" + if output_type(value) in {Scalar, Vector, Matrix}: + value = _cast_for_store(value, self.dtype) rowidx, colidx = resolved_indexes.indices rowsize = rowidx.size rows = rowidx.index diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index e4c41b95e..cb8ba3746 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -192,7 +192,14 @@ def _finalize_udt_op(parent_op, dtype, dtype2, ret_type, wrapper, wrapper_sig, t from ``typed_user_cls.opclass``. ``dtype2`` is ``None`` for unary ops; the rest pass both. Returns the cached ``TypedUser*Op``. """ - wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) + try: + wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) + except NumbaError as exc: + dtypes_str = str(dtype) if dtype2 is None else f"{dtype}, {dtype2}" + raise UdfParseError( + f"{parent_op._modname}.{parent_op.name} does not work with ({dtypes_str}): " + f"{_summarize_numba_typing_error(exc)}" + ) from exc c_typename = _GB_OBJ_C_TYPENAME[typed_user_cls.opclass] error_label = c_typename.removeprefix("GrB_").removeprefix("GxB_") gb_obj = ffi.new(f"{c_typename}*") diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 2cdab6eff..43a64110b 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -64,7 +64,12 @@ def _make_udt_comparison(dtype, dtype2, *, is_eq): UDT. Mismatched-shape UDT pairs raise ``KeyError``. Returns ``(wrapper_func, wrapper_sig)``. """ - from .udt_utils import _check_udt_pair, _get_udt_info, _iter_record_leaves + from .udt_utils import ( + _array_operand_source, + _check_udt_pair, + _get_udt_info, + _iter_record_leaves, + ) info_x = _get_udt_info(dtype) info_y = _get_udt_info(dtype2) @@ -81,65 +86,62 @@ def _make_udt_comparison(dtype, dtype2, *, is_eq): join = " and " if is_eq else " or " if udt_info[0] == "array": - # Array UDT: unroll element-by-element so the scalar broadcast - # case can drop ``numba.carray`` on that side entirely. - # ``_get_udt_info`` already flattens the multi-dim shape. - # The wrapper sees the UDT side as a flat pointer-to-element, - # not pointer-to-Record; passing the Record numba_type here - # makes ``numba.carray`` try to read N record-sized chunks. - _base, N = udt_info[1] - base_numba = numba.from_dtype(_base) - x_ptr_type = ( - nt.CPointer(numba.from_dtype(dtype.np_type)) - if x_is_scalar - else nt.CPointer(base_numba) + # Array UDT: a scalar side is read through its pointer, with no + # ``numba.carray``. The wrapper sees an array side as a flat + # pointer-to-element, not pointer-to-Record; passing the Record + # numba_type here makes ``numba.carray`` try to read N + # record-sized chunks. Two array UDTs may differ in base dtype, + # so each side gets its own; typing both as the left one would + # reinterpret the bytes. Their shapes broadcast, as for the + # arithmetic ops, so eq is True when every pair of elements numpy + # would compare is equal. + shape_x = None if x_is_scalar else dtype.np_type.subdtype[1] + shape_y = None if y_is_scalar else dtype2.np_type.subdtype[1] + shape = np.broadcast_shapes(*(s for s in (shape_x, shape_y) if s is not None)) + N = reduce(mul, shape) + x_ptr_type = nt.CPointer( + numba.from_dtype(dtype.np_type if x_is_scalar else info_x[1][0]) ) - y_ptr_type = ( - nt.CPointer(numba.from_dtype(dtype2.np_type)) - if y_is_scalar - else nt.CPointer(base_numba) + y_ptr_type = nt.CPointer( + numba.from_dtype(dtype2.np_type if y_is_scalar else info_y[1][0]) ) wrapper_sig = nt.void(nt.CPointer(INT8.numba_type), x_ptr_type, y_ptr_type) - if x_is_scalar: - terms = [f"(x_ptr[0] {cmp_op} y[{i}])" for i in range(N)] - arrays = f" y = numba.carray(y_ptr, {N})\n" - elif y_is_scalar: - terms = [f"(x[{i}] {cmp_op} y_ptr[0])" for i in range(N)] - arrays = f" x = numba.carray(x_ptr, {N})\n" - else: - terms = [f"(x[{i}] {cmp_op} y[{i}])" for i in range(N)] - arrays = f" x = numba.carray(x_ptr, {N})\n y = numba.carray(y_ptr, {N})\n" - body = join.join(terms) - src = f"def wrapper(z_ptr, x_ptr, y_ptr):\n{arrays} z_ptr[0] = {body}\n" + ns = {} + x_setup, x_ref = _array_operand_source("x", shape_x, shape, ns) + y_setup, y_ref = _array_operand_source("y", shape_y, shape, ns) + # A loop, not a chain of N terms, whose typing time grows faster + # than linearly with N (see ``_make_array_wrapper``). It reads every + # element rather than stopping at the first that decides, which + # lets the compiler vectorize it: faster on the equal values that + # need every element read anyway. + accumulate = "&=" if is_eq else "|=" + src = ( + f"def wrapper(z_ptr, x_ptr, y_ptr):\n{x_setup}{y_setup}" + f" z = {is_eq}\n" + f" for i in range({N}):\n" + f" z {accumulate} {x_ref} {cmp_op} {y_ref}\n" + f" z_ptr[0] = z\n" + ) wrapper = _compile_codegen( src, func_name="wrapper", source_label=f"", + extra_ns=ns, ) return wrapper, wrapper_sig - # Record UDT (possibly nested): chain leaf comparisons. For an - # array-valued sub-field, unroll its elements so the scalar side - # never sees a numpy fancy-broadcast (Numba's record access can't - # express that as a single expression). + # Record UDT (possibly nested): chain leaf comparisons. An + # array-valued sub-field compares as a whole, against the other + # operand's field or a scalar, and reduces to one bool. np_type = udt_dtype.np_type terms = [] for py_path, _c, leaf_dtype in _iter_record_leaves(np_type): + x_expr = "x_ptr[0]" if x_is_scalar else f"x[0]{py_path}" + y_expr = "y_ptr[0]" if y_is_scalar else f"y[0]{py_path}" if leaf_dtype.subdtype is not None: - _base, shape = leaf_dtype.subdtype - sub_N = int(reduce(mul, shape)) - if x_is_scalar: - subterms = [f"(x_ptr[0] {cmp_op} y[0]{py_path}[{i}])" for i in range(sub_N)] - elif y_is_scalar: - subterms = [f"(x[0]{py_path}[{i}] {cmp_op} y_ptr[0])" for i in range(sub_N)] - else: - reducer = ".all()" if is_eq else ".any()" - terms.append(f"(x[0]{py_path} {cmp_op} y[0]{py_path}){reducer}") - continue - terms.append("(" + join.join(subterms) + ")") + reducer = ".all()" if is_eq else ".any()" + terms.append(f"({x_expr} {cmp_op} {y_expr}){reducer}") else: - x_expr = "x_ptr[0]" if x_is_scalar else f"x[0]{py_path}" - y_expr = "y_ptr[0]" if y_is_scalar else f"y[0]{py_path}" terms.append(f"({x_expr} {cmp_op} {y_expr})") expr = join.join(terms) if terms else ("True" if is_eq else "False") lines = ["def wrapper(z_ptr, x_ptr, y_ptr):"] @@ -618,20 +620,32 @@ def _compile_udt(self, dtype, dtype2): if dtypes in self._udt_types: return self._udt_ops[dtypes] + # Built-in ops lift to UDTs by name. A UDF registered under one of those + # names (``BinaryOp.register_anonymous(func, "plus", is_udt=True)``) is + # anonymous, and is compiled from its own function below. + is_lifted = not self._anonymous and ( + self.name in _BUILTIN_UDT_BINARY_OPS or self.name in ("eq", "ne") + ) + # Built-in arithmetic ops set ``_udt_types = {}`` to enable UDT # dispatch, which also routes plain-scalar misses (e.g. ``plus[BOOL]``) # here. Re-raise with the legacy single-dtype message so callers like # ``mapnumpy`` see the same error as before UDT auto-lift. - if self.name in _BUILTIN_UDT_BINARY_OPS and not dtype._is_udt and not dtype2._is_udt: + if ( + is_lifted + and self.name in _BUILTIN_UDT_BINARY_OPS + and not dtype._is_udt + and not dtype2._is_udt + ): raise KeyError(f"{self.name} does not work with {dtype}") - if self.name in ("eq", "ne") and not self._anonymous and _has_numba: + if is_lifted and self.name in ("eq", "ne") and _has_numba: binary_wrapper, wrapper_sig = _make_udt_comparison( dtype, dtype2, is_eq=(self.name == "eq") ) ret_type = BOOL - elif _has_numba and self.name in _BUILTIN_UDT_BINARY_OPS: + elif _has_numba and is_lifted: # Auto-generate a per-leaf wrapper for built-in arithmetic. Most # of these names (plus, minus, times, truediv, min, max) have no # ``_numba_func`` and would fall through to the KeyError below @@ -661,9 +675,15 @@ def _compile_udt(self, dtype, dtype2): ) # Set JIT C definition for auto-generated built-in UDT ops. Only # same-type pairs apply: mixed UDT+scalar ops don't have a - # meaningful single-type C definition. ``_has_jit_set`` gates access - # to ``lib.GrB_BinaryOp_set_String``, which is absent on SS < 9. - is_jittable = self.name in _BUILTIN_UDT_BINARY_OPS or self.name in ("eq", "ne") + # meaningful single-type C definition. The arithmetic kernel also + # declares ``z`` with the operands' type, so it fits only an op whose + # result is that type too, which ``truediv`` on integers is not. + # ``_has_jit_set`` gates access to ``lib.GrB_BinaryOp_set_String``, + # which is absent on SS < 9. + if is_lifted and self.name in _BUILTIN_UDT_BINARY_OPS: + is_jittable = ret_type == dtype + else: + is_jittable = is_lifted if _has_numba and _has_jit_set and dtype == dtype2 and is_jittable: if self.name in _BUILTIN_UDT_BINARY_OPS: op._jit_c_info = set_jit_c_on_op( diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index c06dbbee5..14f67ebc7 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -10,7 +10,7 @@ import ast import itertools import linecache -from functools import reduce +from functools import lru_cache, reduce from operator import mul import numpy as np @@ -315,7 +315,8 @@ def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): ``None`` (one side is a scalar broadcast), no check fires. Without this gate the generated wrapper code raises a cryptic Numba ``TypingError`` on the first field access, or, for array fields that do not broadcast, - fails inside the cfunc without telling the caller. + fails inside the cfunc without telling the caller. Element dtypes may + differ; :func:`_udt_op_types` promotes them. """ if info_x is None or info_y is None: return @@ -330,8 +331,8 @@ def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): if detail_x != detail_y: raise KeyError( f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"record UDTs must share field names; got {list(detail_x)} vs " - f"{list(detail_y)}." + f"record UDTs must share field names, in the same order; got " + f"{list(detail_x)} vs {list(detail_y)}." ) # Matching top-level names is not enough. The codegen reads both # operands at the left one's leaf paths, so a field that is a @@ -346,47 +347,266 @@ def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): if [py for py, _c, _d in leaves_x] != [py for py, _c, _d in leaves_y]: raise KeyError( f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"record UDTs must nest the same way, with the same field names at " - f"every level; got {[c for _py, c, _d in leaves_x]} vs " + f"record UDTs must nest the same way, with the same field names in the " + f"same order at every level; got {[c for _py, c, _d in leaves_x]} vs " f"{[c for _py, c, _d in leaves_y]}." ) - # Array fields combine as numpy arrays do, by broadcasting. Two shapes - # that do not broadcast raise inside the cfunc, where the caller never - # sees it, and the result keeps whatever was in the buffer. Each field's - # code is generated from the left operand's leaf, so a scalar field on - # the left of an array one never compiles, and the arithmetic ops write - # the broadcast into a field shaped like the left one, so it must fit - # there (``_fits_by_broadcast``, the rule the wrapper applies to a UDF's - # return; left to the wrapper, the error names a UDF the user did not - # write). ``eq`` and ``ne`` reduce the field to one bool, so they only - # need the shapes to broadcast. - from .base import _fits_by_broadcast - + # Each field must also have the same shape in both (field dtypes may + # differ), as numpy requires before it promotes or compares two record + # dtypes. Shapes that do not broadcast would raise inside the cfunc, + # unseen, and leave whatever was in the output buffer. Shapes that do + # broadcast would work one way round only, since the arithmetic ops + # write into a field shaped like the left operand's: ``(3,)`` plus + # ``(1,)`` works, ``(1,)`` plus ``(3,)`` does not. for (_py, c, leaf_x), (_py_y, _c_y, leaf_y) in zip(leaves_x, leaves_y, strict=True): - try: - combined = np.broadcast_shapes(leaf_x.shape, leaf_y.shape) - except ValueError: - rule = "must have fields whose shapes broadcast together" - else: - if not leaf_x.shape and leaf_y.shape: - rule = ( - "can pair an array field with a scalar field only when the array is " - "on the left" - ) - elif op_name not in ("eq", "ne") and not _fits_by_broadcast(combined, leaf_x.shape): - rule = "must have fields that broadcast to the left field's shape" - else: - continue - raise KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"record UDTs {rule}; field {c} is {leaf_x.shape} vs {leaf_y.shape}." - ) - if kind_x == "array" and detail_x != detail_y: + if leaf_x.shape != leaf_y.shape: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"record UDTs must have the same shape for each field; field {c} is " + f"{leaf_x.shape} vs {leaf_y.shape}." + ) + return + # Array UDTs broadcast as numpy arrays do, so ``(3, 1)`` with ``(1, 4)`` + # gives ``(3, 4)``; the generated code reads each operand through a flat + # index map (``_array_operand_source``). Shapes such as ``(2, 3)`` with + # ``(3, 2)`` do not broadcast, as numpy refuses them. + (base_x, shape_x), (base_y, shape_y) = dtype.np_type.subdtype, dtype2.np_type.subdtype + try: + np.broadcast_shapes(shape_x, shape_y) + except ValueError: raise KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"array UDTs must share base dtype and flat size; " - f"got {detail_x} vs {detail_y}." + f"binary.{op_name} does not work with ({dtype}, {dtype2}): array UDTs must " + f"have shapes that broadcast together; got {base_x}{list(shape_x)} vs " + f"{base_y}{list(shape_y)}." + ) from None + + +def _udt_cast_error(dtype, to_dtype): + """Return why an element of UDT ``dtype`` cannot be stored as UDT ``to_dtype``, or ``None``. + + Records must line up as the operands of a lifted op do + (:func:`_check_udt_pair`): the same field names, order and nesting and + the same shape for each field. Arrays must have the same shape apart from + leading axes of length 1; operands broadcast, but a stored element keeps + every value in its place. numpy casts records by position whatever their + names, and repeats or drops elements of an array field of another length, + which would store values in places nobody chose. Each pair of elements + must be built-in numbers, or the same type. + """ + np_x, np_z = dtype.np_type, to_dtype.np_type + if (np_x.names is None) != (np_z.names is None): + return "a record UDT and an array UDT do not correspond" + if np_x.names is not None: + leaves_x = list(_iter_record_leaves(np_x)) + leaves_z = list(_iter_record_leaves(np_z)) + if [py for py, _c, _d in leaves_x] != [py for py, _c, _d in leaves_z]: + return "records must have the same field names, in the same order, at every level" + if any(x.shape != z.shape for (_, _, x), (_, _, z) in zip(leaves_x, leaves_z, strict=True)): + return "records must have the same shape for each field" + pairs = [(x.base, z.base) for (_, _, x), (_, _, z) in zip(leaves_x, leaves_z, strict=True)] + else: + (base_x, shape_x), (base_z, shape_z) = np_x.subdtype, np_z.subdtype + if _strip_leading_ones(shape_x) != _strip_leading_ones(shape_z): + return "array UDTs must have the same shape, apart from leading axes of length 1" + pairs = [(base_x, base_z)] + for x, z in pairs: + if x != z and not (x.kind in "biufc" and z.kind in "biufc"): + return f"elements of {x} do not cast to {z}" + return None + + +def _strip_leading_ones(shape): + i = 0 + while i < len(shape) and shape[i] == 1: + i += 1 + return shape[i:] + + +def _builtin_dtype(np_type): + """Return the built-in ``DataType`` for ``np_type``, or ``None`` if it is not one.""" + from ..dtypes import _registry + + dt = _registry.get(np_type) + return None if dt is None or dt._is_udt else dt + + +def _complex_op_error(op_name, dtype, dtype2): + return KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"this op is not defined on complex fields. Use ``binary.plus``, " + f"``minus``, ``times``, or ``truediv`` for complex element-wise " + f"arithmetic, or register a custom binary op." + ) + + +def _element_types(op_name, dtype, dtype2, elem_x, elem_y): + """Return ``(in_x, in_y, out)``, the numpy dtypes ``binary.{op_name}`` uses on one element. + + They come from the op's typing on the built-in dtypes, so each element of an + array UDT, or each record leaf, is promoted as a Vector of that dtype would + be, and ``truediv`` on integers gives floats. ``in_x`` and ``in_y`` are the + types GraphBLAS casts the operands to before it applies the op. The + generated code casts to them too, because Numba promotes differently: it + adds uint64 and int64 in int64, where GraphBLAS uses float64. ``dtype`` and + ``dtype2`` are the operands, named in errors. Two elements that are not + both built-in dtypes (bytes, datetime, records inside an array field) have + no promotion rule, so they must be the same dtype, which the result keeps. + """ + # ``min``/``max``/``floordiv`` have no defined semantics on complex + # operands (no ordering, no integer mod). Reject early with a clear + # message; otherwise Numba's cfunc compile blows up several frames + # down with ``NotImplementedError: No definition for lowering lt``. + if op_name in _OPS_NOT_FOR_COMPLEX and "c" in (elem_x.kind, elem_y.kind): + raise _complex_op_error(op_name, dtype, dtype2) + builtin_x = _builtin_dtype(elem_x) + builtin_y = _builtin_dtype(elem_y) + if builtin_x is None or builtin_y is None: + if elem_x == elem_y: + return elem_x, elem_y, elem_x + # Builds without complex dtypes (the vanilla backend) still lift complex + # fields, so promote them as the built-in typing does where it has them. + if "c" in (elem_x.kind, elem_y.kind) and {elem_x.kind, elem_y.kind} <= set("biufc"): + promoted = np.promote_types(elem_x, elem_y) + return promoted, promoted, promoted + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): elements of " + f"{elem_x} and {elem_y} have no common type." ) + from ... import binary + from .utils import get_typed_op + + typed = get_typed_op(getattr(binary, op_name), builtin_x, builtin_y) + return typed.type.np_type, typed.type2.np_type, typed.return_type.np_type + + +def _udt_op_types(op_name, dtype, dtype2): + """Return ``(ret_type, input_types)`` for ``binary.{op_name}`` on ``dtype`` and ``dtype2``. + + At least one operand is a record or array UDT, and ``_check_udt_pair`` has + accepted the pair; the other may be a plain scalar type. Each element, or + record leaf, follows :func:`_element_types`, so the result never narrows: + an int64 array UDT plus ``0.5`` is ``FP64[3]``. ``ret_type`` is an + operand's own type when that holds the result, so same-type arithmetic + keeps its UDT. Otherwise it is the promoted layout, which is the UDT a user + registered with that layout if there is one, else an anonymous one. + ``input_types`` holds the ``(in_x, in_y)`` pair of each record leaf, in + leaf order, or the single pair for an array UDT's elements. + """ + from ..dtypes import lookup_dtype + + np_x, np_y = dtype.np_type, dtype2.np_type + if np_x.names is not None or np_y.names is not None: + leaves_x = list(_iter_record_leaves(np_x)) if np_x.names is not None else None + leaves_y = list(_iter_record_leaves(np_y)) if np_y.names is not None else None + input_types = [] + result_leaves = [] + for i, (_py, _c, leaf) in enumerate(leaves_x or leaves_y): + elem_x = np_x if leaves_x is None else leaves_x[i][2].base + elem_y = np_y if leaves_y is None else leaves_y[i][2].base + in_x, in_y, elem = _element_types(op_name, dtype, dtype2, elem_x, elem_y) + input_types.append((in_x, in_y)) + result_leaves.append(np.dtype((elem, leaf.shape)) if leaf.shape else elem) + for operand, leaves in [(dtype, leaves_x), (dtype2, leaves_y)]: + if leaves is not None and all( + result_leaf == leaf + for result_leaf, (_py, _c, leaf) in zip(result_leaves, leaves, strict=True) + ): + return operand, input_types + records = [np_type for np_type in (np_x, np_y) if np_type.names is not None] + return lookup_dtype(_promoted_record(records, iter(result_leaves))), input_types + base_x, shape_x = np_x.subdtype or (np_x, ()) + base_y, shape_y = np_y.subdtype or (np_y, ()) + in_x, in_y, elem = _element_types(op_name, dtype, dtype2, base_x, base_y) + result = np.dtype((elem, np.broadcast_shapes(shape_x, shape_y))) + for operand in (dtype, dtype2): + if operand.np_type == result: + return operand, [(in_x, in_y)] + return lookup_dtype(result), [(in_x, in_y)] + + +def _promoted_record(records, result_leaves): + """Build a record laid out like ``records`` with ``result_leaves`` for its leaves. + + ``records`` holds one or two record dtypes with the same field names and + nesting. The new record aligns its fields, at each level, if any of them + does there, so a C-compatible operand gives a C-compatible result. + """ + first = records[0] + fields = [] + for name in first.names: + field = first.fields[name][0] + if field.names is not None: + field = _promoted_record([r.fields[name][0] for r in records], result_leaves) + else: + field = next(result_leaves) + fields.append((name, field)) + return np.dtype(fields, align=any(r.isalignedstruct for r in records)) + + +# The kinds of element a weak Python literal fits without changing kind, and +# the dtype it takes otherwise: numpy 2's rules for Python scalars (NEP 50). +_WEAK_LITERAL_FITS = {"b": "biufc", "i": "iufc", "f": "fc", "c": "c"} +_WEAK_LITERAL_DEFAULT = { + "b": np.dtype(np.bool_), + "i": np.dtype(np.int64), + "f": np.dtype(np.float64), + "c": np.dtype(np.complex128), +} + + +def _weak_literal_element(value, kind, elem): + """Return the dtype a Python literal of ``kind`` takes beside an element of ``elem``. + + It is ``elem`` when that kind of element holds the literal (``1`` beside + int8, ``0.5`` beside float32), else the literal's default dtype (``0.5`` + beside int8 is float64), and complex of ``elem``'s precision for a complex + literal beside a float. An int literal outside an integer element's range + raises ``OverflowError``, as numpy 2 does, rather than wrapping. + """ + if elem.kind in _WEAK_LITERAL_FITS[kind]: + if kind == "i" and elem.kind in "iu": + info = np.iinfo(elem) + if not info.min <= value <= info.max: + raise OverflowError(f"Python integer {value} out of bounds for {elem}") + return elem + if kind == "c" and elem.kind == "f": + return np.result_type(elem, np.complex64) + return _WEAK_LITERAL_DEFAULT[kind] + + +def _weak_literal_udt(dtype, value): + """Return the UDT a Python number becomes beside UDT ``dtype``: weak, element by element. + + The literal takes, in each element or record leaf, the dtype numpy 2 gives + a Python scalar beside that element (:func:`_weak_literal_element`), so the + pair then promotes as two UDTs do: ``int8_udt + 1`` stays ``int8_udt``, and + ``int8_udt * 2.5`` gives float64 elements. A leaf that is not a number + (bytes, datetime) takes the literal's default dtype and fails to pair. + """ + from ..dtypes import lookup_dtype + + kind = ( + "b" + if isinstance(value, bool) + else "i" if isinstance(value, int) else "f" if isinstance(value, float) else "c" + ) + + def element(elem): + if elem.kind not in "biufc": + return _WEAK_LITERAL_DEFAULT[kind] + return _weak_literal_element(value, kind, elem) + + np_type = dtype.np_type + if np_type.names is not None: + leaves = [] + for _py, _c, leaf in _iter_record_leaves(np_type): + elem = element(leaf.base) + leaves.append(np.dtype((elem, leaf.shape)) if leaf.shape else elem) + literal_type = _promoted_record([np_type], iter(leaves)) + else: + base, shape = np_type.subdtype + literal_type = np.dtype((element(base), shape)) + return dtype if literal_type == np_type else lookup_dtype(literal_type) def _iter_record_leaves(np_type, python_prefix="", c_prefix=""): @@ -577,7 +797,22 @@ def _expr_unary(py_op, operand): return f"{py_op}({operand})" return f"{py_op}{operand}" - def _make_record_func(leaf_paths, arity, py_op, *, x_is_scalar=False, y_is_scalar=False): + def _cast_expr(expr, to_dtype, ns, *, shaped=False): + """Return Numba source converting ``expr`` to numpy dtype ``to_dtype``. + + ``to_dtype`` of ``None`` leaves ``expr`` alone. A ``shaped`` operand is + an array-valued record field, converted element by element. The + conversion function is added to ``ns``, the generated code's namespace. + """ + if to_dtype is None: + return expr + name = f"_to_{to_dtype.name}" + ns[name] = numba.from_dtype(to_dtype) + return f"{expr}.astype({name})" if shaped else f"{name}({expr})" + + def _make_record_func( + leaf_paths, arity, py_op, *, x_is_scalar=False, y_is_scalar=False, casts=None + ): """Build a Numba njit function for a record UDT. ``leaf_paths`` is a sequence of Python access strings ``"['a']"`` or, @@ -590,12 +825,27 @@ def _make_record_func(leaf_paths, arity, py_op, *, x_is_scalar=False, y_is_scala When ``x_is_scalar`` or ``y_is_scalar`` is True, that argument is a plain scalar (not a record), so it is used directly for all leaves. + ``casts`` holds an ``(x_dtype, y_dtype, shaped)`` triple per leaf: the + numpy dtype to convert each operand to before the op (``None`` for + none), and whether the leaf is an array field. """ + ns = {} if arity == 2: parts = [] - for path in leaf_paths: - x_expr = "x" if x_is_scalar else f"x{path}" - y_expr = "y" if y_is_scalar else f"y{path}" + for i, path in enumerate(leaf_paths): + x_cast, y_cast, shaped = casts[i] if casts else (None, None, False) + x_expr = _cast_expr( + "x" if x_is_scalar else f"x{path}", + x_cast, + ns, + shaped=shaped and not x_is_scalar, + ) + y_expr = _cast_expr( + "y" if y_is_scalar else f"y{path}", + y_cast, + ns, + shaped=shaped and not y_is_scalar, + ) parts.append(_expr_binary(py_op, x_expr, y_expr)) sig = "x, y" else: @@ -609,54 +859,96 @@ def _make_record_func(leaf_paths, arity, py_op, *, x_is_scalar=False, y_is_scala src, func_name="_op", source_label=f"", + extra_ns=ns, ) return numba.njit(op_func, error_model="numpy") + def _array_operand_source(side, operand_shape, shape, ns): + """Return ``(setup, ref)``, Numba source that reads operand ``side`` at element ``i``. + + ``i`` is a flat position in a result of ``shape``. ``operand_shape`` is + the operand's shape, or ``None`` for a plain scalar, read for every + element. An operand that broadcasts to ``shape`` reads through a + constant flat index map, added to ``ns``. + """ + if operand_shape is None: + return "", f"{side}_ptr[0]" + size = reduce(mul, operand_shape) + setup = f" {side} = numba.carray({side}_ptr, {size})\n" + # Broadcasting repeats no element when the sizes agree, so the shapes + # differ at most by leading axes of length 1 and the flat order is the same. + if size == reduce(mul, shape): + return setup, f"{side}[i]" + ns[f"{side}_index"] = np.broadcast_to(np.arange(size).reshape(operand_shape), shape).ravel() + return setup, f"{side}[{side}_index[i]]" + def _make_array_wrapper( - size, + shape, base_numba_type, arity, py_op, *, - x_scalar_type=None, - y_scalar_type=None, + x_type=None, + y_type=None, + x_shape=None, + y_shape=None, + x_is_scalar=False, + y_is_scalar=False, + x_cast=None, + y_cast=None, ): """Build a cfunc-ready wrapper for an array UDT (element-by-element). - When ``x_scalar_type`` or ``y_scalar_type`` is set, that side is a plain - scalar pointer (broadcast to all elements). + ``shape`` and ``base_numba_type`` are the shape and element type of the + result. ``x_type`` and ``y_type`` are the operands' element types, and + ``x_shape`` and ``y_shape`` their shapes, which default to the result's + and broadcast to it; when ``x_is_scalar`` or ``y_is_scalar`` is set, + that side is a plain scalar of that type, broadcast to all elements. + ``x_cast`` and ``y_cast`` are numpy dtypes to convert each operand to + before the op, if any. Returns (wrapper_func, wrapper_sig). """ nt = numba.types + ns = {} + size = reduce(mul, shape) + if x_type is None: + x_type = base_numba_type + if x_shape is None: + x_shape = shape + if y_shape is None: + y_shape = shape + # A loop, not one statement per element: Numba's typing time grows + # faster than linearly with the number of statements, so unrolling took + # seconds for a 32 by 32 array and minutes for a 64 by 64 one. + arrays = f" z = numba.carray(z_ptr, {size})\n" if arity == 2: - x_ref = "x_ptr[0]" if x_scalar_type else "x[{i}]" - y_ref = "y_ptr[0]" if y_scalar_type else "y[{i}]" - assigns = "\n".join( - f" z[{i}] = {_expr_binary(py_op, x_ref.format(i=i), y_ref.format(i=i))}" - for i in range(size) - ) + if y_type is None: + y_type = base_numba_type + # Each operand is read once per element, into ``xv`` and ``yv``, + # so an op that names an operand more than once reads it once. + body = "" + for side, operand_shape, cast in [ + ("x", None if x_is_scalar else x_shape, x_cast), + ("y", None if y_is_scalar else y_shape, y_cast), + ]: + setup, ref = _array_operand_source(side, operand_shape, shape, ns) + arrays += setup + body += f" {side}v = {_cast_expr(ref, cast, ns)}\n" + body += f" z[i] = {_expr_binary(py_op, 'xv', 'yv')}\n" params = "z_ptr, x_ptr, y_ptr" - arrays = f" z = numba.carray(z_ptr, {size})\n" - if not x_scalar_type: - arrays += f" x = numba.carray(x_ptr, {size})\n" - if not y_scalar_type: - arrays += f" y = numba.carray(y_ptr, {size})\n" - x_numba = nt.CPointer(x_scalar_type) if x_scalar_type else nt.CPointer(base_numba_type) - y_numba = nt.CPointer(y_scalar_type) if y_scalar_type else nt.CPointer(base_numba_type) - sig = nt.void(nt.CPointer(base_numba_type), x_numba, y_numba) + sig = nt.void(nt.CPointer(base_numba_type), nt.CPointer(x_type), nt.CPointer(y_type)) else: - assigns = "\n".join( - f" z[{i}] = {_expr_unary(py_op, f'x[{i}]')}" for i in range(size) - ) + body = f" z[i] = {_expr_unary(py_op, 'x[i]')}\n" params = "z_ptr, x_ptr" - arrays = f" z = numba.carray(z_ptr, {size})\n x = numba.carray(x_ptr, {size})\n" - sig = nt.void(nt.CPointer(base_numba_type), nt.CPointer(base_numba_type)) - src = f"def _op({params}):\n{arrays}{assigns}\n" + arrays += f" x = numba.carray(x_ptr, {size})\n" + sig = nt.void(nt.CPointer(base_numba_type), nt.CPointer(x_type)) + src = f"def _op({params}):\n{arrays} for i in range({size}):\n{body}" op_func = _compile_codegen( src, func_name="_op", source_label=f"", + extra_ns=ns, ) return op_func, sig @@ -670,90 +962,206 @@ def compile_udt_binary_wrapper(op_name, py_op, dtype, dtype2): Handles two cases: - - Both sides are the same UDT: a field-by-field (record) or + - Both sides are UDTs of the same shape: a field-by-field (record) or element-by-element (array) op. - One side is a UDT and the other is a scalar type: the scalar is broadcast to all fields or elements (e.g., ``Point + int`` adds the int to every field). + Each element is computed as the op computes it on the operands' + element dtypes (see :func:`_udt_op_types`). + Returns ``(wrapper_func, wrapper_sig, ret_type)``. Raises ``KeyError`` when the dtype combination is not supported, with a clear message for common mistakes like passing two record UDTs with different field names. """ - from .base import _get_udt_wrapper, _resolve_udt_return_type + from .base import _get_udt_wrapper info_x = _get_udt_info(dtype) info_y = _get_udt_info(dtype2) - _check_udt_pair(op_name, dtype, dtype2, info_x, info_y) - - # Pick the UDT side. Both sides may be UDTs; the pre-check above - # ensures they share a shape in that case. - if info_x is not None: - udt_dtype = dtype - udt_info = info_x - elif info_y is not None: - udt_dtype = dtype2 - udt_info = info_y - else: + if info_x is None and info_y is None: raise KeyError( f"binary.{op_name} does not work with ({dtype}, {dtype2}). " f"Element-wise UDT ops require a record dtype (named fields) " f"or an array dtype (e.g., FP64[3])." ) - - # ``min``/``max``/``floordiv`` have no defined semantics on complex - # operands (no ordering, no integer mod). Reject early with a clear - # message; otherwise Numba's cfunc compile blows up several frames - # down with ``NotImplementedError: No definition for lowering lt``. - if not _op_supports_field_dtypes(op_name, udt_dtype.np_type): - raise KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"this op is not defined on complex fields. Use ``binary.plus``, " - f"``minus``, ``times``, or ``truediv`` for complex element-wise " - f"arithmetic, or register a custom binary op." - ) + _check_udt_pair(op_name, dtype, dtype2, info_x, info_y) + ret_type, input_types = _udt_op_types(op_name, dtype, dtype2) x_is_scalar = info_x is None # left side is a plain scalar type y_is_scalar = info_y is None # right side is a plain scalar type - kind, detail = udt_info + udt_dtype = dtype2 if x_is_scalar else dtype - if kind == "record": + # Numba promotes mixed operands its own way (uint64 with int64 is + # int64, where GraphBLAS uses float64) and only then stores the value + # into the result's type, so each operand is converted first to the + # type the built-in op computes in. Same-type operands are left alone. + def cast(in_dtype, elem): + return None if in_dtype == elem.base else in_dtype + + if udt_dtype.np_type.names is not None: from .base import _compile_udf_for_udt # Use leaf paths so the same codegen handles nested-record UDTs # uniformly. A non-nested record's leaves are its top-level # fields, with paths like ``"['a']"``. - leaf_paths = [py for py, _c, _d in _iter_record_leaves(udt_dtype.np_type)] + leaves = list(_iter_record_leaves(udt_dtype.np_type)) + leaves_x = None if x_is_scalar else list(_iter_record_leaves(dtype.np_type)) + leaves_y = None if y_is_scalar else list(_iter_record_leaves(dtype2.np_type)) + ret_leaves = list(_iter_record_leaves(ret_type.np_type)) + casts = [] + for i, ((_py, _c, leaf), (in_x, in_y)) in enumerate( + zip(leaves, input_types, strict=True) + ): + if leaf.shape and (out := ret_leaves[i][2].base) != in_x: + # An array field runs through a Numba ufunc, which picks its + # own loop: int8 / int8 runs in float32. A scalar field divides + # in float64, so convert array operands to the result's dtype. + in_x = in_y = out + casts.append( + ( + cast(in_x, dtype.np_type if x_is_scalar else leaves_x[i][2]), + cast(in_y, dtype2.np_type if y_is_scalar else leaves_y[i][2]), + bool(leaf.shape), + ) + ) func = _make_record_func( - leaf_paths, + [py for py, _c, _d in leaves], 2, py_op, x_is_scalar=x_is_scalar, y_is_scalar=y_is_scalar, + casts=casts, ) sig = (dtype.numba_type, dtype2.numba_type) _compile_udf_for_udt( func, sig, op_kind="binary", op_name=op_name, dtypes=(dtype, dtype2) ) numba_ret_type = func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, udt_dtype) + # Numba can still return a leaf wider than ``ret_type``'s (int8 + # plus int8 is int64 in Numba, INT8 here); the wrapper casts each + # leaf as it stores it, which for those ops gives the same value. wrapper, wrapper_sig = _get_udt_wrapper( func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type ) else: - base_dtype, size = detail - ret_type = udt_dtype + # Each side has its own element type, since the bases may differ. + x_type = dtype.np_type if x_is_scalar else info_x[1][0] + y_type = dtype2.np_type if y_is_scalar else info_y[1][0] + ((in_x, in_y),) = input_types + base, shape = ret_type.np_type.subdtype wrapper, wrapper_sig = _make_array_wrapper( - size, - numba.from_dtype(base_dtype), + shape, + numba.from_dtype(base), 2, py_op, - x_scalar_type=numba.from_dtype(dtype.np_type) if x_is_scalar else None, - y_scalar_type=numba.from_dtype(dtype2.np_type) if y_is_scalar else None, + x_type=numba.from_dtype(x_type), + y_type=numba.from_dtype(y_type), + x_shape=None if x_is_scalar else dtype.np_type.subdtype[1], + y_shape=None if y_is_scalar else dtype2.np_type.subdtype[1], + x_is_scalar=x_is_scalar, + y_is_scalar=y_is_scalar, + x_cast=cast(in_x, x_type), + y_cast=cast(in_y, y_type), ) return wrapper, wrapper_sig, ret_type + @lru_cache + def _saturating_cast(dst): + """Return a Numba function converting a float to integer dtype ``dst`` as GraphBLAS does. + + NaN is 0, a value beyond the type's range is its nearest bound, and + anything else truncates toward zero. Numba's own conversion is LLVM's, + which leaves the first two undefined, and numpy's differs by platform. + """ + info = np.iinfo(dst) + lo = "0.0" if dst.kind == "u" else repr(float(info.min)) + src = ( + "def _sat(x):\n" + " x = _f64(x)\n" + f" if x != x or x <= {lo}:\n" + f" return _to(0 if x != x else {info.min})\n" + f" if x >= {float(info.max)!r}:\n" + f" return _to({info.max})\n" + " return _to(x)\n" + ) + func = _compile_codegen( + src, + func_name="_sat", + source_label=f"", + extra_ns={"_f64": numba.float64, "_to": numba.from_dtype(dst)}, + ) + return numba.njit(func) + + def _cast_leaf_source(src, dst, expr, ns): + """Return Numba source converting ``expr`` from numpy dtype ``src`` to ``dst``. + + It converts as GraphBLAS casts built-in types on store: anything + nonzero is True, a complex number loses its imaginary part, a float + stored as an integer saturates (:func:`_saturating_cast`), and + integers wrap. Helpers are added to ``ns``, the generated code's + namespace. + """ + if src == dst: + return expr + if dst.kind == "b": + return f"({expr} != 0)" + if src.kind == "c" and dst.kind != "c": + expr = f"({expr}).real" + src = np.dtype(f"f{src.itemsize // 2}") + if src.kind == "f" and dst.kind in "iu": + ns[f"_sat_{dst.name}"] = _saturating_cast(dst) + return f"_sat_{dst.name}({expr})" + ns[f"_to_{dst.name}"] = numba.from_dtype(dst) + return f"_to_{dst.name}({expr})" + + def compile_udt_cast_wrapper(dtype, to_dtype): + """Compile a cfunc wrapper that stores an element of UDT ``dtype`` as UDT ``to_dtype``. + + :func:`_udt_cast_error` has accepted the pair. Each element of an array + UDT, and each record leaf, converts by :func:`_cast_leaf_source`. + + Returns (wrapper_func, wrapper_sig). + """ + nt = numba.types + ns = {} + np_x, np_z = dtype.np_type, to_dtype.np_type + if np_x.names is not None: + body = [" x = numba.carray(x_ptr, 1)", " z = numba.carray(z_ptr, 1)"] + for (path, _c, leaf_x), (_path, _cz, leaf_z) in zip( + _iter_record_leaves(np_x), _iter_record_leaves(np_z), strict=True + ): + if leaf_x.shape: + value = _cast_leaf_source(leaf_x.base, leaf_z.base, f"x[0]{path}[i]", ns) + body.append(f" for i in np.ndindex({leaf_x.shape}):") + body.append(f" z[0]{path}[i] = {value}") + else: + value = _cast_leaf_source(leaf_x, leaf_z, f"x[0]{path}", ns) + body.append(f" z[0]{path} = {value}") + sig = nt.void(nt.CPointer(to_dtype.numba_type), nt.CPointer(dtype.numba_type)) + else: + (base_x, shape), base_z = np_x.subdtype, np_z.subdtype[0] + size = reduce(mul, shape) + body = [ + f" x = numba.carray(x_ptr, {size})", + f" z = numba.carray(z_ptr, {size})", + f" for i in range({size}):", + f" z[i] = {_cast_leaf_source(base_x, base_z, 'x[i]', ns)}", + ] + sig = nt.void( + nt.CPointer(numba.from_dtype(base_z)), nt.CPointer(numba.from_dtype(base_x)) + ) + ns["np"] = np + src = "def _cast(z_ptr, x_ptr):\n" + "\n".join(body) + "\n" + wrapper = _compile_codegen( + src, + func_name="_cast", + source_label=f"", + extra_ns=ns, + ) + return wrapper, sig + def compile_udt_unary_wrapper(op_name, py_op, dtype): """Compile a built-in element-wise unary op for a UDT. @@ -770,7 +1178,7 @@ def compile_udt_unary_wrapper(op_name, py_op, dtype): f"or an array dtype (e.g., FP64[3])." ) - kind, detail = info + kind, _detail = info if kind == "record": from .base import _compile_udf_for_udt @@ -784,9 +1192,11 @@ def compile_udt_unary_wrapper(op_name, py_op, dtype): func, ret_type, dtype, numba_ret_type=numba_ret_type ) else: - base_dtype, size = detail + base_dtype, shape = dtype.np_type.subdtype ret_type = dtype - wrapper, wrapper_sig = _make_array_wrapper(size, numba.from_dtype(base_dtype), 1, py_op) + wrapper, wrapper_sig = _make_array_wrapper( + shape, numba.from_dtype(base_dtype), 1, py_op + ) return wrapper, wrapper_sig, ret_type @@ -924,16 +1334,13 @@ def _make_jit_c_definition(op_name, py_op, dtype, arity): else: # array UDT, flattened to v[size] base_dtype, shape = np_type.subdtype size = reduce(mul, shape) + # A loop, as in the Numba wrapper: one statement per element made the + # C compiler take seconds on a 64 by 64 array, once per JIT cache. if arity == 2: - assigns = " ".join( - f"z->v[{i}] = {_c_expr_binary(py_op, f'x->v[{i}]', f'y->v[{i}]', base_dtype)} ;" - for i in range(size) - ) + expr = _c_expr_binary(py_op, "x->v[i]", "y->v[i]", base_dtype) else: - assigns = " ".join( - f"z->v[{i}] = {_c_expr_unary(py_op, f'x->v[{i}]', base_dtype)} ;" - for i in range(size) - ) + expr = _c_expr_unary(py_op, "x->v[i]", base_dtype) + assigns = f"for (int64_t i = 0 ; i < {size} ; i++) {{ z->v[i] = {expr} ; }}" return c_name, f"void {c_name} ({params}) {{ {assigns} }}" @@ -945,8 +1352,8 @@ def _make_jit_c_comparison_definition(op_name, dtype, *, is_eq): expressed in C. The kernel signature is ``void op(_Bool *z, const Udt *x, const Udt *y)``: each leaf field contributes a scalar ``==`` (or ``!=``) comparison; record-UDT leaves - are chained with ``&&`` (eq) or ``||`` (ne). Top-level array UDTs unroll - their elements at codegen time. + are chained with ``&&`` (eq) or ``||`` (ne). Top-level array UDTs loop + over their elements. Array-valued sub-fields inside a record (e.g. ``[("weights", (float64, 3))]``) aren't supported: ``_udt_c_typedef`` @@ -974,17 +1381,22 @@ def _make_jit_c_comparison_definition(op_name, dtype, *, is_eq): op = "==" if is_eq else "!=" join = " && " if is_eq else " || " if np_type.subdtype is not None: - # Array UDT: unroll all elements. + # Array UDT: a loop over every element (see ``_make_jit_c_definition`` + # and the Numba wrapper in ``binary._make_udt_comparison``). _base, shape = np_type.subdtype size = reduce(mul, shape) - terms = [f"((x->v[{i}]) {op} (y->v[{i}]))" for i in range(size)] + accumulate = "&=" if is_eq else "|=" + body = ( + f"_Bool r = {int(is_eq)} ; for (int64_t i = 0 ; i < {size} ; i++) " + f"{{ r {accumulate} ((x->v[i]) {op} (y->v[i])) ; }} *z = r ;" + ) elif np_type.names is not None: terms = [f"((x->{c}) {op} (y->{c}))" for _py, c, _d in _iter_record_leaves(np_type)] + body = f"*z = {join.join(terms) if terms else int(is_eq)} ;" else: return None - body = join.join(terms) if terms else ("1" if is_eq else "0") params = f"_Bool *z, const {type_name} *x, const {type_name} *y" - return c_name, f"void {c_name} ({params}) {{ *z = {body} ; }}" + return c_name, f"void {c_name} ({params}) {{ {body} }}" def _maybe_warn_jit_skipped(jit_info, op_name, dtype_name): diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index fc9c92661..f5bf0fc44 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -124,6 +124,7 @@ def _one(x): from .udt_utils import ( _has_jit_set, _maybe_warn_jit_skipped, + compile_udt_cast_wrapper, compile_udt_unary_wrapper, set_jit_c_on_op, ) @@ -494,3 +495,23 @@ def __init__( ParameterizedUnaryOp._op_class = UnaryOp + + +# {to_dtype: UnaryOp}; each one's typed ops are keyed by the type they cast from. +_udt_casts = {} + + +def _udt_cast_op(dtype, to_dtype): + """Return the UnaryOp storing an element of UDT ``dtype`` as UDT ``to_dtype``. + + ``base._store_cast_op`` decides when to use it, after + ``udt_utils._udt_cast_error`` has accepted the pair. It is compiled once + per pair and has no JIT kernel. + """ + parent = _udt_casts.get(to_dtype) + if parent is None: + parent = _udt_casts[to_dtype] = UnaryOp(f"cast_to_{to_dtype}", anonymous=True, is_udt=True) + if dtype in parent._udt_ops: + return parent._udt_ops[dtype] + wrapper, wrapper_sig = compile_udt_cast_wrapper(dtype, to_dtype) + return _finalize_udt_op(parent, dtype, None, to_dtype, wrapper, wrapper_sig, TypedUserUnaryOp) diff --git a/graphblas/core/operator/utils.py b/graphblas/core/operator/utils.py index 6f2df5535..2f9dc04e4 100644 --- a/graphblas/core/operator/utils.py +++ b/graphblas/core/operator/utils.py @@ -68,6 +68,22 @@ def get_typed_op(op, dtype, dtype2=None, *, is_left_scalar=False, is_right_scala # Handle special cases such as first and second (may have UDTs) if op._custom_dtype is not None and (rv := op._custom_dtype(op, dtype, dtype2)) is not None: return rv + # UDTs do not cast, so never unify a UDT with another type: numpy would + # promote two records to one of them (or a third), and GraphBLAS would + # reject the operand of the other type with GrB_DOMAIN_MISMATCH. An op + # that does not lift to UDTs says so here with a KeyError, where a UDT + # with a built-in type mostly failed to unify with numpy's TypeError. + if dtype is not dtype2 and (dtype._is_udt or dtype2._is_udt): + if op.is_positional: + return op[UINT64] + if op._udt_types is None: + raise KeyError(f"{op.name} does not work with ({dtype}, {dtype2})") + if kind == "binary" and isinstance(op, Monoid) and dtype != dtype2: + # A Monoid takes one type, but used element-wise it is its + # BinaryOp, which takes the pair: ``v.ewise_add(w)`` works where + # ``v + w`` does. Indexing (``monoid.plus[X, Y]``) still raises. + return get_typed_op(op.binaryop, dtype, dtype2) + return op._compile_udt(dtype, dtype2) # Generic case: try to unify the two dtypes try: return op[ diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index ce191f374..14690d231 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -468,7 +468,11 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=None, name=None): else: new_scalar = Scalar(dtype, is_cscalar=is_cscalar, name=name) if not clear and not self._is_empty: - if new_scalar.is_cscalar and not new_scalar.dtype._is_udt: + if self.dtype._is_udt: + # Cast as storing it would (``base._store_cast_op``), not as numpy + # does; a UDT does not cast to a built-in dtype at all. + new_scalar << self + elif new_scalar.is_cscalar and not new_scalar.dtype._is_udt: # Cast value so we don't raise given explicit dup with dtype new_scalar.value = new_scalar.dtype.np_type.type(self.value) else: @@ -668,7 +672,7 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): other = op(other).new() if type(other) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_dtype(self.dtype, other, op) try: other = Scalar.from_value(other, dtype, is_cscalar=False, name="") except TypeError: @@ -682,6 +686,16 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): ) op = get_typed_op(op, self.dtype, other.dtype, kind="binary") self._expect_op(op, ("BinaryOp", "Monoid"), within=method_name, argname="op") + if _ewise_add_converts_operands(op, self.dtype, other.dtype): + # GraphBLAS casts an unpaired entry to the result type, which a UDT + # cannot be. Convert both to the result type first, which changes no + # value, and apply the op in that one type: an int8 UDT plus 0.5 is + # float64, and an empty Scalar still adds as 0. + if self.dtype != op.return_type: + self = _converted_to(self, op.return_type, name="") + if other.dtype != op.return_type: + other = _converted_to(other, op.return_type, name="") + op = get_typed_op(op.parent, self.dtype, other.dtype, kind="binary") return ScalarExpression( method_name, f"GrB_Vector_eWiseAdd_{op.opclass}", @@ -739,7 +753,7 @@ def _ewise_mult(self, other, op=binary.times, is_infix=False): other = op(other).new() if type(other) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_dtype(self.dtype, other, op) try: other = Scalar.from_value(other, dtype, is_cscalar=False, name="") except TypeError: @@ -817,7 +831,9 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): dtype = right_dtype if right_dtype._is_udt else None if type(other) is not Scalar: try: - other = Scalar.from_value(other, dtype, is_cscalar=False, name="") + other = Scalar.from_value( + other, _literal_dtype(right_dtype, other, op), is_cscalar=False, name="" + ) except TypeError: other = self._expect_type( other, @@ -828,6 +844,11 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): op=op, ) else: + # A typed Scalar keeps its dtype beside a lifted op, as it does for + # Vector.ewise_union: ``s8 - sf`` is float, not ``sf`` cut to int8. + if dtype is not None and other.dtype != dtype and _is_lifted_op(op): + dtype = None + _check_scalar_fits(dtype, other) other = _as_scalar(other, dtype, is_cscalar=False) # pragma: is_grbscalar temp_op = get_typed_op(op, self.dtype, other.dtype, kind="binary") @@ -835,6 +856,8 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): left_dtype = temp_op.type dtype = left_dtype if left_dtype._is_udt else None if type(left_default) is not Scalar: + if dtype is not None: + _check_literal_fits(dtype, left_default) try: left = Scalar.from_value( left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -849,10 +872,13 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): op=op, ) else: + _check_scalar_fits(dtype, left_default) left = _as_scalar(left_default, dtype, is_cscalar=False) # pragma: is_grbscalar right_dtype = temp_op.type2 dtype = right_dtype if right_dtype._is_udt else None if type(right_default) is not Scalar: + if dtype is not None: + _check_literal_fits(dtype, right_default) try: right = Scalar.from_value( right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -867,6 +893,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): op=op, ) else: + _check_scalar_fits(dtype, right_default) right = _as_scalar(right_default, dtype, is_cscalar=False) # pragma: is_grbscalar op1 = get_typed_op(op, self.dtype, right.dtype, kind="binary") @@ -1140,6 +1167,241 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=False, name=None, **opts): # End auto-generated code: Scalar +def _literal_dtype(dtype, value, op): + """Return the dtype to build ``value`` with, as the other operand of ``op`` beside ``dtype``. + + ``None`` means the value's own dtype. This is where python-graphblas + decides how a literal operand is typed beside a UDT; built-in dtypes + always get ``None`` and are typed by ``unify``. + + Beside a UDT, a tuple, list or array literal becomes an element of that + UDT, as it must for a user's op, but only when no value changes + (:func:`_check_literal_fits`). With a built-in op that lifts to UDTs, a + number is an operand in its own right, typed as numpy 2 types it (NEP 50): + + - a numpy scalar or 0-d array keeps its dtype (strong); + - a Python ``bool``, ``int``, ``float`` or ``complex`` is weak: in each + element or leaf it takes that element's dtype when it fits the kind, so + ``int8_udt + 1`` stays ``int8_udt`` and ``int8_udt += 1`` can be stored, + while ``int8_udt * 2.5`` has float64 elements (see + ``udt_utils._weak_literal_udt``). Converting every literal to the UDT, + as before, made ``int8_udt * 2.5`` multiply by 2. + + ``eq`` and ``ne`` take a number as it is, strong, since a comparison + needs no result type and ``int8_udt == 300`` should be False, not an error. + """ + if not dtype._is_udt or getattr(op, "is_positional", False): + return None + is_python_number = type(value) in {bool, int, float, complex} + is_numpy_number = isinstance(value, (np.bool_, np.number)) or ( + isinstance(value, np.ndarray) and value.ndim == 0 and value.dtype.kind in "biufc" + ) + if is_python_number or is_numpy_number: + from .operator.binary import BinaryOp + from .operator.udt_utils import BUILTIN_UDT_BINARY_OPS, _weak_literal_udt + + if type(op) is BinaryOp and not op._anonymous: + if op.name in ("eq", "ne") or (is_numpy_number and op.name in BUILTIN_UDT_BINARY_OPS): + return None + if op.name in BUILTIN_UDT_BINARY_OPS: + return _weak_literal_udt(dtype, value) + if getattr(op, "_udt_types", True) is not None: + # An op that does not take UDTs at all (``lt``) reports that itself. + _check_literal_fits(dtype, value) + return dtype + + +def _check_literal_fits(dtype, value): + """Raise if converting literal ``value`` into UDT ``dtype`` would change one of its values. + + A literal becomes an element of the UDT when nothing else types it: a tuple, + list or array beside a UDT, a literal beside a user-defined op, the + defaults of ``ewise_union``. numpy turns ``0.5`` into ``0`` in an integer + field without a word, and the op would then run on a value the caller never + wrote. A number may round to a narrower float, as numpy 2 rounds a Python + number beside float32, but may not lose an imaginary part or overflow to + infinity. + """ + import numbers + + if not isinstance(value, (numbers.Number, tuple, list, dict, np.generic, np.ndarray)): + return + np_type = dtype.np_type + try: + if np_type.names is not None: + if isinstance(value, dict): + value = _dict_to_record(np_type, value) + given = np.zeros(1, dtype=_object_layout(np_type)) + converted = np.zeros(1, dtype=np_type) + else: + base, shape = np_type.subdtype + given = np.empty(shape, dtype=object) + converted = np.zeros(shape, dtype=base) + given[:] = value + import warnings + + with warnings.catch_warnings(): + # numpy warns for some lossy conversions (an out-of-range int on + # numpy 1, a complex array into floats); the comparison below + # reports each of them the same way on every numpy. + warnings.simplefilter("ignore") + converted[:] = _view_if_same_layout(value, converted.dtype) + except (TypeError, ValueError): + return # Not something this can read; Scalar.from_value reports it. + for given_leaf, converted_leaf in zip( + _leaf_arrays(given), _leaf_arrays(converted), strict=True + ): + kind = converted_leaf.dtype.kind + if kind not in "biufc": + continue + for x, y in zip(given_leaf.ravel().tolist(), converted_leaf.ravel().tolist(), strict=True): + if kind in "fc": + z = complex(x) + changed = (kind == "f" and z.imag != 0) or (np.isfinite(z) and not np.isfinite(y)) + else: + changed = x != y + if changed and kind in "iu" and type(x) is int: + # numpy 2 raises this for a Python int itself; numpy 1 wraps. + raise OverflowError( + f"Python integer {x} out of bounds for {converted_leaf.dtype}" + ) + if changed: + raise ValueError( + f"converting {value!r} to {dtype} would change {x!r} to {y!r}; a value is " + "converted into a UDT only when none of it changes." + ) + + +def _check_scalar_fits(dtype, scalar): + """Raise if a typed Scalar would change a value when converted into UDT ``dtype``. + + An ``ewise_union`` default must have the op's input type, so it is + converted; the same rule as for a literal (:func:`_check_literal_fits`). + """ + if dtype is not None and scalar.dtype != dtype and not scalar._is_empty: + _check_literal_fits(dtype, scalar.value) + + +def _is_lifted_op(op): + """Whether ``op`` is a built-in BinaryOp that lifts to UDTs, or a Monoid of one.""" + from .operator.binary import BinaryOp + from .operator.monoid import Monoid + from .operator.udt_utils import BUILTIN_UDT_BINARY_OPS + + if isinstance(op, Monoid): + op = op.binaryop + return type(op) is BinaryOp and not op._anonymous and op.name in BUILTIN_UDT_BINARY_OPS + + +def _object_layout(np_type): + """Return ``np_type`` with every leaf an object, keeping field names, nesting and shapes.""" + if np_type.names is not None: + return np.dtype([(name, _object_layout(np_type.fields[name][0])) for name in np_type.names]) + if np_type.subdtype is not None: + base, shape = np_type.subdtype + return np.dtype((_object_layout(base), shape)) + return np.dtype(object) + + +def _leaf_arrays(arr): + """Yield the array of each leaf of a structured array, in field order.""" + if arr.dtype.names is None: + yield arr + else: + for name in arr.dtype.names: + yield from _leaf_arrays(arr[name]) + + +def _ewise_add_converts_operands(op, left_dtype, right_dtype): + """Whether ``ewise_add`` must convert its operands to ``op``'s result type before the op. + + eWiseAdd copies an entry present in only one input into the result, cast + to the op's output type, and GraphBLAS checks that cast before it looks at + the data. A UDT casts only to itself, so such a pair is a bare + GrB_DOMAIN_MISMATCH even when every entry has a partner. Each lifted + arithmetic op gives the same values on operands already converted to its + result type, since that type is what it promotes their elements to, so + those convert first (``_ewise_add_in_result_type``). Any other op raises + ``DomainMismatch`` here. + """ + out = op.return_type + if not (out._is_udt or left_dtype._is_udt or right_dtype._is_udt): + return False + if left_dtype == out and right_dtype == out: + return False + from .operator.udt_utils import BUILTIN_UDT_BINARY_OPS + + parent = op.parent + if op.opclass == "BinaryOp" and not parent._anonymous and parent.name in BUILTIN_UDT_BINARY_OPS: + return True + from ..exceptions import DomainMismatch + + raise DomainMismatch( + f"ewise_add cannot use {parent!r} on {left_dtype} and {right_dtype}: an entry present " + f"in only one input is cast to the result type {out}, and UDTs do not cast. Use " + "ewise_union, which takes a default for each side, or ewise_mult." + ) + + +def _ewise_add_in_result_type(updater, left, right, op): + """Recipe for ``ewise_add`` when ``_ewise_add_converts_operands`` says to convert. + + Each operand is converted to the result type (:func:`_converted_to`), and + the op then runs on the two operands in that one type. + """ + out = op.return_type + if left.dtype != out: + left = _converted_to(left, out, name="left_temp", **updater.opts) + if right.dtype != out: + right = _converted_to(right, out, name="right_temp", **updater.opts) + updater << left.ewise_add(right, op.parent) + + +def _converted_to(operand, dtype, **new_kwargs): + """Return ``operand`` converted to UDT ``dtype``, which holds its elements, value for value. + + It adds ``dtype``'s zero on the left with the lifted ``plus``. With the + zero on the left, the result is ``dtype`` itself even when the operand has + the same elements in another layout (a packed and an aligned record), + since a left operand that holds the promoted elements is the result type. + The float and complex leaves of that zero are ``-0.0``, because + ``-0.0 + x`` is ``x`` for every float, where ``0.0 + -0.0`` is ``0.0``. + """ + np_type = dtype.np_type + if np_type.names is not None: + zero = np.zeros(1, dtype=np_type) + else: + zero = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) + for leaf in _leaf_arrays(zero): + if leaf.dtype.kind in "fc": + leaf[...] = complex(-0.0, -0.0) if leaf.dtype.kind == "c" else -0.0 + zero = Scalar.from_value(zero[0] if np_type.names is not None else zero, dtype, name="") + rv = operand.apply(binary.plus, left=zero).new(**new_kwargs) + if rv.dtype != dtype: # pragma: no cover (safety) + from ..exceptions import DomainMismatch + + raise DomainMismatch(f"cannot convert {operand.dtype} to {dtype}") + return rv + + +def _cast_for_store(value, dtype): + """Return a Scalar, Vector or Matrix ``value`` as ``dtype``, to be assigned into an object. + + GraphBLAS's assign and setElement take a UDT value as raw bytes, so one of + another type would be read as the wrong type; ``base._store_cast_op`` + casts it first or raises. + """ + from .base import _store_cast_op + + if (cast_op := _store_cast_op(value.dtype, dtype)) is None: + return value + if not isinstance(value, BaseType): + # An expression, such as ``x[[0, 1]] << y + 0.5``: compute it in its own + # type, then cast it. + value = value.new(name="") + return value.apply(cast_op).new(name=value.name) + + def _as_scalar(scalar, dtype=None, *, is_cscalar): if type(scalar) is not Scalar: return Scalar.from_value(scalar, dtype, is_cscalar=is_cscalar, name="") diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index 5d8c8fef9..ba48dce80 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -25,6 +25,12 @@ ScalarExpression, ScalarIndexExpr, _as_scalar, + _cast_for_store, + _check_literal_fits, + _check_scalar_fits, + _ewise_add_converts_operands, + _ewise_add_in_result_type, + _literal_dtype, _scalar_index, ) from .utils import ( @@ -1026,15 +1032,30 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): # Per the spec, op may be a semiring, but this is weird, so don't. self._expect_op(op, ("BinaryOp", "Monoid"), within=method_name, argname="op") - if other.ndim == 2: - # Broadcast columnwise from the left - if other._nrows != self._size: - # Check this before we compute a possibly large matrix below - raise DimensionMismatch( - "Dimensions not compatible for broadcasting Vector from the left " - f"to columns of Matrix in {method_name}. Matrix.nrows (={other._nrows}) " - f"must equal Vector.size (={self._size})." + if other.ndim == 2 and other._nrows != self._size: + # Broadcast columnwise from the left. Check this before we compute a + # possibly large matrix below. + raise DimensionMismatch( + "Dimensions not compatible for broadcasting Vector from the left " + f"to columns of Matrix in {method_name}. Matrix.nrows (={other._nrows}) " + f"must equal Vector.size (={self._size})." + ) + if _ewise_add_converts_operands(op, self.dtype, other.dtype): + # GraphBLAS casts an unpaired entry to the result type, which a UDT + # cannot be. Convert both operands to the result type first and + # apply the op in that one type: an int8 UDT plus a float32 one is + # float32, with every entry of either side. + args = [self, other, _ewise_add_in_result_type, (self, other, op)] + if other.ndim == 2: + return MatrixExpression( + method_name, None, args, nrows=other._nrows, ncols=other._ncols, op=op ) + expr = VectorExpression(method_name, None, args, op=op) + if self._size != other._size: + expr.new(name="") # incompatible shape; raise now + return expr + + if other.ndim == 2: return MatrixExpression( method_name, None, @@ -1207,6 +1228,8 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): left_dtype = temp_op.type dtype = left_dtype if left_dtype._is_udt else None if type(left_default) is not Scalar: + if dtype is not None: + _check_literal_fits(dtype, left_default) try: left = Scalar.from_value( left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -1221,10 +1244,13 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): op=op, ) else: + _check_scalar_fits(dtype, left_default) left = _as_scalar(left_default, dtype, is_cscalar=False) # pragma: is_grbscalar right_dtype = temp_op.type2 dtype = right_dtype if right_dtype._is_udt else None if type(right_default) is not Scalar: + if dtype is not None: + _check_literal_fits(dtype, right_default) try: right = Scalar.from_value( right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -1239,6 +1265,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): op=op, ) else: + _check_scalar_fits(dtype, right_default) right = _as_scalar(right_default, dtype, is_cscalar=False) # pragma: is_grbscalar if is_infix: @@ -1420,6 +1447,11 @@ def apply(self, op, right=None, *, left=None): if isinstance(op, str): op = op_from_string(op) op, opclass = find_opclass(op) + if opclass == "Monoid" and self.dtype._is_udt: + # Applying a Monoid applies its BinaryOp. Typing it as that op lets a + # literal promote as it would with the BinaryOp, where a Monoid takes + # one type and would need the literal converted into the UDT. + op, opclass = op.binaryop, "BinaryOp" if opclass in {"IndexUnaryOp", "SelectOp"}: # Provide default value for index unary if right is None: @@ -1441,7 +1473,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = None elif right is None: if type(left) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_dtype(self.dtype, left, op) try: left = Scalar.from_value(left, dtype, is_cscalar=None, name="") except TypeError: @@ -1477,7 +1509,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = "{1.name}.apply({op}, left={0._expr_name})" elif left is None: if type(right) is not Scalar: - dtype = self.dtype if (self.dtype._is_udt and not op.is_positional) else None + dtype = _literal_dtype(self.dtype, right, op) try: right = Scalar.from_value(right, dtype, is_cscalar=None, name="") except TypeError: @@ -1873,6 +1905,8 @@ def _prep_for_extract(self, resolved_indexes, mask=None, is_submask=False): def _assign_element(self, resolved_indexes, value): idx = resolved_indexes.indices[0] + if type(value) is Scalar: + value = _cast_for_store(value, self.dtype) if type(value) is not Scalar: dtype = self.dtype if self.dtype._is_udt else None try: @@ -1901,6 +1935,8 @@ def _assign_element(self, resolved_indexes, value): def _prep_for_assign(self, resolved_indexes, value, mask, is_submask, replace, opts): method_name = "__setitem__" + if output_type(value) in {Scalar, Vector}: + value = _cast_for_store(value, self.dtype) idx = resolved_indexes.indices[0] size = idx.size cscalar = idx.cscalar diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 981ad2036..b6a5f5424 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -42,7 +42,7 @@ UINT32, UINT64, ) -from graphblas.exceptions import DomainMismatch, UdfParseError +from graphblas.exceptions import DimensionMismatch, DomainMismatch, UdfParseError from .conftest import shouldhave @@ -1263,9 +1263,12 @@ def _udt_index(val, idx, _, thunk): # pragma: no cover (numba) _udt_index = IndexUnaryOp.register_anonymous(_udt_index, "_udt_index", is_udt=True) assert udt in _udt_index - result = v.apply(_udt_index, 3).new() + result = v.apply(_udt_index, (False, 3)).new() expected = Vector.from_coo([0, 1, 2], [3, -3, -3]) assert result.isequal(expected) + # A number fills every field, and 3 in the bool field would be True. + with pytest.raises(ValueError, match="would change 3 to True"): + v.apply(_udt_index, 3) def _udt_first(x, y): return x # pragma: no cover (numba) @@ -1636,7 +1639,7 @@ def _record_expected(udt, rows): ("plus", [(11, 22.0), (33, 44.0), (55, 66.0)]), ("minus", [(-9, -18.0), (-27, -36.0), (-45, -54.0)]), ("times", [(10, 40.0), (90, 160.0), (250, 360.0)]), - ("truediv", [(0, 0.1), (0, 0.1), (0, 0.1)]), + ("truediv", [(0.1, 0.1), (0.1, 0.1), (0.1, 0.1)]), ], ) @pytest.mark.skipif("not supports_udfs") @@ -1645,7 +1648,16 @@ def test_udt_builtin_binary_record(record_udt, op_name, expected_rows): """Per-field arithmetic on a record UDT matches the scalar definition.""" v, w = _record_pair(record_udt) result = getattr(binary, op_name)(v & w).new() - assert result.isequal(_record_expected(record_udt, expected_rows)) + if op_name == "truediv": + # The int64 field divides to float64, as INT64 does, so the result is + # the record with both fields float64. + expected_type = dtypes.lookup_dtype( + np.dtype([("a", np.float64), ("b", np.float64)], align=True) + ) + else: + expected_type = record_udt + assert result.dtype == expected_type + assert result.isequal(_record_expected(expected_type, expected_rows)) @pytest.mark.skipif("not supports_udfs") @@ -1831,8 +1843,8 @@ def test_udt_jit_op_definitions(): lib.GrB_BinaryOp_get_String(typed.gb_obj, buf, lib.GxB_JIT_C_DEFINITION) defn = ffi.string(buf).decode() assert "plus_Vec5Jit" in defn - assert "v[0]" in defn - assert "v[4]" in defn + assert "z->v[i] = (x->v[i]) + (y->v[i])" in defn + assert "i < 5" in defn typed = unary.ainv[arr_udt] lib.GrB_UnaryOp_get_String(typed.gb_obj, buf, lib.GxB_JIT_C_DEFINITION) @@ -3206,79 +3218,1164 @@ def test_udt_record_nesting_mismatch_is_a_keyerror(case, op): v.ewise_mult(w, getattr(binary, op)).new() -def _record_with_field_shape(shape, name): +def _record_with_field_shape(shape, name, b_dtype=np.float64): """Register a record UDT whose ``fsh_a`` field has ``shape`` (a scalar for ``()``).""" field = ("fsh_a", np.float64, shape) if shape else ("fsh_a", np.float64) - return dtypes.register_anonymous(np.dtype([field, ("fsh_b", np.float64)]), name) + return dtypes.register_anonymous(np.dtype([field, ("fsh_b", b_dtype)]), name) @pytest.mark.skipif("not supports_udfs") @pytest.mark.parametrize("op", ["plus", "eq"]) @pytest.mark.parametrize( - ("left", "right", "match"), - [ - ((3,), (4,), "shapes broadcast together"), - ((2, 3), (3, 2), "shapes broadcast together"), - ((), (3,), "only when the array is on the left"), - ], + ("left", "right"), + [((3,), (4,)), ((2, 3), (3, 2)), ((3,), ()), ((3,), (1,)), ((3,), (1, 3))], ) -def test_udt_record_field_shape_mismatch_is_a_keyerror(left, right, match, op): - """Record array fields that cannot be combined are a KeyError. - - The codegen combines array fields as numpy does, by broadcasting, so a - ``(3,)`` field with a ``(4,)`` one raised inside the cfunc, where the - caller never saw it, and the result element kept whatever was in the - buffer (``eq`` read it as its bool). A scalar field on the left of an - array one failed Numba's typing instead. +def test_udt_record_field_shape_mismatch_is_a_keyerror(left, right, op): + """Record fields whose shapes differ are a KeyError, whichever operand is on the left. + + The codegen combines array fields as numpy arrays, so a ``(3,)`` field with + a ``(4,)`` one raised inside the cfunc, where the caller never saw it, and + the result element kept whatever was in the buffer (``eq`` read it as its + bool). Shapes that broadcast worked only one way round: ``plus`` writes + into a field shaped like the left operand's, so a ``(3,)`` field took a + ``(1,)`` one on its right but failed with it on its left, and a scalar + field on the left of an array one never compiled. numpy will not promote + or compare two record dtypes whose fields differ in shape either. """ - v = Vector(_record_with_field_shape(left, "_FshLeft"), size=1) - w = Vector(_record_with_field_shape(right, "_FshRight"), size=1) - with pytest.raises(KeyError, match=match): - v.ewise_mult(w, getattr(binary, op)).new() + X = _record_with_field_shape(left, "_FshLeft") + Y = _record_with_field_shape(right, "_FshRight") + for a, b in [(X, Y), (Y, X)]: + v = Vector(a, size=1) + w = Vector(b, size=1) + with pytest.raises(KeyError, match="must have the same shape for each field"): + v.ewise_mult(w, getattr(binary, op)).new() + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_record_fields_of_equal_shape_combine(): + """Two record UDTs with the same field shapes combine, though a field's dtype differs. + + ``eq`` and ``ne`` used to fail on this pair with GrB_DOMAIN_MISMATCH: + ``get_typed_op`` unified the two records to one of them, since numpy + promotes their field dtypes, so GraphBLAS saw an operand of the other + type. They now compile the pair, as the arithmetic ops do. + """ + X = _record_with_field_shape((2, 3), "_FshLeft") + v = Vector(X, size=2) + v[0] = (np.arange(6.0).reshape(2, 3), 1.0) + v[1] = (np.ones((2, 3)), 2.0) + w = Vector(_record_with_field_shape((2, 3), "_FshRight", b_dtype=np.float32), size=2) + w[0] = (np.full((2, 3), 10.0), 0.5) + w[1] = (np.ones((2, 3)), 2.0) + total = v.ewise_mult(w, binary.plus).new() + assert total.dtype == X + value = total[0].new().value + np.testing.assert_array_equal(value["fsh_a"], np.arange(6.0).reshape(2, 3) + 10.0) + assert value["fsh_b"] == 1.5 + assert binary.eq(v & w).new().isequal(Vector.from_coo([0, 1], [False, True])) + assert binary.ne(v & w).new().isequal(Vector.from_coo([0, 1], [True, False])) + # A semiring dispatches through its BinaryOp, so ``lor_eq`` takes the pair too. + assert v.inner(w, semiring.lor_eq).new().value + assert not v.inner(w, semiring.land_eq).new().value + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_positional_ops_take_two_different_udts(): + """``first``, ``second``, ``any`` and ``pair`` take any two UDTs; they never read a field. + + ``any`` was the odd one out: it had no ``_custom_dtype``, so ``get_typed_op`` + unified two records of the same layout to one of them and GraphBLAS raised + GrB_DOMAIN_MISMATCH on the other operand. ``get_typed_op`` now never unifies + two different UDTs. + """ + # Own field names: the helper's scalar-field record is ``_FshRight``'s dtype, + # and ``register_anonymous`` returns one cached DataType per dtype. + X = dtypes.register_anonymous( + np.dtype([("pos_a", np.float64), ("pos_b", np.float64)]), "_PosLeft" + ) + Y = dtypes.register_anonymous( + np.dtype([("pos_a", np.float64), ("pos_b", np.int64)]), "_PosRight" + ) + v = Vector(X, size=1) + v[0] = (1.0, 3.0) + w = Vector(Y, size=1) + w[0] = (4.0, 6) + assert binary.first(v & w).new().isequal(v) + assert binary.second(v & w).new().isequal(w) + assert binary.any(v & w).new().isequal(w) + assert binary.pair(v & w).new().isequal(Vector.from_coo([0], [1])) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_record_pair_promotes_each_field(): + """Two records whose field dtypes differ give each field the dtype the built-in op gives. + + The result was always the left record, so ``x + y`` truncated an int64 field + that ``y + x`` kept as float64, and a complex field on the right of ``min`` + slipped past the check that rejects complex ``min`` and failed in Numba. + When one operand's record holds every promoted field, the result is that + record, on either side; otherwise it is the record with the promoted + fields, which is the UDT registered with that layout if there is one. Each + record also has an array field, which no Numba return type can be matched + against, so the result type cannot be recovered from what the generated + function returns. + """ + + def record(a_dtype, b_dtype, name): + fields = [("wid_a", a_dtype), ("wid_b", b_dtype), ("wid_v", np.float64, (2,))] + return dtypes.register_anonymous(np.dtype(fields), name) + + f8 = np.float64 + narrow = record(np.int64, f8, "_WidN") + wide = record(f8, f8, "_WidW") + crossed = record(f8, np.int64, "_WidX") + cplx = record(np.complex128, f8, "_WidC") + v = Vector(narrow, size=1) + v[0] = (1, 2.0, [1.0, 2.0]) + w = Vector(wide, size=1) + w[0] = (0.5, 0.25, [10.0, 20.0]) + for a, b in [(v, w), (w, v)]: + result = binary.plus(a & b).new() + assert result.dtype == wide + value = result[0].new().value + assert (value["wid_a"], value["wid_b"]) == (1.5, 2.25) + np.testing.assert_array_equal(value["wid_v"], [11.0, 22.0]) + # Each of ``narrow`` and ``crossed`` is wider in one field, so neither holds + # both. The promoted record has ``wide``'s layout, so it is ``wide``. + x = Vector(crossed, size=1) + x[0] = (1.0, 2, [1.0, 2.0]) + for a, b in [(v, x), (x, v)]: + result = binary.times(a & b).new() + assert result.dtype == wide + value = result[0].new().value + assert (value["wid_a"], value["wid_b"]) == (1.0, 4.0) + np.testing.assert_array_equal(value["wid_v"], [1.0, 4.0]) + assert binary.eq(v & x).new().isequal(Vector.from_coo([0], [True])) + c = Vector(cplx, size=1) + c[0] = (1 + 1j, 2.0, [0.0, 0.0]) + for a, b in [(w, c), (c, w)]: + assert binary.plus(a & b).new().dtype == cplx + with pytest.raises(KeyError, match="not defined on complex fields"): + binary.min(a & b).new() + # The result is ``wide`` on either side. Stored in ``narrow``, its float + # field is cast to int64 as a built-in store would cast it. + v << binary.plus(v & w) + assert v.dtype == narrow + value = v[0].new().value + assert (value["wid_a"], value["wid_b"]) == (1, 2.25) + + +@pytest.mark.skipif("not supports_udfs") +# SS < 9 names a UDT by its numpy repr and warns when that is over 128 chars, +# as these nested records' reprs are. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_udt_record_pair_promotes_to_a_third_record(): + """Records that are each wider in some field give a third record, as numpy would. + + Each field is promoted as the built-in op promotes two Vectors, so int16 + with float32 is float32, and the result keeps the operands' names, nesting + and alignment. This pair used to be a KeyError, since the result had to be + one of the two operands. + """ + i2, i8, f4, f8 = np.int16, np.int64, np.float32, np.float64 + X = dtypes.register_anonymous( + np.dtype([("pro_a", [("pro_i", i2), ("pro_j", f8)]), ("pro_b", f4)], align=True), "_ProX" + ) + Y = dtypes.register_anonymous( + np.dtype([("pro_a", [("pro_i", f4), ("pro_j", i8)]), ("pro_b", i2)], align=True), "_ProY" + ) + expected = np.dtype([("pro_a", [("pro_i", f4), ("pro_j", f8)]), ("pro_b", f4)], align=True) + v = Vector(X, size=1) + v[0] = ((3, 0.5), 1.5) + w = Vector(Y, size=1) + w[0] = ((0.25, 4), 2) + for a, b in [(v, w), (w, v)]: + result = binary.plus(a & b).new() + assert result.dtype.np_type == expected + assert result[0].new().value.tolist() == ((3.25, 4.5), 3.5) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_record_scalar_promotes_each_field(): + """A number beside a record UDT keeps its dtype, and each field promotes as a Vector's would. + + The number used to be converted to the record first, so ``int_record + 0.5`` + added 0 to each field. + """ + ints = dtypes.register_anonymous( + np.dtype([("rsp_x", np.int64), ("rsp_y", np.int64)]), "_RspInts" + ) + floats = dtypes.register_anonymous( + np.dtype([("rsp_f", np.float64), ("rsp_g", np.float64)]), "_RspFloats" + ) + promoted = np.dtype([("rsp_x", np.float64), ("rsp_y", np.float64)]) + v = Vector(ints, size=1) + v[0] = (2, 3) + for expr, expected in [ + (v + 0.5, (2.5, 3.5)), + (1.5 * v, (3.0, 4.5)), + (v.apply(binary.truediv, right=4), (0.5, 0.75)), + ]: + result = expr.new() + assert result.dtype.np_type == promoted + assert result[0].new().value.tolist() == expected + assert (v + 1).new().dtype == ints + w = Vector(floats, size=1) + w[0] = (0.5, 1.5) + result = (w + 1).new() + assert result.dtype == floats + assert result[0].new().value.tolist() == (1.5, 2.5) @pytest.mark.skipif("not supports_udfs") -@pytest.mark.parametrize(("left", "right"), [((1,), (3,)), ((3,), (3, 1)), ((3, 1), (3,))]) -def test_udt_record_field_broadcast_exceeding_left_is_a_keyerror(left, right): - """An arithmetic op writes the broadcast into a field shaped like the left one. +def test_udt_packed_record_with_aligned_takes_an_operand_type(): + """A packed and an aligned record with the same leaves give the left operand's type. - ``_check_udf_fills_output`` caught these pairs when the wrapper was built, - but as "UDT UDF returned an array of shape (3,)" on sample values, naming a - UDF the user did not write. ``eq`` reduces the field to one bool, so it - takes the same pair. + Their fields are the same dtypes at different offsets, so neither holds the + other's layout, but each holds every result field. Taking the left one lets + ``x << x + y`` update ``x`` in place, whichever of the two ``x`` is. """ - v = Vector(_record_with_field_shape(left, "_FshLeft"), size=1) - w = Vector(_record_with_field_shape(right, "_FshRight"), size=1) - with pytest.raises(KeyError, match="broadcast to the left field's shape"): - v.ewise_mult(w, binary.plus).new() - v[0] = (np.ones(left), 1.0) - w[0] = (np.ones(right), 1.0) - assert binary.eq(v & w).new().isequal(Vector.from_coo([0], [True])) + fields = [("pak_a", np.int8), ("pak_b", np.float64)] + aligned = dtypes.register_anonymous(np.dtype(fields, align=True), "_PakAligned") + packed = dtypes.register_anonymous(np.dtype(fields), "_PakPacked") + assert aligned.np_type.itemsize != packed.np_type.itemsize + x = Vector(aligned, size=1) + x[0] = (1, 0.5) + y = Vector(packed, size=1) + y[0] = (2, 0.25) + for a, b in [(x, y), (y, x)]: + assert binary.plus(a & b).new().dtype == a.dtype + a << binary.plus(a & b) + assert x[0].new().value.tolist() == (3, 0.75) + assert y[0].new().value.tolist() == (5, 1.0) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.parametrize("op", ["plus", "eq"]) +@pytest.mark.parametrize(("left", "right"), [((2, 3), (3, 2)), ((6,), (2, 3)), ((3,), (4,))]) +def test_udt_array_pair_shape_mismatch_is_a_keyerror(left, right, op): + """Array UDTs whose shapes do not broadcast are a KeyError, as numpy refuses them. + + Two of these with the same size used to combine over their flat buffers, + so ``(2, 3)`` plus ``(3, 2)`` added elements numpy never pairs. The base + dtypes differ, which is allowed, so only the shapes are at fault. + """ + # int16 and int32: other tests register float64 arrays of these shapes, and + # ``register_anonymous`` returns one cached DataType per dtype. + X = dtypes.register_anonymous(np.dtype((np.int16, left)), "_PairArrL") + Y = dtypes.register_anonymous(np.dtype((np.int32, right)), "_PairArrR") + for a, b in [(X, Y), (Y, X)]: + with pytest.raises(KeyError, match="shapes that broadcast together"): + getattr(binary, op)(Vector(a, size=1) & Vector(b, size=1)).new() @pytest.mark.skipif("not supports_udfs") @pytest.mark.parametrize( - ("left", "right"), [((3,), ()), ((3,), (1,)), ((3,), (1, 3)), ((2, 3), (3,))] + ("x_type", "y_type"), + [ + ((np.uint16, (3, 1)), (np.int32, (1, 4))), + ((np.uint16, (2, 1, 3)), (np.float32, (4, 1))), + ((np.uint16, (1,)), (np.int16, (2, 2))), + ((np.int16, (3, 1)), (np.int16, (2, 1, 4))), + ], ) -def test_udt_record_field_shapes_that_broadcast(left, right): - """Record array fields that broadcast together still combine as numpy arrays do.""" - X = _record_with_field_shape(left, "_FshLeft") - xa = np.arange(1.0, 1.0 + np.prod(left)).reshape(left) - v = Vector(X, size=2) - v[0] = (xa, 2.0) - v[1] = (np.full(left, 10.0), 2.0) - w = Vector(_record_with_field_shape(right, "_FshRight"), size=2) - w[0] = (np.full(right, 10.0), 2.0) - w[1] = (np.full(right, 10.0), 2.0) - total = v.ewise_mult(w, binary.plus).new() - assert total.dtype == X - expected = np.empty(left) - expected[...] = xa + np.full(right, 10.0) - first = total[0].new().value - np.testing.assert_array_equal(first["fsh_a"], expected) - assert first["fsh_b"] == 4.0 - # Element 0 differs only in ``fsh_a``, so its False comes from the arrays. - assert binary.eq(v & w).new().isequal(Vector.from_coo([0, 1], [False, True])) +def test_udt_array_pairs_broadcast_as_numpy(x_type, y_type): + """Two array UDTs broadcast as numpy arrays do, in either order. + + The result has numpy's broadcast shape and the promoted element dtype. It + is an operand's type when one has both, else the structural array type. + ``eq`` is True when every pair of elements numpy compares is equal. + """ + X = dtypes.register_anonymous(np.dtype(x_type), "_BcastX") + Y = dtypes.register_anonymous(np.dtype(y_type), "_BcastY") + xv = np.arange(1, 1 + np.prod(x_type[1])).reshape(x_type[1]).astype(x_type[0]) + yv = 10 * np.arange(1, 1 + np.prod(y_type[1])).reshape(y_type[1]).astype(y_type[0]) + v = Vector(X, size=3) + v[0] = xv + v[1] = xv + w = Vector(Y, size=3) + w[0] = yv + w[2] = yv + M = Matrix(X, nrows=1, ncols=1) + M[0, 0] = xv + N = Matrix(Y, nrows=1, ncols=1) + N[0, 0] = yv + for a, b, A, B, av, bv in [(v, w, M, N, xv, yv), (w, v, N, M, yv, xv)]: + expected = np.subtract(av, bv) + result_type = dtypes.lookup_dtype(np.dtype((expected.dtype, expected.shape))) + for result in [binary.minus(a & b).new(), binary.minus(A & B).new()]: + assert result.dtype == result_type + assert result.nvals == 1 + value = result[0].new().value if result.ndim == 1 else result[0, 0].new().value + np.testing.assert_array_equal(value, expected) + assert not binary.eq(a & b).new()[0].new().value + assert binary.ne(a & b).new()[0].new().value + # ewise_add converts an entry without a partner to the result type, so + # it broadcasts to the result's shape. + added = a.ewise_add(b, binary.minus).new() + assert added.dtype == result_type + np.testing.assert_array_equal(added[0].new().value, expected) + lone = 1 if a is v else 2 + np.testing.assert_array_equal(added[lone].new().value, np.broadcast_to(av, expected.shape)) + # Equal values: every compared pair is equal, as ``(xv == yv).all()``. + w[0] = np.broadcast_to(xv[(0,) * xv.ndim], y_type[1]) + v[0] = np.broadcast_to(xv[(0,) * xv.ndim], x_type[1]) + assert binary.eq(v & w).new()[0].new().value + assert binary.eq(w & v).new()[0].new().value + assert not binary.ne(v & w).new()[0].new().value + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_broadcast_semiring_and_reduce(): + """A semiring on two array UDTs that broadcast sums the broadcast products.""" + X = dtypes.register_anonymous(np.dtype((np.int16, (3, 1))), "_BcastSemiX") + Y = dtypes.register_anonymous(np.dtype((np.int16, (1, 5))), "_BcastSemiY") + a = [np.arange(3).reshape(3, 1) + k for k in range(2)] + b = [10 * np.arange(5).reshape(1, 5) - k for k in range(2)] + A = Matrix(X, nrows=1, ncols=2) + A[0, 0] = a[0] + A[0, 1] = a[1] + B = Matrix(Y, nrows=2, ncols=1) + B[0, 0] = b[0] + B[1, 0] = b[1] + C = A.mxm(B, semiring.plus_times).new() + assert C.dtype.np_type == np.dtype((np.int16, (3, 5))) + expected = a[0] * b[0] + a[1] * b[1] + np.testing.assert_array_equal(C[0, 0].new().value, expected) + s = A.reduce_scalar(monoid.plus).new() + np.testing.assert_array_equal(s.value, a[0] + a[1]) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_record_2d_field_compares_with_a_number(): + """``eq`` and ``ne`` of a record with a 2-D array field and a number compare every element. + + The generated code read the field at flat positions, which in a 2-D field + are rows, so it failed to compile once a number was compared as it is + rather than converted into the record first. + """ + T = dtypes.register_anonymous( + np.dtype([("cf_grid", np.float64, (2, 3)), ("cf_n", np.int64)], align=True), "_CmpField2D" + ) + v = Vector(T, size=2) + v[0] = (np.full((2, 3), 2.0), 2) + v[1] = (np.array([[2.0, 2, 2], [2, 2, np.nan]]), 2) + assert binary.eq(v, 2.0).new().isequal(Vector.from_coo([0, 1], [True, False])) + assert binary.ne(2.0, v).new().isequal(Vector.from_coo([0, 1], [False, True])) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_ops_on_large_elements(): + """Lifted ops on a 64 by 64 array compile in a loop, not one statement per element. + + Unrolled, Numba took about two minutes to type ``minus`` on this element. + """ + T = dtypes.register_anonymous(np.dtype((np.float64, (64, 64))), "_LoopBig") + S = dtypes.register_anonymous(np.dtype((np.float64, (64, 1))), "_LoopBigCol") + x = np.arange(64 * 64, dtype=np.float64).reshape(64, 64) + col = np.arange(64, dtype=np.float64).reshape(64, 1) + v = Vector(T, size=1) + v[0] = x + u = Vector(S, size=1) + u[0] = col + np.testing.assert_array_equal((v - u).new()[0].new().value, x - col) + np.testing.assert_array_equal(unary.ainv(v).new()[0].new().value, -x) + assert binary.eq(v & v).new()[0].new().value + assert not binary.ne(v & v).new()[0].new().value + w = v.dup() + w[0] = np.where(x == 4095, np.nan, x) + assert not binary.eq(w & w).new()[0].new().value + assert binary.ne(w & w).new()[0].new().value + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_pair_leading_ones(): + """Array UDTs that differ only in leading axes of length 1 combine, as numpy broadcasts them. + + This is what a ``Matrix.to_csr`` round trip produces: values of shape + ``(1, 3)`` for a ``(3,)`` UDT. The result takes the operand with more axes, + numpy's broadcast shape, whichever side it is on. + """ + short = dtypes.register_anonymous(np.dtype((np.int16, (3,))), "_PairArrShort") + long = dtypes.register_anonymous(np.dtype((np.int16, (1, 3))), "_PairArrLong") + v = Vector(short, size=1) + v[0] = [1, 2, 3] + w = Vector(long, size=1) + w[0] = [[10, 20, 30]] + for a, b in [(v, w), (w, v)]: + result = binary.plus(a & b).new() + assert result.dtype == long + np.testing.assert_array_equal(result[0].new().value, [[11, 22, 33]]) + assert not binary.eq(a & b).new()[0].new().value + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_pair_promotes_the_base(): + """Array UDTs with different base dtypes combine, each element promoted as the built-in op does. + + The pair used to be a KeyError. The result is the operand whose type holds + the promoted elements, on either side, so ``int64`` with ``float64`` gives + the ``float64`` one. ``eq`` and ``ne`` must read each operand as its own + dtype: reading both as the left one's compares the bits of a float64 as an + int64. + """ + ints = dtypes.register_anonymous(np.dtype((np.int64, (2, 2, 2))), "_PromI222") + floats = dtypes.register_anonymous(np.dtype((np.float64, (2, 2, 2))), "_PromF222") + base = np.arange(8).reshape(2, 2, 2) + v = Vector(ints, size=2) + v[0] = base + v[1] = base + w = Vector(floats, size=2) + w[0] = base + 0.5 + w[1] = base + for a, b in [(v, w), (w, v)]: + result = binary.plus(a & b).new() + assert result.dtype == floats + np.testing.assert_array_equal(result[0].new().value, 2 * base + 0.5) + assert binary.eq(a & b).new().isequal(Vector.from_coo([0, 1], [False, True])) + assert binary.ne(a & b).new().isequal(Vector.from_coo([0, 1], [True, False])) + # A semiring's monoid is typed by its multiplier's result. + assert v.inner(w, semiring.plus_times).new().dtype == floats + # GraphBLAS would cast an entry of ``v`` alone to the result type, which + # UDTs do not do; ``plus`` converts both sides first instead. + assert v.ewise_add(w, binary.plus).new().isequal((v + w).new()) + assert v.ewise_union(w, binary.plus, 0, 0.0).new().dtype == floats + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_scalar_promotes_like_builtin(): + """A number beside an array UDT keeps its dtype, and each element promotes as the op does. + + ``apply`` converted a Python number to the UDT first (``v + 0.5`` is an + ``apply``), and the result took the UDT's base whatever the other operand + was, so ``int_udt + 0.5`` added 0 and ``int_udt / 4`` gave integers. Each + element now has the dtype the op gives on the built-in dtypes, and the + result is the operand's UDT only when that holds it. + """ + ints = dtypes.register_anonymous(np.dtype((np.int64, (2, 4))), "_UpcI24") + int8s = dtypes.register_anonymous(np.dtype((np.int8, (2, 4))), "_UpcB24") + fp32s = dtypes.register_anonymous(np.dtype((np.float32, (2, 4))), "_UpcF24") + fp64s = np.dtype((np.float64, (2, 4))) + base = np.arange(8).reshape(2, 4) + v = Vector(ints, size=1) + v[0] = base + M = Matrix(ints, nrows=1, ncols=1) + M[0, 0] = base + for expr, expected in [ + (v + 0.5, base + 0.5), + (0.5 + v, base + 0.5), + (v.apply(binary.times, right=2.5), base * 2.5), + (v.apply(binary.truediv, right=4), base / 4), + (binary.minus(10.5, v), 10.5 - base), + (v.apply(binary.plus, right=gb.Scalar.from_value(0.5)), base + 0.5), + (v.ewise_mult(Vector.from_coo([0], [0.5]), binary.times), base * 0.5), + (M * 0.5, base * 0.5), + (0.5 * M, base * 0.5), + # A 0-d numpy array is a number too, not a whole element. + (v + np.array(0.5, np.float32), base + 0.5), + (v / np.array(4, np.int8), base / 4), + ]: + result = expr.new() + assert result.dtype.np_type == fp64s + value = result[0].new().value if result.ndim == 1 else result[0, 0].new().value + np.testing.assert_array_equal(value, expected) + # Same element type: the result keeps the UDT, as does a literal that is a + # whole element rather than a number. + for expr in [v + 1, v + [[1] * 4] * 2]: + result = expr.new() + assert result.dtype == ints + np.testing.assert_array_equal(result[0].new().value, base + 1) + # A Python number is weak, as in numpy 2: it takes the elements' dtype when + # their kind holds it, so INT8 elements plus 1 stay INT8 and FP32 elements + # plus 0.5 stay FP32, and an in-place update can store the result. A typed + # number is strong, and an int out of the elements' range is an error. + w = Vector(int8s, size=1) + w[0] = base + assert (w + 1).new().dtype == int8s + w += 1 + np.testing.assert_array_equal(w[0].new().value, base + 1) + assert (w + np.int64(1)).new().dtype == ints # the UDT registered with that layout + with pytest.raises(OverflowError, match="300 out of bounds for int8"): + w + 300 + x = Vector(fp32s, size=1) + x[0] = base + result = (x + 0.5).new() + assert result.dtype == fp32s + np.testing.assert_array_equal(result[0].new().value, base + 0.5) + x *= 2 + np.testing.assert_array_equal(x[0].new().value, base * 2) + assert (x + np.float64(0.5)).new().dtype.np_type == fp64s + if dtypes._supports_complex: + for expr in [lambda: binary.min(x, 1j), lambda: binary.min(1j, x)]: + with pytest.raises(KeyError, match="not defined on complex fields"): + expr() + # eq and ne compare with the number itself; converted to the UDT, 0.5 was 0. + zeros = Vector(ints, size=1) + zeros[0] = np.zeros((2, 4)) + assert not (zeros == 0.5).new()[0].new().value + assert (zeros != 0.5).new()[0].new().value + assert (zeros == 0).new()[0].new().value + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_uint64_with_signed_ints_computes_like_builtin(): + """UINT64 elements with signed ones are computed in FP64, as the built-in ops compute them. + + The result type was right, but the generated code let Numba do the + arithmetic, and Numba computes uint64 with int64 in int64: ``2**63 + 1`` + came out as ``-9.22e18`` and ``2**63 * 2`` as 0. Each operand is now + converted to the type the built-in op computes in first. + """ + big = 2**63 + uints = dtypes.register_anonymous(np.dtype((np.uint64, (2, 2))), "_MixU22") + ints = dtypes.register_anonymous(np.dtype((np.int64, (2, 2))), "_MixI22") + u = Vector(uints, size=1) + u[0] = [[big, 7], [0, 1]] + i = Vector(ints, size=1) + i[0] = [[1, -1], [0, 0]] + for expr, expected in [ + (binary.plus(u & i), [[big + 1.0, 6.0], [0.0, 1.0]]), + (binary.plus(i & u), [[big + 1.0, 6.0], [0.0, 1.0]]), + (binary.floordiv(i & u), [[0.0, -1.0], [np.nan, 0.0]]), + # A numpy int64 is strong, so it is a signed operand like ``i``. + (u + np.int64(1), [[big + 1.0, 8.0], [1.0, 2.0]]), + (binary.plus(np.int64(1), u), [[big + 1.0, 8.0], [1.0, 2.0]]), + (u * np.int64(2), [[2.0 * big, 14.0], [0.0, 2.0]]), + ]: + result = expr.new() + assert result.dtype.np_type == np.dtype((np.float64, (2, 2))) + np.testing.assert_array_equal(result[0].new().value, expected) + # A Python int is weak, so it takes the uint64 elements' type, exactly. + result = (u + 1).new() + assert result.dtype == uints + np.testing.assert_array_equal(result[0].new().value, [[big + 1, 8], [1, 2]]) + # BOOL with UINT64 is UINT64, so True // 2**63 is 0, not -1 through int64. + result = binary.floordiv(True, u).new() + assert result.dtype == uints + np.testing.assert_array_equal(result[0].new().value[0], [0, 0]) + record = dtypes.register_anonymous( + np.dtype([("mxs_u", np.uint64), ("mxs_f", np.float64)]), "_MixRec" + ) + other = dtypes.register_anonymous( + np.dtype([("mxs_u", np.int64), ("mxs_f", np.float64)]), "_MixRecI" + ) + r = Vector(record, size=1) + r[0] = (big, 1.5) + s = Vector(other, size=1) + s[0] = (-1, 0.5) + assert (r * np.int64(2)).new()[0].new().value.tolist() == (2.0 * big, 3.0) + assert binary.times(np.int64(2), r).new()[0].new().value.tolist() == (2.0 * big, 3.0) + # A weak 2 keeps the uint64 field, which wraps as numpy 2's uint64 does. + assert (r * 2).new()[0].new().value.tolist() == (0, 3.0) + assert binary.plus(r & s).new()[0].new().value.tolist() == (big - 1.0, 2.0) + assert binary.floordiv(s & r).new()[0].new().value.tolist() == (-1.0, 0.0) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_user_op_named_like_a_lifted_op(): + """A UDF registered under a lifted op's name runs its own function, and literals become UDTs. + + ``_compile_udt`` matched built-in ops by name only, so a UDF registered as + ``"plus"`` or ``"max"`` was compiled as the built-in, its function ignored, + and on a built-in dtype it raised KeyError. A literal beside a UDT is + converted to the UDT for it, since nothing says how a UDF combines a scalar + with an element, but only when that changes no value: ``0.5`` became an + element of zeros. + """ + udt = dtypes.register_anonymous(np.dtype((np.int64, (2, 3))), "_NamedLikeLifted") + + def scaled_plus(x, y): + return x * 10 + y + + def difference(x, y): + return x - y + + plus_like = BinaryOp.register_anonymous(scaled_plus, "plus", is_udt=True) + max_like = BinaryOp.register_anonymous(difference, "max", is_udt=True) + v = Vector(udt, size=1) + v[0] = np.ones((2, 3), dtype=np.int64) + for udf_op, right, expected in [ + (plus_like, 3, 13), + (plus_like, 2.0, 12), + (max_like, 3, -2), + ]: + result = v.apply(udf_op, right=right).new() + assert result.dtype == udt + np.testing.assert_array_equal(result[0].new().value, np.full((2, 3), expected)) + assert udf_op[udt].jit_c_source is None + with pytest.raises(ValueError, match="would change 0.5 to 0"): + v.apply(plus_like, right=0.5) + assert plus_like(Vector.from_coo([0], [5]), 2).new().isequal(Vector.from_coo([0], [52])) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_python_number_is_weak_per_leaf(): + """A Python number beside a UDT is typed leaf by leaf, as numpy 2 types it. + + Each leaf keeps its dtype when its kind holds the number (an int beside + int8, a float beside float32) and takes the number's default dtype when it + does not (a float beside int8 is float64). So an in-place update with a + number of the leaf's kind keeps the UDT, where typing every Python int as + INT64 and every float as FP64 made ``int8_udt += 1`` and ``f32_udt *= 2`` + unstorable. numpy scalars are strong, and so are typed Scalars. + """ + + # Aligned, so the same-type ops get a JIT kernel and warn nothing. + def record(i_dtype, f_dtype): + return np.dtype([("wk_i", i_dtype), ("wk_f", f_dtype)], align=True) + + rec = dtypes.register_anonymous(record(np.int8, np.float32), "_WeakRec") + v = Vector(rec, size=1) + v[0] = (1, 2.0) + v += 1 + v *= 2 + assert v.dtype == rec + assert v[0].new().value.tolist() == (4, 6.0) + result = (v + 0.5).new() + assert result.dtype.np_type == record(np.float64, np.float32) + assert result[0].new().value.tolist() == (4.5, 6.5) + assert (v + np.float32(0.5)).new().dtype.np_type == record(np.float32, np.float32) + # The result has a float64 leaf, which is cast back to int8 to be stored. + v += 0.5 + assert v.dtype == rec + assert v[0].new().value.tolist() == (4, 6.5) + with pytest.raises(OverflowError, match="200 out of bounds for int8"): + v + 200 + # A UDT Scalar follows the same rule, through infix and apply alike. + arr = dtypes.register_anonymous(np.dtype((np.int16, (5,))), "_WeakArr") + s = gb.Scalar(arr) + s.value = [1, 2, 3, 4, 5] + for expr in [s * 2.5, 2.5 * s, s.apply(binary.times, right=2.5)]: + result = expr.new() + assert result.dtype.np_type == np.dtype((np.float64, (5,))) + np.testing.assert_array_equal(result.value, np.arange(1, 6) * 2.5) + assert (s + 1).new().dtype == arr + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_infix_plus_on_mixed_types_is_the_union_with_zero(): + """``ewise_add`` with a lifted op on UDT operands of different types runs without a cast. + + GraphBLAS casts an entry present in only one input to the result type, + which a UDT cannot be: a UDT Scalar plus ``0.5`` (infix ``+`` is + ``ewise_add``), an int8 array UDT vector plus a float32 one, and + ``truediv`` on two int UDT vectors of the same type (float64 result) all + raised a bare GrB_DOMAIN_MISMATCH. Each lifted op gives the same values on + operands already converted to its result type, so ``ewise_add`` converts + them (subtracting zero, which keeps ``-0.0``) and applies the op there. + """ + int8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_UniPlusI8") + fp32s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_UniPlusF32") + s = gb.Scalar.from_value([1, 2, 3], dtype=int8s) + for expr in [s + 0.5, 0.5 + s]: + result = expr.new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + np.testing.assert_array_equal(result.value, [1.5, 2.5, 3.5]) + result = (gb.Scalar(int8s) + 0.5).new() # an empty Scalar adds as 0 + np.testing.assert_array_equal(result.value, [0.5, 0.5, 0.5]) + v = Vector(int8s, size=3) + v[0] = [1, 2, 3] + v[1] = [4, 5, 6] + w = Vector(fp32s, size=3) + w[1] = [0.5, 0.5, 0.5] + w[2] = [7, 8, 9] + for a, b in [(v, w), (w, v)]: + result = (a + b).new() + assert result.dtype == fp32s + assert result.nvals == 3 + np.testing.assert_array_equal(result[0].new().value, [1, 2, 3]) + np.testing.assert_array_equal(result[1].new().value, [4.5, 5.5, 6.5]) + np.testing.assert_array_equal(result[2].new().value, [7, 8, 9]) + A = Matrix(int8s, nrows=1, ncols=2) + A[0, 0] = [1, 2, 3] + B = Matrix(fp32s, nrows=1, ncols=2) + B[0, 1] = [0.5, 0.5, 0.5] + result = (A + B).new() + assert result.dtype == fp32s + assert result.nvals == 2 + # A broadcast of the wrong size raises up front, as for a same-type pair, + # not once both operands have been converted. + with pytest.raises(DimensionMismatch, match="Matrix.nrows"): + v.ewise_add(B, binary.plus) + with pytest.raises(DimensionMismatch, match="Matrix.ncols"): + B.ewise_add(v, binary.plus) + # A same-type pair still goes through ewise_add itself. + assert (v + v).new().dtype == int8s + # The method does the same, for each lifted op; an unpaired entry is kept as is. + assert v.ewise_add(w, binary.plus).new().isequal((v + w).new()) + for binop, both in [(binary.minus, [3.5, 4.5, 5.5]), (binary.max, [4, 5, 6])]: + result = v.ewise_add(w, binop).new() + assert result.dtype == fp32s + np.testing.assert_array_equal(result[0].new().value, [1, 2, 3]) + np.testing.assert_array_equal(result[1].new().value, both) + np.testing.assert_array_equal(result[2].new().value, [7, 8, 9]) + u = Vector(int8s, size=3) + u[1] = [2, 2, 2] + u[2] = [-3, 0, 4] + result = v.ewise_add(u, binary.truediv).new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + assert result.to_coo()[1].tolist() == [[1, 2, 3], [2, 2.5, 3], [-3, 0, 4]] + # Converting keeps the sign of a zero, as GraphBLAS's cast does. + f64s = dtypes.lookup_dtype(np.dtype((np.float64, (3,)))) + neg = Vector(fp32s, size=2) + neg[0] = [-0.0, 0.0, -1.5] + pos = Vector(f64s, size=2) + pos[1] = [1, 1, 1] + result = (neg + pos).new() + assert result.dtype == f64s + np.testing.assert_array_equal(np.signbit(result[0].new().value), [True, False, True]) + # Any other op cannot run in the result type, so it is a clear error up front. + udf = BinaryOp.register_anonymous(lambda x, y: x, "_uni_first", is_udt=True) + with pytest.raises(DomainMismatch, match="ewise_add cannot use .* on _UniPlusI8 and"): + v.ewise_add(Vector(f64s, size=3), udf) + # A result of another type is cast to the object's type before it is + # accumulated, as for built-in types. + x = v.dup() + x(accum=binary.plus) << x + w + assert x.dtype == int8s + assert x.to_coo()[1].tolist() == [[2, 4, 6], [8, 10, 12], [7, 8, 9]] + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_literal_converts_only_when_exact(): + """A literal that has to become an element of a UDT must keep every value. + + Tuples, lists and arrays beside a UDT, literals beside a user-defined op, + and ``ewise_union`` defaults are converted into the UDT, which numpy does + silently: ``(0.5, 1.5, 2.5)`` added ``(0, 1, 2)`` to an int8 array UDT, and + a default of ``0.5`` stood in as ``0``. A Monoid in ``apply`` is typed as + its BinaryOp, so its literal promotes instead of being converted. + """ + arr = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_ExactArr") + rec = dtypes.register_anonymous( + np.dtype([("ex_i", np.int16), ("ex_f", np.float32)], align=True), "_ExactRec" + ) + v = Vector(arr, size=3) + v[0] = [1, 2, 3] + w = Vector(arr, size=3) + w[1] = [4, 5, 6] + r = Vector(rec, size=1) + r[0] = (1, 1.5) + # Exact values convert, including a float of a narrower float field. + plus = binary.plus + assert v.apply(plus, right=(1, 2, 3)).new()[0].new().value.tolist() == [2, 4, 6] + assert v.apply(plus, right=[2.0, 2.0, 2.0]).new().dtype == arr + assert (v + np.ones(3, dtype=np.int64)).new().dtype == arr + expected = (2, np.float32(1.5) + np.float32(0.1)) + assert r.apply(plus, right=(1, 0.1)).new()[0].new().value.tolist() == expected + assert r.apply(binary.plus, right={"ex_i": 2, "ex_f": 0.5}).new()[0].new().value.tolist() == ( + 3, + 2.0, + ) + for expr in [ + lambda: v.apply(plus, right=(0.5, 1.5, 2.5)), + lambda: v.apply(binary.times, left=np.array([0.5, 1, 1])), + lambda: r.apply(plus, left=(0.5, 1)), + lambda: v.ewise_union(w, binary.plus, 0.5, 0), + lambda: v.ewise_union(w, binary.plus, 0, 0.5), + lambda: gb.Scalar.from_value([1, 2, 3], arr).ewise_union( + gb.Scalar(arr), binary.plus, 0, 0.5 + ), + lambda: v == (0.5, 2, 3), + ]: + with pytest.raises(ValueError, match="would change 0.5 to 0"): + expr() + with pytest.raises(OverflowError, match="300 out of bounds for int8"): + v.apply(plus, right=(1, 2, 300)) + # Defaults that fit still work, and fill in for the missing side. + result = v.ewise_union(w, binary.plus, 10, 20).new() + assert result.to_coo()[1].tolist() == [[21, 22, 23], [14, 15, 16]] + # apply with a Monoid promotes as its BinaryOp does, on either side. + for expr in [v.apply(monoid.plus, right=0.5), v.apply(monoid.plus, left=0.5)]: + result = expr.new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + np.testing.assert_array_equal(result[0].new().value, [1.5, 2.5, 3.5]) + assert v.apply(monoid.max, right=2).new().dtype == arr + + +@pytest.mark.skipif("not supports_udfs") +# SS < 9 names a UDT by its numpy repr and warns when that is over 128 chars, +# as the records with a field per dtype pair are. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_udt_store_casts_each_element_as_builtin_types_do(monkeypatch): + """A UDT result stored in an object of another UDT type is cast element by element. + + GraphBLAS cannot cast a UDT, so storing across UDT types raised, and + setting one element from a Scalar of another UDT type read its bytes as the + object's type. Each array element or record leaf is now cast as GraphBLAS + casts built-in types on store (compared below), for every way of storing: + ``<<``, a mask, an accumulator, ``replace``, ``.new(dtype=)``, Scalars and + assignment. Layouts that do not correspond, and a built-in type on either + side, raise ``DomainMismatch``. + """ + # One record casts every pair of leaf dtypes, so it compiles once. + pairs = [ + ("f8", "i1"), + ("f8", "u1"), + ("f8", "i8"), + ("f8", "u8"), + ("f8", "?"), + ("f8", "f4"), + ("f4", "u4"), + ("i8", "i1"), + ("u8", "i8"), + ("i1", "u8"), + ("?", "f8"), + ] + if dtypes._supports_complex: + pairs += [("c16", "f8"), ("c16", "i2"), ("c16", "?"), ("i4", "c8")] + samples = { + "f": [np.nan, np.inf, -np.inf, 300.7, -1.7, 2.0**63, -(2.0**64)], + "c": [complex(np.nan, 1), 1 + 2j, -1.7 - 3j, 300.5j, 2.0**64, 0j, -2.5], + "i": [2**63 - 1, -(2**63), 128, -129, -1, 0, 255], + "u": [2**64 - 1, 2**63, 255, 256, 1, 0, 65535], + "b": [True, False, True, True, False, True, False], + } + src_type = np.dtype([(f"sc{i}", src) for i, (src, _dst) in enumerate(pairs)]) + dst_type = np.dtype([(f"sc{i}", dst) for i, (_src, dst) in enumerate(pairs)]) + values = np.zeros(7, dtype=src_type) + for i, (src, _dst) in enumerate(pairs): + kind = np.dtype(src).kind + values[f"sc{i}"] = np.array(samples[kind], np.uint64 if kind == "u" else None).astype(src) + x = Vector.from_coo(np.arange(7), values, dtype=dtypes.register_anonymous(src_type)) + z = Vector(dtypes.register_anonymous(dst_type), size=7) + z << x + cast = z.to_coo()[1] + for i, (src, dst) in enumerate(pairs): + expected = Vector(dst, size=7) + expected << Vector.from_coo(np.arange(7), values[f"sc{i}"], dtype=src) + np.testing.assert_array_equal(cast[f"sc{i}"], expected.to_coo()[1], err_msg=f"{src} {dst}") + + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3, 2))), "_StoreCastI8") + f64s = dtypes.register_anonymous(np.dtype((np.float64, (3, 2))), "_StoreCastF64") + v = Vector(f64s, size=3) + v[0] = [[1.7, -1.7], [300.0, np.nan], [-np.inf, 2.5]] + v[2] = np.full((3, 2), 9.5) + first = [[1, -1], [127, 0], [-128, 2]] + nines = [[9, 9]] * 3 + w = Vector(i8s, size=3) + w << v + assert w.to_coo()[1].tolist() == [first, nines] + assert (v * 1).new(dtype=i8s).isequal(w) + # The mask, accumulator and replace act on the cast values, as for built-in + # types: ``1 + 127`` wraps in int8. + mask = Vector.from_coo([0, 1], [True, True], size=3) + ones = [[1, 1]] * 3 + w = Vector.from_coo([0, 1, 2], [ones] * 3, dtype=i8s) + w(mask.S, accum=binary.plus) << v + assert w.to_coo()[1].tolist() == [[[2, 0], [-128, 1], [-127, 3]], ones, ones] + w(mask.S) << v * 1 # an expression, computed into a temporary first + assert w.to_coo()[1].tolist() == [first, ones] + w(mask.S, replace=True) << v + assert w.to_coo()[1].tolist() == [first] + # Scalars, as GraphBLAS scalars or not, and assignment of one element or many. + for is_cscalar in [True, False]: + s = gb.Scalar(i8s, is_cscalar=is_cscalar) + s << v[0].new() + assert s.value.tolist() == first + s << v.reduce(monoid.plus) # summed in float64, then cast + assert s.value.tolist() == [[11, 7], [127, 0], [-128, 12]] + assert v[0].new(is_cscalar=is_cscalar).dup(i8s).value.tolist() == first + assert v.dup(i8s).isequal((v * 1).new(dtype=i8s)) + w = Vector(i8s, size=3) + w[1] = v[0].new() # read as int8 bytes before + w[[0, 2]] = v[2].new() + assert w.to_coo()[1].tolist() == [nines, first, nines] + A = Matrix.from_coo([0], [1], [v[0].new().value], dtype=f64s, nrows=2, ncols=2) + C = Matrix(i8s, nrows=2, ncols=2) + C << A.T + assert C.to_coo()[2].tolist() == [first] + C[0, :] = v[[0, 1]].new() + assert C[0, 0].new().value.tolist() == first + # Leading axes of length 1 do not change the layout. + lead = dtypes.register_anonymous(np.dtype((np.float32, (1, 3, 2))), "_StoreCastLead") + u = Vector(lead, size=3) + u << w + assert u.to_coo()[1].tolist() == [[nines], [first], [nines]] + + other_shape = Vector(dtypes.register_anonymous(np.dtype((np.int16, (6,)))), size=3) + with pytest.raises(DomainMismatch, match="same shape, apart from leading axes"): + other_shape << v + renamed = np.dtype([(f"sr{i}", dst) for i, (_src, dst) in enumerate(pairs)]) + with pytest.raises(DomainMismatch, match="same field names"): + Vector(dtypes.register_anonymous(renamed), size=7) << x + for expr in [ + lambda: Vector(f64s, size=3) << Vector.from_coo([0], [1.5], size=3), + lambda: Vector(dtypes.FP64, size=3) << v, + lambda: w.__setitem__(0, gb.Scalar.from_value(1.5)), + ]: + with pytest.raises(DomainMismatch, match="a UDT casts only to another UDT"): + expr() + # The switch back to hard: every UDT store across types raises. + monkeypatch.setattr(gb.core.base, "_udt_store_casts", False) + for expr in [lambda: w << v, lambda: w.__setitem__(0, v[0].new())]: + with pytest.raises(DomainMismatch, match="UDTs do not cast"): + expr() + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_udt_records_differing_only_in_layout_combine(): + """A packed and an aligned record with the same leaves add, either way round. + + The result is the left operand's type, since it holds every leaf. Converting + the other operand to it for ``ewise_add`` subtracted a zero, which came back + as the other operand's own type, so the conversion recursed forever. + """ + fields = [("lay_f", np.float64), ("lay_i", np.int16), ("lay_n", [("lay_k", np.int32)])] + packed = dtypes.register_anonymous(np.dtype(fields), "_LayPacked") + aligned = dtypes.register_anonymous(np.dtype(fields, align=True), "_LayAligned") + x = Vector(packed, size=3) + x[0] = (1.5, 2, (3,)) + y = Vector(aligned, size=3) + y[0] = (-0.0, 1, (1,)) + y[2] = (10.0, 20, (30,)) + for a, b in [(x, y), (y, x)]: + for expr in [a + b, a.ewise_add(b, binary.max), a.ewise_add(b)]: + result = expr.new() + assert result.dtype == a.dtype + assert result.nvals == 2 + assert result[2].new().value.tolist() == (10.0, 20, (30,)) + assert (a + b).new()[0].new().value.tolist() == (1.5, 3, (4,)) + s = gb.Scalar.from_value((1.5, 2, (3,)), dtype=packed) + t = gb.Scalar.from_value((1.0, 1, (1,)), dtype=aligned) + assert (s + t).new().value.tolist() == (2.5, 3, (4,)) + # Converting keeps a negative zero where there is no partner. + z = Vector(aligned, size=1) + z[0] = (-0.0, 0, (0,)) + result = (Vector(packed, size=1) + z).new() + assert np.signbit(result[0].new().value["lay_f"]) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_record_array_field_truediv_is_float64(): + """``truediv`` on a small integer array field divides in float64, as a scalar field does. + + Numba runs an array field through a ufunc whose loop it picks itself, and + for int8 or int16 it picked float32, so ``1 / 46`` came back as + 0.021739130839705467 in an FP64 field. + """ + for np_type in [np.int8, np.uint8, np.int16, np.uint16]: + name = f"_TdArr{np.dtype(np_type).name}" + rec = dtypes.register_anonymous( + np.dtype([(f"td_a_{name}", np_type, (2,)), (f"td_s_{name}", np_type)]), name + ) + v = Vector(rec, size=1) + v[0] = ([1, 46], 46) + w = Vector(rec, size=1) + w[0] = ([46, 7], 3) + result = (v / w).new()[0].new().value + np.testing.assert_array_equal(result[0], np.array([1, 46]) / np.array([46, 7])) + assert result[1] == 46 / 3 + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_typed_scalar_operands_and_defaults(): + """Typed Scalars keep their values beside a UDT, as literals do. + + ``Scalar.ewise_union`` converted its other operand into the left UDT + (``s8 - sf`` was int8), and an ``ewise_union`` default given as a typed + Scalar was converted into the operand's UDT however much that changed it. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_TsI8") + f8s = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_TsF8") + s8 = gb.Scalar.from_value([1, 2, 3], dtype=i8s) + sf = gb.Scalar.from_value([0.5, 0.5, 0.5], dtype=f8s) + for expr, expected in [(s8 - sf, [0.5, 1.5, 2.5]), (sf - s8, [-0.5, -1.5, -2.5])]: + result = expr.new() + assert result.dtype == f8s + np.testing.assert_array_equal(result.value, expected) + v = Vector(i8s, size=2) + v[0] = [1, 2, 3] + w = Vector(i8s, size=2) + w[1] = [4, 5, 6] + A = Matrix(i8s, nrows=1, ncols=2) + A[0, 0] = [1, 2, 3] + B = Matrix(i8s, nrows=1, ncols=2) + B[0, 1] = [4, 5, 6] + for a, b in [(v, w), (A, B)]: + with pytest.raises(ValueError, match="would change 0.5 to 0"): + a.ewise_union(b, binary.plus, gb.Scalar.from_value(0.5), 0) + with pytest.raises(ValueError, match="would change 0.5 to 0"): + a.ewise_union(b, binary.plus, 0, sf) + assert a.ewise_union(b, binary.plus, gb.Scalar.from_value(2.0), 0).new().dtype == i8s + # Matrix.apply with a Monoid promotes as Vector.apply does. + result = A.apply(monoid.plus, right=0.5).new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_errors_name_the_real_problem(): + """Errors that the UDT rules had made misleading. + + Infix on two UDTs of different sizes said ``ewise_add cannot use + binary.first``, the op it uses to provoke the size error, and an op that + does not take UDTs at all reported the literal check's ValueError. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_ErrI8") + f8s = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_ErrF8") + v = Vector(i8s, size=2) + w = Vector(f8s, size=3) + for expr in [lambda: binary.plus(v | w), lambda: binary.first(v | w), lambda: v + w]: + with pytest.raises(DimensionMismatch): + expr().new() + with pytest.raises(KeyError, match="lt does not work with _ErrI8"): + v.apply(binary.lt, right=2.5) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_literal_rules_edge_cases(): + """Edges of the weak rule and of the exactness check, each pinned on its own. + + numpy 2 raises OverflowError itself when ``300`` is converted to int8, which + hid that the weak rule's own range check matters on numpy 1, where the + conversion wraps. + """ + from decimal import Decimal + from fractions import Fraction + + from graphblas.core.operator.udt_utils import _weak_literal_udt + + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_EdgeI8") + f4s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_EdgeF4") + c8s = dtypes.register_anonymous(np.dtype((np.complex64, (3,))), "_EdgeC8") + with pytest.raises(OverflowError, match="300 out of bounds for int8"): + _weak_literal_udt(i8s, 300) + assert _weak_literal_udt(i8s, 1) is i8s + assert _weak_literal_udt(f4s, 1j).np_type == np.dtype((np.complex64, (3,))) + v = Vector(i8s, size=1) + v[0] = [1, 2, 3] + for value in [Fraction(1, 2), Decimal("0.5")]: + with pytest.raises(ValueError, match="would change"): + v + value + assert (v + Fraction(2)).new()[0].new().value.tolist() == [3, 4, 5] + c = Vector(c8s, size=1) + c[0] = [1, 2, 3] + with pytest.raises(ValueError, match="would change 1e[+]300 to"): + c.apply(binary.plus, right=(1e300, 0, 0)) + assert c.apply(binary.plus, right=(1 + 2j, 0, 0)).new()[0].new().value[0] == 2 + 2j + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_store_casts_reach_every_store_path(): + """Each way of storing into a UDT object casts as ``<<`` does, or says why it cannot. + + ``agg.first`` into a Scalar of another UDT handed GraphBLAS the output as it + was, which read the int8 bytes as float64 (24 bytes from a 3-byte + element). Indexed assignment of an expression skipped the cast, a UDT + Scalar's ``s += 0.5`` converted 0.5 into the UDT where a Vector casts the + sum, and ``dup`` failed for a dtype given as a numpy dtype or a string. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_PathI8") + f8s = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_PathF8") + v = Vector(i8s, size=3) + v[0] = [1, 2, 3] + v[2] = [4, 5, 6] + aggregators = [(agg.sum, [5.0, 7.0, 9.0])] + if gb.backend == "suitesparse": + aggregators += [(agg.ss.first, [1.0, 2.0, 3.0]), (agg.ss.last, [4.0, 5.0, 6.0])] + for aggregator, expected in aggregators: + s = gb.Scalar(f8s) + s << v.reduce(aggregator) + np.testing.assert_array_equal(s.value, expected) + with pytest.raises(DomainMismatch, match="a UDT casts only to another UDT"): + gb.Scalar(FP64) << v.reduce(agg.sum) + w = Vector(f8s, size=3) + w[0] = [0.5, 0.5, 0.5] + w[1] = [1.5, 1.5, 1.5] + x = v.dup() + x[[0, 1]] << w[[0, 1]] + assert x.to_coo()[1].tolist() == [[0, 0, 0], [1, 1, 1], [4, 5, 6]] + x[[0, 2]] << (v + 0.5).new()[[0, 2]] + assert x.to_coo()[1].tolist() == [[1, 2, 3], [1, 1, 1], [4, 5, 6]] + s = gb.Scalar.from_value([1, 2, 3], dtype=i8s) + s += 0.5 + np.testing.assert_array_equal(s.value, [1, 2, 3]) + s += 1 + np.testing.assert_array_equal(s.value, [2, 3, 4]) + result = s.dup(dtype=np.dtype((np.float64, (3,)))) + assert result.dtype == f8s + np.testing.assert_array_equal(result.value, [2, 3, 4]) + with pytest.raises(DomainMismatch, match="a UDT casts only to another UDT"): + s.dup(dtype="FP64") + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_jit_kernel_only_for_a_same_type_result(): + """A same-type pair gets the arithmetic JIT kernel only when its result is that type too. + + The kernel declares its result with the operands' type, so for ``truediv`` + on an integer UDT, whose result is float64, it would write integer + quotients into a float64 buffer. That op runs through the cfunc instead. + 64 entries keep SuiteSparse off its iso and short-vector shortcuts, so a + JIT kernel would run. + """ + from graphblas.core.operator.udt_utils import _has_jit_set + + ints = dtypes.register_anonymous(np.dtype((np.int32, (11,))), "_JitDivI11") + N = 64 + v = Vector(ints, size=N) + w = Vector(ints, size=N) + for i in range(N): + v[i] = np.arange(11) + i + w[i] = np.full(11, 4) + result = v.ewise_mult(w, binary.truediv).new() + assert result.dtype.np_type == np.dtype((np.float64, (11,))) + for i in [0, 1, N - 1]: + np.testing.assert_array_equal(result[i].new().value, (np.arange(11) + i) / 4) + if suitesparse and _has_jit_set: + assert binary.truediv[ints].jit_c_source is None + assert binary.plus[ints].jit_c_source is not None + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_eq_compile_failure_is_a_udfparseerror(): + """``eq`` reports a pair the codegen cannot type as a UdfParseError, as ``plus`` does. + + Its cfunc was compiled without the wrapper that turns Numba's errors into + a one-line UdfParseError, so a bytes field against a float field surfaced + as Numba's full TypingError. ``plus`` rejects the pair sooner: bytes and + float64 have no common dtype for the result field. + """ + text = dtypes.register_anonymous(np.dtype([("ueq_a", "S4"), ("ueq_b", np.float64)]), "_UeqS") + num = dtypes.register_anonymous( + np.dtype([("ueq_a", np.float64), ("ueq_b", np.float64)]), "_UeqF" + ) + v = Vector(text, size=1) + v[0] = (b"ab", 1.0) + w = Vector(num, size=1) + w[0] = (1.0, 1.0) + with pytest.raises(UdfParseError, match="binary.eq does not work with"): + binary.eq(v & w).new() + with pytest.raises(KeyError, match=r"elements of \|S4 and float64 have no common type"): + binary.plus(v & w).new() + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_pair_is_never_unified(): + """Two different UDTs reach the op as they are, never unified to one type. + + numpy promotes two records of the same layout to one of them or to a third, + which ``unify`` registered as a new type. An op that does not lift to UDTs + then named that third type in its KeyError, and ``ewise_add`` with a monoid + passed GraphBLAS two operands it rejected with GrB_DOMAIN_MISMATCH. A + Monoid takes one type, so on a mixed pair it runs as its BinaryOp. + """ + X = dtypes.register_anonymous(np.dtype([("uni_a", np.float32), ("uni_b", np.int64)]), "_UniX") + Y = dtypes.register_anonymous(np.dtype([("uni_a", np.int64), ("uni_b", np.float32)]), "_UniY") + with pytest.raises(KeyError, match=r"lt does not work with \(_UniX, _UniY\)"): + binary.lt[X, Y] + v = Vector(X, size=1) + v[0] = (1.0, 2) + w = Vector(Y, size=1) + w[0] = (1, 2.0) + for expr in [v.ewise_add(w, monoid.plus), v.ewise_add(w), monoid.plus(v | w)]: + result = expr.new() + assert result.dtype.np_type == np.dtype([("uni_a", np.float64), ("uni_b", np.float64)]) + assert result[0].new().value.tolist() == (2.0, 4.0) + assert v.ewise_mult(w, monoid.max).new()[0].new().value.tolist() == (1.0, 2.0) + with pytest.raises(TypeError, match="Monoid inputs must be the same dtype"): + monoid.plus[X, Y] @pytest.mark.skipif("not supports_udfs") diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index eeeeac206..024b4117a 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -957,8 +957,7 @@ def test_udt_jit_c_source_introspection(): times_arr = binary.times[arr_udt] src_arr = times_arr.jit_c_source assert src_arr is not None - assert "z->v[0] = (x->v[0]) * (y->v[0])" in src_arr - assert "z->v[10] = (x->v[10]) * (y->v[10])" in src_arr + assert "for (int64_t i = 0 ; i < 11 ; i++) { z->v[i] = (x->v[i]) * (y->v[i]) ; }" in src_arr # Builtin scalar op: no JIT source assert binary.plus[int].jit_c_source is None From 3bb56021631714d61f3e17cd0a15e33d69025b73 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 7 Oct 2026 22:58:53 -0500 Subject: [PATCH 05/54] Cast UDT stores without temporaries, type UDT literals, fold in #592-#595 Store casts: an existing object casts in one pass, and element-wise lifted ops (apply, ewise_mult, ewise_union) write the target type directly. Any other store into another UDT type raises DomainMismatch and suggests obj << expr.new(). ewise_add on mixed-type Vectors/Matrices raises; infix + takes the union with each operand's zero (-0.0 in float fields). Literals beside a UDT: Python numbers and tuples/lists/dicts of them are weak per field (NEP 50), numpy values and Scalars are strong. A value that must become a UDT element (user op literal, select/IndexUnaryOp thunk, ewise_union default) must fit by type; one that fits only by value (0 for a bool field, 2.0 for an int field) still converts, with a DeprecationWarning. Lossy values raise. From #592-#595: min/max (fmin semantics, fixed -0.0 tie), floordiv, integer and complex division, and complex multiplication give the same bits on the Numba and C JIT paths. Array fields in records compute element by element. Also: extracting an element into another UDT type casts it instead of copying raw bytes (a segfault on main); Monoid literals are typed as their BinaryOp's; outer works on mixed types; the user guide's UDT rules are two tables. --- docs/user_guide/udt.rst | 207 ++-- graphblas/core/base.py | 90 +- graphblas/core/infix.py | 15 +- graphblas/core/infixmethods.py | 27 +- graphblas/core/matrix.py | 61 +- graphblas/core/operator/binary.py | 8 +- graphblas/core/operator/udt_utils.py | 917 +++++++++++++---- graphblas/core/operator/utils.py | 13 +- graphblas/core/scalar.py | 359 ++++--- graphblas/core/utils.py | 24 + graphblas/core/vector.py | 53 +- graphblas/tests/test_matrix.py | 5 +- graphblas/tests/test_op.py | 1401 +++++++++++++++++++++++--- graphblas/tests/test_ssjit.py | 175 +++- graphblas/tests/test_vector.py | 5 +- 15 files changed, 2702 insertions(+), 658 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 5c7a597e5..61eb97900 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -134,84 +134,135 @@ For example: plus_times = semiring.plus_times[edge_dtype] # field-wise multiply-add semiring total = v.reduce(agg.sum[edge_dtype]).new() # field-wise reduce -The lift is field-by-field for record UDTs and element-by-element for array -UDTs. Each field or element of the result has the dtype the op gives on the -built-in dtypes, so a lifted op does not narrow: ``truediv`` on an ``INT64[3]`` -UDT gives ``FP64[3]``, as ``truediv`` on two INT64 vectors gives FP64. Field -types that don't support the operation raise ``KeyError`` on the first lookup. -``binary.min``, ``binary.max``, and ``binary.floordiv`` reject UDTs with any -complex leaf (no ordering, no integer modulus); use ``plus``, ``minus``, -``times``, or ``truediv`` for complex arithmetic, or register a custom binary -op. - -A scalar operand combines with every field or element. A vector, ``Scalar``, -numpy scalar or 0-d array of a built-in dtype keeps its dtype, so with ``v`` of -an ``INT8[3]`` UDT, ``v + np.int64(1)`` is ``INT64[3]``. A Python number is -weak, as numpy 2 treats it (NEP 50): in each field or element it takes that -field's dtype when the field's kind can hold it, and its own default dtype -(INT64, FP64 or FC64) when not. So ``v + 1`` stays ``INT8[3]``, ``v * 2.5`` -is ``FP64[3]``, and ``v + 300`` raises ``OverflowError``. A record with an ``INT8`` field and an ``FP32`` -field, plus ``0.5``, has an ``FP64`` field and an ``FP32`` one. python-graphblas -applies this rule itself, so it does not depend on the installed numpy (vectors -of built-in dtypes still type a Python number through numpy, which differs -between numpy 1 and 2). ``eq`` and ``ne`` compare with the number as it is. - -A literal is converted to the UDT instead when it is a whole element (a tuple, -a list, or a numpy array that is not 0-d), when the op is user-defined or a -Monoid, and for the defaults of ``ewise_union``, which must have the op's input -types (a default given as a Scalar of another type is converted too). -``apply`` with a Monoid applies its BinaryOp, so there a literal is typed as -that BinaryOp types it. The conversion must keep every value: -``v + (1, 2, 3)`` works, but ``v + (0.5, 1.5, 2.5)`` raises ``ValueError`` -rather than adding ``(0, 1, 2)``, and ``v.ewise_union(w, binary.plus, 0.5, 0)`` -raises rather than using ``0`` for a missing ``v`` entry. A number may still -round to a narrower float field, as it does beside an ``FP32`` field above. - -The lifted binary ops also take two different UDTs whose shapes match. Two -record UDTs match when they nest the same way, with the same field names in the -same order at every level, and each array field has the same shape in both; -array fields do not broadcast against each other. Two array UDTs match when -their shapes broadcast together as numpy arrays do, and the result has the -broadcast shape: ``FP64[3]`` pairs with ``INT64[2, 3]`` to give ``FP64[2, 3]``, -and ``FP64[3, 1]`` with ``FP64[1, 4]`` gives ``FP64[3, 4]``. ``eq`` is True -when every pair of elements numpy would compare is equal. Any other pair, such -as ``FP64[2, 3]`` with ``FP64[3, 2]``, raises ``KeyError`` on the first lookup. -Field and element dtypes may differ, and each pair of them follows the -built-in dtype rules; a dtype that is not a built-in one, such as bytes, pairs -only with itself. A Monoid takes one type, so ``ewise_add``, ``ewise_mult`` and -``ewise_union`` with a Monoid on two different UDTs use its BinaryOp. - -The result is an operand's UDT when that type holds it, on either side. -Otherwise it is the promoted type: for array UDTs, the structural type such as -``FP64[3]``, and for records, the record with the same names and nesting and -the promoted field dtypes, so ``{"x": INT64, "y": FP64}`` plus -``{"x": FP64, "y": INT64}`` gives ``{"x": FP64, "y": FP64}``. That is the UDT -registered with that layout if there is one, else an anonymous UDT. ``eq`` and -``ne`` compare the values and return ``BOOL``. - -A result stored in an object of another UDT type is cast to that type, as -GraphBLAS casts built-in dtypes when it stores a result: the object's type is -the request. GraphBLAS cannot cast a UDT itself, so python-graphblas casts each -element or record leaf the way GraphBLAS casts that pair of built-in dtypes. -Integers wrap; a float stored as an integer is truncated toward zero, saturates -at the integer type's bounds, and is 0 if NaN; a complex number loses its -imaginary part; and anything nonzero is ``True``. So -``int_udt_vec << int_udt_vec + 0.5``, ``int_udt_vec += 0.5`` and -``(int_udt_vec + 0.5).new(dtype=int_udt)`` truncate, as they do for an ``INT8`` -vector, and so does assigning, accumulating or masking into an object of -another UDT type. The two types must correspond as the operands of a lifted op -must, above (same field names and nesting, same shapes). numpy also casts -records by position whatever their names, and arrays to another length, but -those raise ``DomainMismatch`` here, as does storing a UDT result in an object -of a built-in dtype or the reverse. The cast is an extra pass over the result, -which is first computed in its own type. ``ewise_add`` has a related limit: -GraphBLAS casts a value that has no partner to the op's output type, so it -takes UDT operands only when both have the result's type. With a lifted -arithmetic op, it converts each operand to the result type first, which changes -no value, and applies the op there. So ``v + w`` and -``v.ewise_add(w, binary.max)`` work for two UDTs of different types, and -``truediv`` works on an integer UDT. With any other op it raises -``DomainMismatch``; use ``ewise_union`` with a default for each side instead. +How results are typed +~~~~~~~~~~~~~~~~~~~~~ + +A lifted op types each field or element as it types a Vector of that built-in +dtype, so it never narrows. Below, ``v`` is a Vector of an ``INT8[3]`` UDT. + +.. list-table:: + :header-rows: 1 + :widths: 25 45 30 + + * - Operands + - Rule + - Example + * - One UDT + - Each field or element as on its built-in dtype. + - ``v / v`` is ``FP64[3]`` + * - Two UDTs + - Records pair when they have the same field names, in the same order, at + every level, and the same shape for each field. Array UDTs pair when + their shapes broadcast as numpy's do. Any other pair raises + ``KeyError``. + - ``INT8[3]`` with ``INT64[2, 3]`` gives ``INT64[2, 3]`` + * - Result type + - An operand's type when it holds the result; else the promoted layout, + which is the UDT registered with that layout if there is one, else an + anonymous UDT. + - ``{x: INT64, y: FP64}`` with ``{x: FP64, y: INT64}`` gives + ``{x: FP64, y: FP64}`` + * - A built-in Vector, ``Scalar``, numpy scalar or numpy array + - Strong: keeps its dtype. A scalar combines with every field or element; + an array pairs as an array UDT of its own shape. + - ``v + np.int64(1)`` is ``INT64[3]`` + * - A Python number, or a tuple, list or dict of them laid out like an + element + - Weak, field by field: a field's dtype when its kind holds the number, + else INT64, FP64 or FC64 (numpy 2's NEP 50 rule, on numpy 1 too, and + extended to sequences). An int out of an integer field's range raises + ``OverflowError``. + - ``v + 1`` is ``INT8[3]``, ``v * 2.5`` and ``v + (0.5, 1, 2)`` are + ``FP64[3]``, ``v + 300`` raises + * - ``eq`` and ``ne`` + - Compare the values as given and return ``BOOL``. + - ``v == 300`` is ``False`` + * - A literal that must become an element: beside a user-defined op, an + ``IndexUnaryOp`` thunk, an ``ewise_union`` default + - Its type by the rules above must fit the UDT: the UDT itself, or a + numpy value or ``Scalar`` that casts to it safely. A weak int does not + fit a bool field: write ``True`` or ``False``. A value that does not fit + 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``. + - ``v.ewise_union(w, binary.plus, 0.5, 0)`` raises + +A Monoid used element-wise (``apply``, ``ewise_mult``, ``ewise_union``) uses +its BinaryOp, for literals and for two different UDTs. ``min``, ``max`` and ``floordiv`` reject UDTs +with a complex field; use a custom op for those. + +Each element, including each element of an array field, is computed as the +built-in op computes it, and the Numba cfunc and a C JIT kernel give the same +bits: ``min`` and ``max`` ignore NaN (C's ``fmin`` and ``fmax``) and resolve a +tie between ``-0.0`` and ``0.0`` to ``-0.0`` for ``min`` and ``0.0`` for +``max``, ``floordiv`` is numpy's, with 0 for an integer divided by zero, and +``truediv`` is numpy's (Smith's method for complex numbers), with numpy's +infinities for a zero divisor. A few edge cases can differ from built-in +vectors, where the UDT gives numpy's answer: ``INT64_MIN // -1`` is +``INT64_MIN`` (an ``INT64`` vector gives 0), a complex product or quotient with +an infinite or NaN part, a complex quotient with a ``-0.0`` divisor, and the +signed-zero tie, which SuiteSparse's built-in ``min`` leaves to the C library. + +Storing results +~~~~~~~~~~~~~~~ + +GraphBLAS casts built-in dtypes when it stores a result, but it cannot cast a +UDT, so python-graphblas does it, without ever making a temporary copy of the +data. + +.. list-table:: + :header-rows: 1 + :widths: 40 60 + + * - Store + - What happens + * - Into an object of another UDT type of the same layout: records that + pair as operands do, or arrays of the same shape apart from leading + axes of length 1 + - Each field or element is cast as GraphBLAS casts that pair of built-in + dtypes: integers wrap; a float stored as an integer truncates toward + zero, saturates at the bounds, and is 0 if NaN; a complex number loses + its imaginary part; anything nonzero is ``True``. + * - ...from an existing object: ``w << v``, ``v.dup(dtype=...)``, with a + mask, an accumulator or ``replace`` + - One pass through a cast op. + * - ...from an element-wise lifted op: ``apply``, ``ewise_mult``, + ``ewise_union``, and infix operators that use them + - The op writes the object's type itself, in the same call: + ``int_udt_vec += 0.5`` truncates, as it does for an ``INT8`` vector. + * - ...from anything else: ``mxm``, ``reduce`` into a Vector, ``ewise_add``, + a user-defined op, or a whole Vector or Matrix assigned into part of an + object + - ``DomainMismatch``: this would need a converted copy of the whole + result. Make it explicitly: compute the result, then store it + (``w << expr.new()``); to assign, convert the value with + ``.dup(dtype=...)`` first. + * - ...from a ``Scalar`` or one element + - Computed in its own type, then cast. + * - Between a UDT and a built-in dtype, or layouts that do not correspond + - ``DomainMismatch``. + +For the same reason, ``ewise_add`` on Vectors and Matrices raises +``DomainMismatch`` when an operand's type differs from the result's, as for two +different UDTs or ``truediv`` on an integer UDT: it would copy a value that has +no partner into the result type. ``ewise_union`` and ``ewise_mult`` work. Infix +``+`` on two such operands takes the union instead, with each operand's zero for +a missing value, as ``-`` takes 0. That zero is ``-0.0`` in a float field, so an +entry of one operand alone keeps its value, as ``ewise_add`` would copy it; +only a ``-0.0`` beside an integer field of the other operand becomes ``0.0``. A +Scalar is one element, so ``ewise_add`` converts its operands to the result type +instead. + +.. code-block:: python + + v = Vector(int8x3, size=2) # an INT8[3] UDT + f = Vector(fp32x3, size=2) # an FP32[3] UDT + (v * 2.5).new() # FP64[3]: the literal is weak, 2.5 needs FP64 + (v * f).new() # FP32[3], as INT8 times FP32 + v += 1 # stays INT8[3] + v << v * 2.5 # computes in FP64, stores as INT8 (truncates) + (v + f).new() # FP32[3], by ewise_union; v.ewise_add(f) raises Composite aggregators (``agg.hypot``, ``agg.L1norm``, ``agg.Linfnorm``, ``agg.sum_of_squares``, ``agg.sum_of_inverses``) do *not* auto-lift to UDTs; diff --git a/graphblas/core/base.py b/graphblas/core/base.py index 04fb0ae04..8dac7eda5 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -162,12 +162,6 @@ def _expect_op(self, op, values, *, within, **kwargs): raise TypeError(message) from None -# Whether a UDT result may be stored in an object of another UDT type, cast -# element by element as GraphBLAS casts built-in types. False makes every -# store that would need a UDT cast raise DomainMismatch. -_udt_store_casts = True - - def _store_cast_op(result_dtype, output_dtype): """Return the UnaryOp that converts a result for an object of ``output_dtype``, or ``None``. @@ -186,14 +180,14 @@ def _store_cast_op(result_dtype, output_dtype): from . import _has_numba reason = "UDTs do not cast" - if _udt_store_casts and _has_numba and result_dtype._is_udt and output_dtype._is_udt: + if _has_numba and result_dtype._is_udt and output_dtype._is_udt: from .operator.udt_utils import _udt_cast_error if (reason := _udt_cast_error(result_dtype, output_dtype)) is None: from .operator.unary import _udt_cast_op return _udt_cast_op(result_dtype, output_dtype) - elif _udt_store_casts and _has_numba: + elif _has_numba: reason = "a UDT casts only to another UDT" from ..exceptions import DomainMismatch @@ -204,6 +198,47 @@ def _store_cast_op(result_dtype, output_dtype): ) +def _fused_store_expr(expr, output_dtype): + """Return ``expr`` computing straight into ``output_dtype``, or raise ``DomainMismatch``. + + A result of another UDT type is stored by the op that computes it, which + casts each element as it writes it (``udt_utils.fused_store_op``), so no + temporary of the result's own type is made: data GraphBLAS works on is + often too large to hold twice. That works for an element-wise + computation by a built-in op that lifts to UDTs (``apply``, + ``ewise_mult``, ``ewise_union``). ``ewise_add`` copies an entry present in + only one input without the op, and a semiring or reduction would add in + the narrower type, so those and any other computation raise. + """ + import copy + + from .operator.udt_utils import fused_store_op + + op, cfunc_name = expr.op, expr.cfunc_name + if op is not None and op.opclass == "Monoid": + # Element-wise, a Monoid is its BinaryOp. + op, cfunc_name = op.binaryop, cfunc_name.replace("_Monoid", "_BinaryOp") + fused = None + if any(k in cfunc_name for k in ("_apply", "eWiseMult", "eWiseUnion")): + fused = fused_store_op(op, output_dtype) + if fused is None: + from ..exceptions import DomainMismatch + + method = "extraction" if expr.method_name == "__getitem__" else expr.method_name + what = method if op is None else f"{method} with {op.name}" + raise DomainMismatch( + f"cannot store a result of type {expr.dtype} in an object of type " + f"{output_dtype}: GraphBLAS cannot cast a UDT, and only a built-in op that lifts " + "to UDTs, in apply, ewise_mult or ewise_union, can write another type as it " + f"computes ({what} cannot), so this would need a temporary copy of the whole " + "result. Make it explicitly by storing the computed result, as in " + "`obj << expr.new()`, which casts it in one more pass." + ) + rv = copy.copy(expr) + rv.op, rv.cfunc_name, rv.dtype = fused, cfunc_name, output_dtype + return rv + + def _scalar_accum_op(scalar, accum): """Return the op to accumulate into ``scalar`` with, as ``scalar.ewise_add`` runs it. @@ -456,16 +491,20 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * from .matrix import Matrix, MatrixExpression, TransposedMatrix if type(expr) is TransposedMatrix and type(self) is Matrix: - # Transpose (C << A.T) - expr = MatrixExpression( - "transpose", - "GrB_transpose", - [expr], - expr_repr="{0}", - dtype=expr.dtype, - nrows=expr._nrows, - ncols=expr._ncols, - ) + if (cast_op := _store_cast_op(expr.dtype, self.dtype)) is not None: + # C << A.T of another UDT type: the cast reads A transposed. + expr = expr.apply(cast_op) + else: + # Transpose (C << A.T) + expr = MatrixExpression( + "transpose", + "GrB_transpose", + [expr], + expr_repr="{0}", + dtype=expr.dtype, + nrows=expr._nrows, + ncols=expr._ncols, + ) elif isinstance(expr, InfixExprBase): # w << (v & v) expr = expr._to_expr() @@ -508,15 +547,16 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * # hands GraphBLAS the output as it is: agg.first on another UDT read its # bytes as this object's type. if (cast_op := _store_cast_op(expr.dtype, self.dtype)) is not None: - # Compute in the result's type, then cast as the result is stored, - # so the mask, accumulator and replace apply to the cast values. - # Masking the temporary too only skips values the store drops. if self._is_scalar: + # One element: compute it in its own type, then cast it as it is + # stored, so the accumulator applies to the cast value. result = expr.new(name="s_result", **opts) - else: - result = expr.new(mask=mask, name="result", **opts) - self._update(result.apply(cast_op), mask, accum, replace, opts=opts) - return + self._update(result.apply(cast_op), mask, accum, replace, opts=opts) + return + if expr.cfunc_name is not None: + # A recipe (cfunc_name None) stores through ``<<`` itself, and + # each of its stores is checked there. + expr = _fused_store_expr(expr, self.dtype) if expr.op is not None and expr.op.opclass == "Aggregator": updater = self(mask=mask, accum=accum, replace=replace, **opts) expr.op._new(updater, expr) diff --git a/graphblas/core/infix.py b/graphblas/core/infix.py index bc66836fb..73b3b91b5 100644 --- a/graphblas/core/infix.py +++ b/graphblas/core/infix.py @@ -1,6 +1,6 @@ from .. import backend, binary from ..dtypes import BOOL -from ..exceptions import DimensionMismatch, DomainMismatch +from ..exceptions import DimensionMismatch from ..monoid import land, lor from ..semiring import any_pair from . import automethods, recorder, utils @@ -492,16 +492,9 @@ def _dummy(obj, obj_type): def _mismatched(left, right, method, op): # Create dummy expression to raise on incompatible dimensions - try: - getattr(_dummy(left) if isinstance(left, InfixExprBase) else left, method)( - _dummy(right) if isinstance(right, InfixExprBase) else right, op - ) - except DomainMismatch: - # GraphBLAS reports a UDT that does not cast before the shapes, but the - # shapes are what this is here to report. - raise DimensionMismatch( - f"Dimensions not compatible for {method}: {left.shape} and {right.shape}" - ) from None + getattr(_dummy(left) if isinstance(left, InfixExprBase) else left, method)( + _dummy(right) if isinstance(right, InfixExprBase) else right, op + ) raise DimensionMismatch # pragma: no cover diff --git a/graphblas/core/infixmethods.py b/graphblas/core/infixmethods.py index 67248223a..b0fbf0d3d 100644 --- a/graphblas/core/infixmethods.py +++ b/graphblas/core/infixmethods.py @@ -2,11 +2,26 @@ from ..dtypes import BOOL from .infix import MatrixInfixExpr, ScalarInfixExpr, VectorInfixExpr from .matrix import Matrix, MatrixExpression, MatrixIndexExpr, TransposedMatrix -from .scalar import Scalar, ScalarExpression, ScalarIndexExpr +from .operator import get_typed_op +from .scalar import Scalar, ScalarExpression, ScalarIndexExpr, _ewise_add_needs_cast, _plus_zero from .utils import output_type from .vector import Vector, VectorExpression, VectorIndexExpr +def _ewise_add_casts_udt(left, right, op): + """Whether ``left.ewise_add(right, op)`` would need GraphBLAS to cast a UDT.""" + try: + left_dtype, right_dtype = left.dtype, right.dtype + except TypeError: + # An infix operand that has no dtype here (``x & y`` of non-BOOL); + # ewise_add reports it as before. + return False + if not (left_dtype._is_udt or right_dtype._is_udt): + return False + typed_op = get_typed_op(op, left_dtype, right_dtype, kind="binary") + return _ewise_add_needs_cast(typed_op, left_dtype, right_dtype) + + def call_op(self, other, op, *, outer=False, union=False): type1 = output_type(self) type2 = output_type(other) @@ -14,6 +29,16 @@ def call_op(self, other, op, *, outer=False, union=False): if type1 in types: if type2 in types: if outer: + if _ewise_add_casts_udt(self, other, op): + # ewise_add copies an entry present in only one input into + # the result type, which GraphBLAS cannot do for a UDT, so + # infix takes the union with each operand's zero for the + # missing value, as ``-`` takes 0. Its float fields are + # -0.0, which plus adds without changing any value; it is + # false for lor and lxor too. ``ewise_add`` itself raises. + return self.ewise_union( + other, op, _plus_zero(self.dtype), _plus_zero(other.dtype) + ) return self.ewise_add(other, op) if union: return self.ewise_union(other, op, False, False) diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 4e0f7e895..0a0992b7f 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -29,8 +29,8 @@ _cast_for_store, _check_literal_fits, _check_scalar_fits, - _ewise_add_converts_operands, - _ewise_add_in_result_type, + _ewise_add_cast_error, + _ewise_add_needs_cast, _literal_dtype, _scalar_index, ) @@ -113,6 +113,16 @@ def _power(updater, A, n, op): if n == 1: updater << A return + if (out_dtype := updater.parent.dtype) is not A.dtype and ( + out_dtype._is_udt or A.dtype._is_udt + ): + # The last product is stored by mxm, which cannot cast a UDT as it writes; + # raise before computing the others. + from .base import _fused_store_expr, _store_cast_op + + expr = op(A @ A) + if _store_cast_op(expr.dtype, out_dtype) is not None: + _fused_store_expr(expr, out_dtype) # Use repeated squaring: compute A^2, A^4, A^8, etc., and combine terms as needed. # See `numpy.linalg.matrix_power` for a simpler implementation to understand how this works. # We reuse `result` and `square` outputs, and use `square_expr` so masks can be applied. @@ -713,7 +723,11 @@ def dup(self, dtype=None, *, clear=False, mask=None, name=None, **opts): dtype = self.dtype rv = Matrix(dtype, nrows=self._nrows, ncols=self._ncols, name=name) if not clear: - rv(mask=mask, **opts)[...] = self + if rv.dtype._is_udt or self.dtype._is_udt: + # ``<<``, not assign: a UDT of another type casts in one pass. + rv(mask=mask, **opts) << self + else: + rv(mask=mask, **opts)[...] = self else: if opts: # Ignore opts for now @@ -1934,20 +1948,13 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): f"to rows of Matrix in {method_name}. Matrix.ncols (={self._ncols}) " f"must equal Vector.size (={other._size})." ) - if _ewise_add_converts_operands(op, self.dtype, other.dtype): - # See Vector._ewise_add: convert both operands to the result type, - # then apply the op in that one type, since GraphBLAS cannot cast a UDT. - expr = MatrixExpression( - method_name, - None, - [self, other, _ewise_add_in_result_type, (self, other, op)], - nrows=self._nrows, - ncols=self._ncols, - op=op, - ) + if _ewise_add_needs_cast(op, self.dtype, other.dtype): + # GraphBLAS would report the UDT cast before the shapes. if other.ndim == 2 and self.shape != other.shape: - expr.new(name="") # incompatible shape; raise now - return expr + raise DimensionMismatch( + f"Matrix shapes must match in {method_name}: {self.shape} and {other.shape}" + ) + raise _ewise_add_cast_error(op, self.dtype, other.dtype) if other.ndim == 1: return MatrixExpression( @@ -2123,8 +2130,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): left_dtype = temp_op.type dtype = left_dtype if left_dtype._is_udt else None if type(left_default) is not Scalar: - if dtype is not None: - _check_literal_fits(dtype, left_default) + _check_literal_fits(dtype, left_default) try: left = Scalar.from_value( left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -2144,8 +2150,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): right_dtype = temp_op.type2 dtype = right_dtype if right_dtype._is_udt else None if type(right_default) is not Scalar: - if dtype is not None: - _check_literal_fits(dtype, right_default) + _check_literal_fits(dtype, right_default) try: right = Scalar.from_value( right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -2629,7 +2634,8 @@ def select(self, op, thunk=None): if thunk is None: thunk = False # most basic form of 0 when unifying dtypes if type(thunk) is not Scalar: - dtype = self.dtype if (self.dtype._is_udt and not op.is_positional) else None + # A thunk beside a UDT becomes an element of it, so it must fit, as in apply. + dtype = _literal_dtype(self.dtype, thunk, op) try: thunk = Scalar.from_value(thunk, dtype, is_cscalar=None, name="") except TypeError: @@ -3058,6 +3064,13 @@ def _extract_element( dtype = self.dtype else: dtype = lookup_dtype(dtype) + if dtype is not self.dtype and (dtype._is_udt or self.dtype._is_udt): + # See Vector._extract_element: extract in this type, then cast. + value = self._extract_element(resolved_indexes, None, opts, is_cscalar=is_cscalar) + if result is None: + result = Scalar(dtype, is_cscalar=is_cscalar, name=name) + result << value + return result rowidx, colidx = resolved_indexes.indices if self._is_transposed: rowidx, colidx = colidx, rowidx @@ -3119,8 +3132,6 @@ def _prep_for_extract(self, resolved_indexes): def _assign_element(self, resolved_indexes, value): rowidx, colidx = resolved_indexes.indices - if type(value) is Scalar: - value = _cast_for_store(value, self.dtype) if type(value) is not Scalar: dtype = self.dtype if self.dtype._is_udt else None try: @@ -3133,6 +3144,10 @@ def _assign_element(self, resolved_indexes, value): argname="value", extra_message="Literal scalars also accepted.", ) + if value.dtype._is_udt or self.dtype._is_udt: + # After the conversion: an autocomputed element of another UDT type + # arrives here as a Scalar too. + value = _cast_for_store(value, self.dtype) if value._is_cscalar: if value._empty: call("GrB_Matrix_removeElement", [self, rowidx.index, colidx.index]) diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 43a64110b..25b48b8e3 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -99,12 +99,8 @@ def _make_udt_comparison(dtype, dtype2, *, is_eq): shape_y = None if y_is_scalar else dtype2.np_type.subdtype[1] shape = np.broadcast_shapes(*(s for s in (shape_x, shape_y) if s is not None)) N = reduce(mul, shape) - x_ptr_type = nt.CPointer( - numba.from_dtype(dtype.np_type if x_is_scalar else info_x[1][0]) - ) - y_ptr_type = nt.CPointer( - numba.from_dtype(dtype2.np_type if y_is_scalar else info_y[1][0]) - ) + x_ptr_type = nt.CPointer(numba.from_dtype(dtype.np_type.base)) + y_ptr_type = nt.CPointer(numba.from_dtype(dtype2.np_type.base)) wrapper_sig = nt.void(nt.CPointer(INT8.numba_type), x_ptr_type, y_ptr_type) ns = {} x_setup, x_ref = _array_operand_source("x", shape_x, shape, ns) diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index 14f67ebc7..87805ecbf 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -68,7 +68,11 @@ def _compile_codegen(src, *, func_name, source_label, extra_ns=None): filename, ) code = compile(src, filename, "exec") - namespace = {"min": min, "max": max, "abs": abs} + # Generated code never calls Python's ``min`` / ``max`` (still reachable as + # builtins): they order a NaN by argument position, where SuiteSparse's + # ``fmin`` ignores it, so :func:`_minmax_expr` spells the comparison out, + # with ``signbit`` for the signed-zero tie. + namespace = {"abs": abs, "np": np, "signbit": np.signbit} if _has_numba: namespace["numba"] = numba if extra_ns: @@ -114,17 +118,27 @@ def _compile_codegen(src, *, func_name, source_label, extra_ns=None): # JIT include chain pulls in GraphBLAS.h, so these names are in scope. # Plus, minus, times, truediv, ainv compile natively on _Complex; abs # uses cabs / cabsf (see _c_expr_unary). min, max, floordiv don't make - # sense on complex and are skipped via _op_supports_field_dtype. + # sense on complex; ``_element_types`` rejects them before any codegen. np.dtype(np.complex64): "GxB_FC32_t", np.dtype(np.complex128): "GxB_FC64_t", } -# Ops whose JIT C codegen would produce uncompilable kernels on complex fields -# (no ordering for min/max, no integer-mod for floordiv). +# Ops with no meaning on complex fields (no ordering for min/max, no integer +# mod for floordiv); ``_element_types`` rejects them on both paths. _OPS_NOT_FOR_COMPLEX = frozenset({"min", "max", "floordiv"}) -# C operator equivalents for JIT code generation -_C_INFIX_OPS = {"+": "+", "-": "-", "*": "*", "/": "/", "//": "/"} +# ``_MIN`` spellings from , keyed by itemsize. Signed integer +# division traps on ``MIN / -1``; the generated C tests for it explicitly. +_C_INT_MIN = {1: "INT8_MIN", 2: "INT16_MIN", 4: "INT32_MIN", 8: "INT64_MIN"} + + +def _is_float(np_dtype): + return np_dtype.kind == "f" + + +def _is_signed_int(np_dtype): + return np_dtype.kind == "i" + # Vanilla strips GxB callables but keeps GxB constants, so the bare # ``hasattr`` would lie; gate on the backend too. @@ -381,38 +395,49 @@ def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): ) from None -def _udt_cast_error(dtype, to_dtype): - """Return why an element of UDT ``dtype`` cannot be stored as UDT ``to_dtype``, or ``None``. +def _corresponding_leaves(np_x, np_z): + """Return ``(reason, pairs)`` for putting an element of ``np_x`` into one of ``np_z``. - Records must line up as the operands of a lifted op do - (:func:`_check_udt_pair`): the same field names, order and nesting and - the same shape for each field. Arrays must have the same shape apart from - leading axes of length 1; operands broadcast, but a stored element keeps - every value in its place. numpy casts records by position whatever their - names, and repeats or drops elements of an array field of another length, - which would store values in places nobody chose. Each pair of elements - must be built-in numbers, or the same type. + ``reason`` says why the two layouts do not correspond, or is ``None``, and + then ``pairs`` holds the base dtypes of each pair of leaves. Records must + line up as the operands of a lifted op do (:func:`_check_udt_pair`): the + same field names, order and nesting and the same shape for each field. + Arrays must have the same shape apart from leading axes of length 1; + operands broadcast, but a stored element keeps every value in its place. + numpy casts records by position whatever their names, and repeats or + drops elements of an array field of another length, which would put + values in places nobody chose. """ - np_x, np_z = dtype.np_type, to_dtype.np_type if (np_x.names is None) != (np_z.names is None): - return "a record UDT and an array UDT do not correspond" - if np_x.names is not None: - leaves_x = list(_iter_record_leaves(np_x)) - leaves_z = list(_iter_record_leaves(np_z)) - if [py for py, _c, _d in leaves_x] != [py for py, _c, _d in leaves_z]: - return "records must have the same field names, in the same order, at every level" - if any(x.shape != z.shape for (_, _, x), (_, _, z) in zip(leaves_x, leaves_z, strict=True)): - return "records must have the same shape for each field" - pairs = [(x.base, z.base) for (_, _, x), (_, _, z) in zip(leaves_x, leaves_z, strict=True)] - else: + return "a record UDT and an array UDT do not correspond", None + if np_x.names is None: + if np_x.subdtype is None or np_z.subdtype is None: + return "only record and array UDTs cast", None (base_x, shape_x), (base_z, shape_z) = np_x.subdtype, np_z.subdtype if _strip_leading_ones(shape_x) != _strip_leading_ones(shape_z): - return "array UDTs must have the same shape, apart from leading axes of length 1" - pairs = [(base_x, base_z)] - for x, z in pairs: + return "array UDTs must have the same shape, apart from leading axes of length 1", None + return None, [(base_x, base_z)] + leaves_x = list(_iter_record_leaves(np_x)) + leaves_z = list(_iter_record_leaves(np_z)) + if [py for py, _c, _d in leaves_x] != [py for py, _c, _d in leaves_z]: + return "records must have the same field names, in the same order, at every level", None + pairs = [(x, z) for (_, _, x), (_, _, z) in zip(leaves_x, leaves_z, strict=True)] + if any(x.shape != z.shape for x, z in pairs): + return "records must have the same shape for each field", None + return None, [(x.base, z.base) for x, z in pairs] + + +def _udt_cast_error(dtype, to_dtype): + """Return why an element of UDT ``dtype`` cannot be stored as UDT ``to_dtype``, or ``None``. + + The layouts must correspond (:func:`_corresponding_leaves`), and each pair + of elements must be built-in numbers, or the same type. + """ + reason, pairs = _corresponding_leaves(dtype.np_type, to_dtype.np_type) + for x, z in pairs or (): if x != z and not (x.kind in "biufc" and z.kind in "biufc"): return f"elements of {x} do not cast to {z}" - return None + return reason def _strip_leading_ones(shape): @@ -574,39 +599,163 @@ def _weak_literal_element(value, kind, elem): return _WEAK_LITERAL_DEFAULT[kind] -def _weak_literal_udt(dtype, value): - """Return the UDT a Python number becomes beside UDT ``dtype``: weak, element by element. +_LITERAL_KINDS = {bool: "b", int: "i", float: "f", complex: "c"} + + +class _LiteralDoesNotFit(Exception): + """A tuple or list literal is not laid out like an element of the UDT.""" - The literal takes, in each element or record leaf, the dtype numpy 2 gives - a Python scalar beside that element (:func:`_weak_literal_element`), so the - pair then promotes as two UDTs do: ``int8_udt + 1`` stays ``int8_udt``, and - ``int8_udt * 2.5`` gives float64 elements. A leaf that is not a number - (bytes, datetime) takes the literal's default dtype and fails to pair. + +def _as_python_number(value): + """Return a ``Fraction`` or ``Decimal`` as the int or float it stands for; anything else as is. + + numpy would convert either into an integer field by truncating it (``Fraction(1, 2)`` + becomes 0), so they are typed as the Python number they are worth instead. + """ + from decimal import Decimal + from fractions import Fraction + + if isinstance(value, (Fraction, Decimal)): + try: + return int(value) if value == int(value) else float(value) + except (OverflowError, ValueError): # infinite or NaN + return float(value) + return value + + +def _literal_leaf_values(value, shape): + """Return the Python numbers in literal ``value`` read as an array of ``shape``, or ``None``. + + ``shape`` is ``()`` for a scalar leaf. A number fills any shape, as numpy + broadcasts it. A sequence that does not fill the shape, or holds anything + that is not a Python number, gives ``None``. + """ + if shape: + arr = np.empty(shape, dtype=object) + try: + arr[...] = value + except (TypeError, ValueError): + return None + values = [_as_python_number(x) for x in arr.ravel().tolist()] + else: + values = [_as_python_number(value)] + return values if all(type(x) in _LITERAL_KINDS for x in values) else None + + +def _weak_literal_udt(dtype, value, *, weak=True): + """Return the UDT a Python number, or a tuple or list of them, becomes beside UDT ``dtype``. + + The literal is read as an element of ``dtype``: a number stands in every + leaf, a sequence fills an array UDT or an array field, and a tuple or list + fills a record by position, nested for a nested record. Each leaf then takes + the dtype numpy 2 gives a Python scalar beside that leaf's element + (:func:`_weak_literal_element`), of the widest kind among the leaf's + numbers, so the pair promotes as two UDTs do: ``int8_udt + 1`` and + ``int8_udt + (1, 2, 3)`` stay ``int8_udt``, while ``int8_udt * 2.5`` and + ``int8_udt + (0.5, 1.5, 2.5)`` give float64 elements. A leaf that is not a + number (bytes, datetime) takes the literal's default dtype and fails to + pair. With ``weak=False`` every leaf takes the numbers' default dtype + (INT64, FP64, FC64), as ``eq`` and ``ne`` take a number as it is. A sequence + not laid out like an element gives ``None``. """ from ..dtypes import lookup_dtype - kind = ( - "b" - if isinstance(value, bool) - else "i" if isinstance(value, int) else "f" if isinstance(value, float) else "c" - ) + def leaf_dtype(leaf, part): + base, shape = leaf.subdtype or (leaf, ()) + values = _literal_leaf_values(part, shape) + if values is None: + raise _LiteralDoesNotFit + kind = max((_LITERAL_KINDS[type(x)] for x in values), key="bifc".index) + if weak and base.kind in "biufc": + # One dtype for the leaf; an int out of the element's range raises. + for x in values: + elem = _weak_literal_element(x, kind, base) + else: + elem = _WEAK_LITERAL_DEFAULT[kind] + if elem.kind == "b" and shape and base.kind in "biufc": + # A bool array cannot be a UDT, and a bool compares as 0 or 1. + elem = base + return np.dtype((elem, shape)) if shape else elem + + def record_leaves(np_type, part): + if type(_as_python_number(part)) in _LITERAL_KINDS: + parts = [part] * len(np_type.names) + elif isinstance(part, (tuple, list)) and len(part) == len(np_type.names): + parts = part + else: + raise _LiteralDoesNotFit + for name, item in zip(np_type.names, parts, strict=True): + field = np_type.fields[name][0] + if field.names is not None: + yield from record_leaves(field, item) + else: + yield leaf_dtype(field, item) + + np_type = dtype.np_type + try: + if np_type.names is not None: + literal_type = _promoted_record([np_type], record_leaves(np_type, value)) + else: + literal_type = leaf_dtype(np_type, value) + except _LiteralDoesNotFit: + return None + return dtype if literal_type == np_type else lookup_dtype(literal_type) - def element(elem): - if elem.kind not in "biufc": - return _WEAK_LITERAL_DEFAULT[kind] - return _weak_literal_element(value, kind, elem) + +def _array_literal_udt(dtype, value): + """Return the UDT a numpy array ``value`` is beside UDT ``dtype``, or ``None`` to convert it. + + An array has a dtype, so it is strong, as numpy 2 treats it: + ``int8_udt + np.array([0.5, 1, 2])`` has float64 elements and + ``+ np.array([1, 2, 3])`` int64 ones. Its type is the array layout of its + dtype, in its own shape when that pairs with the UDT's and in the UDT's + shape when it broadcasts to that (``np.array([1])`` beside ``INT8[3]``). + A 0-d structured array or ``np.void`` beside a record is its record dtype. + Anything else, such as a numeric array beside a record, is converted into + the UDT. + """ + from ..dtypes import lookup_dtype np_type = dtype.np_type + if value.dtype.names is not None: + return lookup_dtype(value.dtype) if value.ndim == 0 else None + # A bool array converts as it is (a BOOL array UDT cannot be registered), + # and bool casts safely into every numeric element. + if np_type.names is not None or value.dtype.kind not in "iufc": + return None + shape = np_type.subdtype[1] + if _strip_leading_ones(value.shape) == _strip_leading_ones(shape): + shape = value.shape + else: + try: + np.broadcast_to(value, shape) + except ValueError: + return None + return lookup_dtype(np.dtype((value.dtype, shape))) + + +def _leaf_dtypes(np_type): + """Yield the base dtype of each leaf of ``np_type``: a record's fields, an array's element.""" if np_type.names is not None: - leaves = [] for _py, _c, leaf in _iter_record_leaves(np_type): - elem = element(leaf.base) - leaves.append(np.dtype((elem, leaf.shape)) if leaf.shape else elem) - literal_type = _promoted_record([np_type], iter(leaves)) + yield leaf.base else: - base, shape = np_type.subdtype - literal_type = np.dtype((element(base), shape)) - return dtype if literal_type == np_type else lookup_dtype(literal_type) + yield np_type.base + + +def _casts_safely_into(literal_type, dtype): + """Whether each leaf of ``literal_type`` casts safely (``np.can_cast``) into UDT ``dtype``'s. + + A scalar type fills every leaf; otherwise the layouts must correspond + (:func:`_corresponding_leaves`). + """ + if literal_type == dtype: + return True + np_from, np_to = literal_type.np_type, dtype.np_type + if np_from.names is None and np_from.subdtype is None: + return all(np.can_cast(np_from, leaf, "safe") for leaf in _leaf_dtypes(np_to)) + reason, pairs = _corresponding_leaves(np_from, np_to) + return reason is None and all(np.can_cast(x, z, "safe") for x, z in pairs) def _iter_record_leaves(np_type, python_prefix="", c_prefix=""): @@ -760,58 +909,103 @@ def _is_c_compatible_layout(np_type): return True -def _op_supports_field_dtypes(op_name, np_type): - """Return True if the JIT codegen for ``op_name`` can handle every field - type in ``np_type``. - - Currently the only restriction is complex fields: ``min`` / ``max`` / - ``floordiv`` have no defined semantics on C99 ``_Complex`` (no ordering, - no integer mod), so we skip JIT and let SuiteSparse use the Numba cfunc - instead. The cfunc itself errors on these combinations, so this gate - only moves the failure earlier and gives a clearer message. - - Recurses into nested record fields so a complex field nested inside an - outer record is still detected. - """ - if op_name not in _OPS_NOT_FOR_COMPLEX: - return True - if np_type.names is not None: - return not any(leaf.kind == "c" for _, _, leaf in _iter_record_leaves(np_type)) - if np_type.subdtype is not None: - return np_type.subdtype[0].kind != "c" - return True - - # Numba function generators, called lazily from each op's ``_compile_udt``. if _has_numba: - def _expr_binary(py_op, x_expr, y_expr): - """Python-source builder; sibling of :func:`_c_expr_binary` for JIT C.""" + def _expr_binary(py_op, x_expr, y_expr, x_dtype, y_dtype): + """Python-source builder; sibling of :func:`_c_expr_binary` for the C JIT. + + ``x_expr`` and ``y_expr`` are scalars, never arrays. ``x_dtype`` and + ``y_dtype`` are the numpy dtypes the op computes in for each operand: + the built-in op's input types, which the operands are converted to + first. Three ops consult them: + + - ``min`` / ``max`` follow SuiteSparse's own ``GrB_MIN_FP64``, which + is C99 ``fmin`` and ignores a NaN operand rather than ordering or + propagating it; see :func:`_minmax_expr`. Only floating-point + operands carry the NaN and signed-zero rules, so the dtype decides + which of the two forms is emitted. + - ``//`` on two signed integers. Numba deliberately returns 0 for + ``INT64_MIN // -1`` to dodge the SIGFPE that x86 raises on the + unrepresentable quotient, while numpy wraps to ``INT64_MIN``. Since + ``a // -1`` is exactly ``-a``, routing that divisor through + negation (which wraps) reaches numpy's answer without the trap. + - ``/`` where either side is complex. Numba's complex division raises + ``ZeroDivisionError`` unconditionally, outside the error model's + control, so a zero divisor left the element unwritten while the C + kernel returned numpy's infinities. Spelling out the zero case here + makes both paths match numpy. + + ``/`` on two integers needs nothing: the built-in typing makes its + result float64, as ``truediv`` on two INT64 vectors is FP64, so + ``x / 0`` is an infinity, as there. + """ if py_op in _FUNC_BINARY_OPS: - return f"{py_op}({x_expr}, {y_expr})" + return _minmax_expr(py_op, x_expr, y_expr, x_dtype, y_dtype) + if py_op == "//" and _is_signed_int(x_dtype) and _is_signed_int(y_dtype): + return f"(-{x_expr} if {y_expr} == -1 else {x_expr} // {y_expr})" + if py_op == "/" and "c" in (x_dtype.kind, y_dtype.kind): + return ( + f"({x_expr} / {y_expr} if {y_expr} != 0 " + f"else complex({x_expr}.real / 0.0, {x_expr}.imag / 0.0))" + ) return f"{x_expr} {py_op} {y_expr}" + def _minmax_expr(py_op, x_expr, y_expr, x_dtype, y_dtype): + """Return a Python expression for ``min`` / ``max`` matching C ``fmin``. + + ``GrB_MIN_FP64`` is C99 ``fmin``, so this builder and + :func:`_c_minmax_expr` both reproduce it or ``binary.min`` means one + thing on FP64 and another on a UDT: a NaN operand is ignored unless + both are NaN. C99 leaves the tie between ``-0.0`` and ``0.0`` to the + C library, which differs by platform, so both builders fix it: it + resolves to ``-0.0`` for ``min``, ``0.0`` for ``max``, whichever side + each sits on (Apple's libm answers so; glibc's x86-64 ``fmin`` answers + by position). + + Neither ``min`` nor ``np.fmin`` gets both rules right under Numba. + Python's builtin returns its first argument whenever the comparison + is false, so it keeps a NaN on the left and drops one on the right. + ``np.fmin`` fixes the NaN rule but Numba lowers it without the + signed-zero tie-break. Spelling the comparison out, here and in the C + kernel alike, is what keeps the two execution paths equal on every + platform. + + Integers have neither NaN nor a signed zero, so they take the same + plain comparison the C side emits. + """ + cmp_op = "<" if py_op == "min" else ">" + if not _is_float(x_dtype) and not _is_float(y_dtype): + return f"({x_expr} if {x_expr} {cmp_op} {y_expr} else {y_expr})" + # ``y != y`` is the NaN test on the right operand: when it holds, the + # left one wins whatever it is, which is fmin's "ignore the NaN" rule + # and also (correctly) returns NaN when both are NaN. + tie = f"signbit({x_expr})" if py_op == "min" else f"not signbit({x_expr})" + return ( + f"({x_expr} if ({x_expr} {cmp_op} {y_expr} or {y_expr} != {y_expr}" + f" or ({x_expr} == {y_expr} and {tie})) else {y_expr})" + ) + def _expr_unary(py_op, operand): """Python-source builder; sibling of :func:`_c_expr_unary` for JIT C.""" - if py_op in _FUNC_UNARY_OPS: + if py_op in _FUNC_UNARY_OPS or py_op == "np.abs": return f"{py_op}({operand})" return f"{py_op}{operand}" - def _cast_expr(expr, to_dtype, ns, *, shaped=False): - """Return Numba source converting ``expr`` to numpy dtype ``to_dtype``. + def _cast_expr(expr, to_dtype, ns): + """Return Numba source converting scalar ``expr`` to numpy dtype ``to_dtype``. - ``to_dtype`` of ``None`` leaves ``expr`` alone. A ``shaped`` operand is - an array-valued record field, converted element by element. The - conversion function is added to ``ns``, the generated code's namespace. + ``to_dtype`` of ``None`` leaves ``expr`` alone. The conversion function + is added to ``ns``, the generated code's namespace. """ if to_dtype is None: return expr name = f"_to_{to_dtype.name}" ns[name] = numba.from_dtype(to_dtype) - return f"{expr}.astype({name})" if shaped else f"{name}({expr})" + return f"{name}({expr})" def _make_record_func( - leaf_paths, arity, py_op, *, x_is_scalar=False, y_is_scalar=False, casts=None + leaf_paths, arity, py_op, *, x_is_scalar=False, y_is_scalar=False, leaf_types=None ): """Build a Numba njit function for a record UDT. @@ -825,36 +1019,54 @@ def _make_record_func( When ``x_is_scalar`` or ``y_is_scalar`` is True, that argument is a plain scalar (not a record), so it is used directly for all leaves. - ``casts`` holds an ``(x_dtype, y_dtype, shaped)`` triple per leaf: the - numpy dtype to convert each operand to before the op (``None`` for - none), and whether the leaf is an array field. + A binary op needs ``leaf_types``, one ``(x_cast, y_cast, x_dtype, + y_dtype, result)`` tuple per leaf: the numpy dtype to convert each + operand to before the op (``None`` for none), the dtypes the op then + computes in, which :func:`_expr_binary` keys its semantics on, and the + result leaf's dtype. + + An array-valued leaf is computed element by element, with the same + scalar expression as a scalar leaf or an array UDT element, into an + array of the result leaf's dtype. As a whole-array expression it ran + through Numba's ufuncs, which pick their own loop (``int8 / int8`` in + float32), raise on a complex zero divisor, and take none of the scalar + forms of ``min``, ``max`` and ``//``. """ ns = {} + lines = [] if arity == 2: parts = [] for i, path in enumerate(leaf_paths): - x_cast, y_cast, shaped = casts[i] if casts else (None, None, False) - x_expr = _cast_expr( - "x" if x_is_scalar else f"x{path}", - x_cast, - ns, - shaped=shaped and not x_is_scalar, - ) - y_expr = _cast_expr( - "y" if y_is_scalar else f"y{path}", - y_cast, - ns, - shaped=shaped and not y_is_scalar, - ) - parts.append(_expr_binary(py_op, x_expr, y_expr)) + x_cast, y_cast, x_dtype, y_dtype, result = leaf_types[i] + if not result.shape: + x_expr = _cast_expr("x" if x_is_scalar else f"x{path}", x_cast, ns) + y_expr = _cast_expr("y" if y_is_scalar else f"y{path}", y_cast, ns) + parts.append(_expr_binary(py_op, x_expr, y_expr, x_dtype, y_dtype)) + continue + # Two array leaves have the same shape (``_check_udt_pair``). + ns[f"_r{i}"] = result.base + x_expr = _cast_expr("x" if x_is_scalar else f"x{path}[k]", x_cast, ns) + y_expr = _cast_expr("y" if y_is_scalar else f"y{path}[k]", y_cast, ns) + lines += [ + f" r{i} = np.empty({result.shape}, _r{i})", + f" for k in np.ndindex({result.shape}):", + f" xv = {x_expr}", + f" yv = {y_expr}", + f" r{i}[k] = {_expr_binary(py_op, 'xv', 'yv', x_dtype, y_dtype)}", + ] + parts.append(f"r{i}") sig = "x, y" else: - parts = [_expr_unary(py_op, f"x{path}") for path in leaf_paths] + # Python's ``abs`` does not take an array field in Numba; ``np.abs`` + # takes both. + unary_op = "np.abs" if py_op == "abs" else py_op + parts = [_expr_unary(unary_op, f"x{path}") for path in leaf_paths] sig = "x" body = ", ".join(parts) # Single-leaf tuple needs the trailing comma to remain a tuple. ret = f"({body},)" if len(leaf_paths) == 1 else f"({body})" - src = f"def _op({sig}):\n return {ret}\n" + lines.append(f" return {ret}") + src = f"def _op({sig}):\n" + "\n".join(lines) + "\n" op_func = _compile_codegen( src, func_name="_op", @@ -888,58 +1100,63 @@ def _make_array_wrapper( arity, py_op, *, - x_type=None, + x_type, y_type=None, x_shape=None, y_shape=None, - x_is_scalar=False, - y_is_scalar=False, x_cast=None, y_cast=None, + store_cast=None, ): """Build a cfunc-ready wrapper for an array UDT (element-by-element). ``shape`` and ``base_numba_type`` are the shape and element type of the - result. ``x_type`` and ``y_type`` are the operands' element types, and - ``x_shape`` and ``y_shape`` their shapes, which default to the result's - and broadcast to it; when ``x_is_scalar`` or ``y_is_scalar`` is set, - that side is a plain scalar of that type, broadcast to all elements. + result. ``x_type`` and ``y_type`` are the operands' element types. For + a binary op, ``x_shape`` and ``y_shape`` are the operands' shapes, which + broadcast to the result's; ``None`` makes that side a plain scalar of + its type, broadcast to all elements. ``x_cast`` and ``y_cast`` are numpy dtypes to convert each operand to - before the op, if any. + before the op, if any. ``store_cast``, a ``(computed, stored)`` pair of + numpy dtypes, makes each element computed as ``computed`` and stored as + ``stored``, cast as a store casts (:func:`_cast_leaf_source`); then + ``base_numba_type`` is the stored element type. Returns (wrapper_func, wrapper_sig). """ nt = numba.types ns = {} + + def store(expr): + if store_cast is None: + return expr + computed, stored = store_cast + return _cast_leaf_source(computed, stored, _cast_expr(expr, computed, ns), ns) + size = reduce(mul, shape) - if x_type is None: - x_type = base_numba_type - if x_shape is None: - x_shape = shape - if y_shape is None: - y_shape = shape # A loop, not one statement per element: Numba's typing time grows # faster than linearly with the number of statements, so unrolling took # seconds for a 32 by 32 array and minutes for a 64 by 64 one. arrays = f" z = numba.carray(z_ptr, {size})\n" if arity == 2: - if y_type is None: - y_type = base_numba_type # Each operand is read once per element, into ``xv`` and ``yv``, # so an op that names an operand more than once reads it once. body = "" - for side, operand_shape, cast in [ - ("x", None if x_is_scalar else x_shape, x_cast), - ("y", None if y_is_scalar else y_shape, y_cast), + computed = [] + for side, operand_shape, operand_type, cast in [ + ("x", x_shape, x_type, x_cast), + ("y", y_shape, y_type, y_cast), ]: setup, ref = _array_operand_source(side, operand_shape, shape, ns) arrays += setup body += f" {side}v = {_cast_expr(ref, cast, ns)}\n" - body += f" z[i] = {_expr_binary(py_op, 'xv', 'yv')}\n" + computed.append( + numba.np.numpy_support.as_dtype(operand_type) if cast is None else cast + ) + body += f" z[i] = {store(_expr_binary(py_op, 'xv', 'yv', *computed))}\n" params = "z_ptr, x_ptr, y_ptr" sig = nt.void(nt.CPointer(base_numba_type), nt.CPointer(x_type), nt.CPointer(y_type)) else: - body = f" z[i] = {_expr_unary(py_op, 'x[i]')}\n" + body = f" z[i] = {store(_expr_unary(py_op, 'x[i]'))}\n" params = "z_ptr, x_ptr" arrays += f" x = numba.carray(x_ptr, {size})\n" sig = nt.void(nt.CPointer(base_numba_type), nt.CPointer(x_type)) @@ -957,13 +1174,13 @@ def _make_array_wrapper( # array path with near-identical scaffolding. When a third op family # (ternary, indexbinary, ...) needs the same treatment, fold the three # into one shape-parametrized helper instead of pasting a third copy. - def compile_udt_binary_wrapper(op_name, py_op, dtype, dtype2): + def compile_udt_binary_wrapper(op_name, py_op, dtype, dtype2, out=None): """Compile a built-in element-wise binary op for a UDT. Handles two cases: - - Both sides are UDTs of the same shape: a field-by-field (record) or - element-by-element (array) op. + - Both sides are UDTs that pair (:func:`_check_udt_pair`): a + field-by-field (record) or element-by-element (array) op. - One side is a UDT and the other is a scalar type: the scalar is broadcast to all fields or elements (e.g., ``Point + int`` adds the int to every field). @@ -971,6 +1188,9 @@ def compile_udt_binary_wrapper(op_name, py_op, dtype, dtype2): Each element is computed as the op computes it on the operands' element dtypes (see :func:`_udt_op_types`). + ``out``, a UDT of the result's layout, makes the op store each element + into that type instead, cast as a store casts (:func:`fused_store_op`). + Returns ``(wrapper_func, wrapper_sig, ret_type)``. Raises ``KeyError`` when the dtype combination is not supported, with a clear message for common mistakes like passing two record UDTs with different field @@ -1010,34 +1230,33 @@ def cast(in_dtype, elem): leaves_x = None if x_is_scalar else list(_iter_record_leaves(dtype.np_type)) leaves_y = None if y_is_scalar else list(_iter_record_leaves(dtype2.np_type)) ret_leaves = list(_iter_record_leaves(ret_type.np_type)) - casts = [] - for i, ((_py, _c, leaf), (in_x, in_y)) in enumerate( - zip(leaves, input_types, strict=True) - ): - if leaf.shape and (out := ret_leaves[i][2].base) != in_x: - # An array field runs through a Numba ufunc, which picks its - # own loop: int8 / int8 runs in float32. A scalar field divides - # in float64, so convert array operands to the result's dtype. - in_x = in_y = out - casts.append( - ( - cast(in_x, dtype.np_type if x_is_scalar else leaves_x[i][2]), - cast(in_y, dtype2.np_type if y_is_scalar else leaves_y[i][2]), - bool(leaf.shape), - ) + leaf_types = [ + ( + cast(in_x, dtype.np_type if x_is_scalar else leaves_x[i][2]), + cast(in_y, dtype2.np_type if y_is_scalar else leaves_y[i][2]), + in_x, + in_y, + ret_leaves[i][2], ) + for i, (in_x, in_y) in enumerate(input_types) + ] func = _make_record_func( [py for py, _c, _d in leaves], 2, py_op, x_is_scalar=x_is_scalar, y_is_scalar=y_is_scalar, - casts=casts, + leaf_types=leaf_types, ) sig = (dtype.numba_type, dtype2.numba_type) _compile_udf_for_udt( func, sig, op_kind="binary", op_name=op_name, dtypes=(dtype, dtype2) ) + if out is not None: + wrapper, wrapper_sig = _make_record_store_wrapper( + func, ret_type, out, dtype, dtype2 + ) + return wrapper, wrapper_sig, out numba_ret_type = func.overloads[sig].signature.return_type # Numba can still return a leaf wider than ``ret_type``'s (int8 # plus int8 is int64 in Numba, INT8 here); the wrapper casts each @@ -1047,25 +1266,75 @@ def cast(in_dtype, elem): ) else: # Each side has its own element type, since the bases may differ. - x_type = dtype.np_type if x_is_scalar else info_x[1][0] - y_type = dtype2.np_type if y_is_scalar else info_y[1][0] + x_type = dtype.np_type.base + y_type = dtype2.np_type.base ((in_x, in_y),) = input_types base, shape = ret_type.np_type.subdtype + stored = base if out is None else out.np_type.subdtype[0] wrapper, wrapper_sig = _make_array_wrapper( shape, - numba.from_dtype(base), + numba.from_dtype(stored), 2, py_op, x_type=numba.from_dtype(x_type), y_type=numba.from_dtype(y_type), - x_shape=None if x_is_scalar else dtype.np_type.subdtype[1], - y_shape=None if y_is_scalar else dtype2.np_type.subdtype[1], - x_is_scalar=x_is_scalar, - y_is_scalar=y_is_scalar, + x_shape=None if x_is_scalar else dtype.np_type.shape, + y_shape=None if y_is_scalar else dtype2.np_type.shape, x_cast=cast(in_x, x_type), y_cast=cast(in_y, y_type), + store_cast=None if out is None else (base, stored), ) - return wrapper, wrapper_sig, ret_type + return wrapper, wrapper_sig, ret_type if out is None else out + + def _make_record_store_wrapper(func, result_type, out_type, dtype, dtype2=None): + """Build a cfunc-ready wrapper storing ``func``'s result as UDT ``out_type``. + + ``func`` returns one value per leaf of ``result_type``, the type the op + computes in. Each value becomes that leaf's dtype and is then cast into + the matching leaf of ``out_type`` as a store casts it + (:func:`_cast_leaf_source`), so storing into an object of another type + needs no temporary of the result's own type. + + Returns (wrapper_func, wrapper_sig). + """ + from .base import _input_operand + + nt = numba.types + ns = {"numba_func": func} + lines = [" z = numba.carray(z_ptr, 1)"] + params, args, sig_args = ["z_ptr"], [], [nt.CPointer(out_type.numba_type)] + for var, operand in [("x", dtype), ("y", dtype2)]: + if operand is None: + continue + setup, deref, ptr_type = _input_operand(operand, var) + if setup: + lines.append(setup.rstrip("\n")) + params.append(f"{var}_ptr") + args.append(deref) + sig_args.append(ptr_type) + lines.append(f" r = numba_func({', '.join(args)})") + leaves = zip( + _iter_record_leaves(result_type.np_type), + _iter_record_leaves(out_type.np_type), + strict=True, + ) + for k, ((path, _c, computed), (_p, _cz, stored)) in enumerate(leaves): + if computed.shape: + value = _cast_expr(f"r[{k}][i]", computed.base, ns) + value = _cast_leaf_source(computed.base, stored.base, value, ns) + lines.append(f" for i in np.ndindex({computed.shape}):") + lines.append(f" z[0]{path}[i] = {value}") + else: + value = _cast_leaf_source(computed, stored, _cast_expr(f"r[{k}]", computed, ns), ns) + lines.append(f" z[0]{path} = {value}") + src = f"def _store({', '.join(params)}):\n" + "\n".join(lines) + "\n" + wrapper = _compile_codegen( + src, + func_name="_store", + source_label=f"", + extra_ns=ns, + ) + return wrapper, nt.void(*sig_args) @lru_cache def _saturating_cast(dst): @@ -1152,7 +1421,6 @@ def compile_udt_cast_wrapper(dtype, to_dtype): sig = nt.void( nt.CPointer(numba.from_dtype(base_z)), nt.CPointer(numba.from_dtype(base_x)) ) - ns["np"] = np src = "def _cast(z_ptr, x_ptr):\n" + "\n".join(body) + "\n" wrapper = _compile_codegen( src, @@ -1162,9 +1430,12 @@ def compile_udt_cast_wrapper(dtype, to_dtype): ) return wrapper, sig - def compile_udt_unary_wrapper(op_name, py_op, dtype): + def compile_udt_unary_wrapper(op_name, py_op, dtype, out=None): """Compile a built-in element-wise unary op for a UDT. + ``out`` stores the result as another UDT, as in + :func:`compile_udt_binary_wrapper`. + Returns (wrapper_func, wrapper_sig, ret_type). Raises KeyError if the dtype is not supported. """ @@ -1188,72 +1459,296 @@ def compile_udt_unary_wrapper(op_name, py_op, dtype): _compile_udf_for_udt(func, sig, op_kind="unary", op_name=op_name, dtypes=(dtype,)) numba_ret_type = func.overloads[sig].signature.return_type ret_type = _resolve_udt_return_type(numba_ret_type, dtype) + if out is not None: + wrapper, wrapper_sig = _make_record_store_wrapper(func, ret_type, out, dtype) + return wrapper, wrapper_sig, out wrapper, wrapper_sig = _get_udt_wrapper( func, ret_type, dtype, numba_ret_type=numba_ret_type ) else: base_dtype, shape = dtype.np_type.subdtype ret_type = dtype + stored = base_dtype if out is None else out.np_type.subdtype[0] wrapper, wrapper_sig = _make_array_wrapper( - shape, numba.from_dtype(base_dtype), 1, py_op + shape, + numba.from_dtype(stored), + 1, + py_op, + x_type=numba.from_dtype(base_dtype), + store_cast=None if out is None else (base_dtype, stored), ) - return wrapper, wrapper_sig, ret_type + return wrapper, wrapper_sig, ret_type if out is None else out + + _fused_store_parents = {} + + def fused_store_op(typed_op, out_type): + """Return ``typed_op`` storing its result as ``out_type``, or ``None`` if it cannot. + + The returned op computes what ``typed_op`` computes, in the same type, + and casts each element or leaf into ``out_type`` as it writes it, as a + store casts (``base._store_cast_op``). So a result goes into an object + of another UDT type in the one GraphBLAS call that computes it, with no + temporary. Only the built-in ops that lift to UDTs have such a form; + ``base._store_cast_op`` has already accepted the pair of types. + """ + from .base import _finalize_udt_op + + parent = getattr(typed_op, "parent", None) + if parent is None or parent._anonymous: + return None + if typed_op.opclass == "BinaryOp" and parent.name in BUILTIN_UDT_BINARY_OPS: + from .binary import BinaryOp, TypedUserBinaryOp + + cls, typed_cls = BinaryOp, TypedUserBinaryOp + dtype2 = typed_op.type2 + py_op = BUILTIN_UDT_BINARY_OPS[parent.name] + elif typed_op.opclass == "UnaryOp" and parent.name in BUILTIN_UDT_UNARY_OPS: + from .unary import TypedUserUnaryOp, UnaryOp + + cls, typed_cls = UnaryOp, TypedUserUnaryOp + dtype2 = None + py_op = BUILTIN_UDT_UNARY_OPS[parent.name] + else: + return None + key = (cls, parent.name, out_type) + fused = _fused_store_parents.get(key) + if fused is None: + fused = _fused_store_parents[key] = cls( + f"{parent.name}_as_{out_type.name}", anonymous=True, is_udt=True + ) + dtype = typed_op.type + if (types := dtype if dtype2 is None else (dtype, dtype2)) in fused._udt_ops: + return fused._udt_ops[types] + if dtype2 is None: + wrapper, sig, _ = compile_udt_unary_wrapper(parent.name, py_op, dtype, out=out_type) + else: + wrapper, sig, _ = compile_udt_binary_wrapper( + parent.name, py_op, dtype, dtype2, out=out_type + ) + return _finalize_udt_op(fused, dtype, dtype2, out_type, wrapper, sig, typed_cls) # JIT C code generators below. -def _c_expr_binary(py_op, lhs, rhs, field_dtype=None): +def _c_assign_binary(py_op, target, lhs, rhs, field_dtype): + """Return a C *statement* assigning ``lhs py_op rhs`` to ``target``. + + Almost every op is a single expression, but floating-point ``//`` and + complex ``/`` and ``*`` need temporaries (see :func:`_c_float_floordiv_stmt`, + :func:`_c_complex_truediv_stmt` and :func:`_c_complex_times_stmt`), so this + is the entry point the kernel builder uses rather than :func:`_c_expr_binary`. + """ + if py_op == "//" and _is_float(field_dtype): + return _c_float_floordiv_stmt(target, lhs, rhs, field_dtype) + if py_op == "/" and field_dtype.kind == "c": + return _c_complex_truediv_stmt(target, lhs, rhs, field_dtype) + if py_op == "*" and field_dtype.kind == "c": + return _c_complex_times_stmt(target, lhs, rhs, field_dtype) + return f"{target} = {_c_expr_binary(py_op, lhs, rhs, field_dtype)} ;" + + +def _c_expr_binary(py_op, lhs, rhs, field_dtype): """Return a C expression for a binary op: e.g., ``(x->a) + (y->a)``. - ``field_dtype`` is the numpy dtype of the *result* element. It is only - consulted for ``floordiv`` (``//``), which needs Python ``//`` semantics - rather than C ``/`` (trunc toward zero for ints, true division for - floats). Other ops are type-agnostic at the C level. + ``field_dtype`` is the numpy dtype of the operands and of the element + being written: a C kernel exists only for an op whose result has the + operands' type. Two of the ops need it because the C spelling that looks + obvious disagrees with what the Numba cfunc computes: ``//`` is not C + ``/``, and ``min`` / ``max`` have to ignore NaN and break a signed-zero tie + as :func:`_minmax_expr` does. It is required rather than optional so a new + caller can't silently get the wrong kernel. :func:`_expr_binary` is the + sibling for the Numba cfunc; SuiteSparse picks between the two on its + own, so they must agree. + + Floating-point ``//`` and complex ``/`` and ``*`` have no single-expression + form; go through :func:`_c_assign_binary`. """ - if py_op == "min": - # Match Python ``min(a, b) = b if b < a else a`` so NaN propagates - # from the first operand (cfunc / numba follows the same rule). - # The naive ``(a < b ? a : b)`` would silently swallow NaN to the - # right-hand side and disagree with the cfunc path. - return f"(({rhs}) < ({lhs}) ? ({rhs}) : ({lhs}))" - if py_op == "max": - return f"(({rhs}) > ({lhs}) ? ({rhs}) : ({lhs}))" + if py_op in _FUNC_BINARY_OPS: + return _c_minmax_expr(py_op, lhs, rhs, field_dtype) + if py_op == "/": + return _c_truediv_expr(lhs, rhs, field_dtype) if py_op == "//": - return _c_floordiv_expr(lhs, rhs, field_dtype) - c_op = _C_INFIX_OPS.get(py_op, py_op) - return f"({lhs}) {c_op} ({rhs})" + return _c_int_floordiv_expr(lhs, rhs, field_dtype) + # Everything left (``+``, ``-``, and ``*`` on real fields) spells the same + # in C as in Python and needs no dtype-specific handling. + return f"({lhs}) {py_op} ({rhs})" + + +def _c_minmax_expr(py_op, lhs, rhs, field_dtype): + """Return a C expression for ``min`` / ``max`` on a field, as the Numba cfunc computes it. + + ``GrB_MIN_FP64`` is C99 ``fmin``: it ignores a NaN operand rather than + ordering it, and a UDT field does the same. C99 leaves the tie between + ``-0.0`` and ``0.0`` to the implementation, though: Apple's libm returns + ``-0.0`` for ``fmin``, while glibc's x86-64 ``fmin`` (one ``minsd``) + returns an operand by position. So the comparison is spelled out, the same + as :func:`_minmax_expr` spells it for the cfunc: a NaN on the right loses, + and a tie resolves to ``-0.0`` for ``min`` and ``0.0`` for ``max``, + whichever side each sits on. The C kernel and the cfunc then agree on + every platform; only SuiteSparse's built-in FP64 ``min`` may break that + tie the other way, where the C library does. ``signbit`` comes from + ````, which reaches the C JIT kernel through ``GraphBLAS.h``, the + same route :func:`_c_float_floordiv_stmt` relies on for ``floor`` and + ``fmod``. + + Integers have neither NaN nor a signed zero, so a plain comparison is + exact for them. + """ + cmp_op = "<" if py_op == "min" else ">" + if not _is_float(field_dtype): + return f"(({lhs}) {cmp_op} ({rhs}) ? ({lhs}) : ({rhs}))" + tie = f"signbit({lhs})" if py_op == "min" else f"!signbit({lhs})" + return ( + f"((({lhs}) {cmp_op} ({rhs}) || ({rhs}) != ({rhs}) " + f"|| (({lhs}) == ({rhs}) && {tie})) ? ({lhs}) : ({rhs}))" + ) + + +def _c_truediv_expr(lhs, rhs, field_dtype): + """Return a C expression for true division on a float field. + + Only float and complex fields reach a C kernel. True division of integers + or bools gives float64, as on built-in vectors, so its result never has + the operands' type, which a C kernel needs (``BinaryOp._compile_udt``). + Complex division is a statement (:func:`_c_complex_truediv_stmt`). + """ + if not _is_float(field_dtype): + raise ValueError(f"no C expression divides {field_dtype}") + return f"({lhs}) / ({rhs})" + + +def _c_complex_truediv_stmt(target, lhs, rhs, field_dtype): + """Return a C block dividing two complex fields into ``target``, as the Numba cfunc does. + + C's own ``/`` on ``_Complex`` cannot be used. SuiteSparse's C JIT compiles + with ``-fcx-limited-range`` under GCC, which replaces C99 Annex G division + with the naive formula: it squares the divisor's parts, so ``1e20 / 1e20`` + in complex64 gave ``nan`` and ``(1+1j) / (1e160+1e160j)`` gave ``0`` on + Linux. This block is Numba's complex division, which is CPython's + ``_Py_c_quot`` (Smith's method: scale by the larger part of the divisor), + computed in the field's own precision with real arithmetic only, so no + flag changes it. Each product is its own statement: clang contracts a + multiply and an add within one expression into a fused multiply-add, + which rounds once where Numba rounds twice, and GCC under ``-std=c11`` + does not contract at all. A zero divisor divides the parts by ``0.0`` as + reals, for numpy's infinities, as ``_expr_binary`` spells it for the + cfunc; ``rhs == 0`` on a ``_Complex`` compares both parts, mirroring the + cfunc's ``y != 0``. + + The constructor is SuiteSparse's ``GB_CMPLX32``/``GB_CMPLX64`` rather + than C11's ``CMPLXF``/``CMPLX``, which macOS does not define. + ``GB_complex.h`` is on the C JIT's include path, names that exact gap, and + its fallback builds the value through a two-element array, so it keeps + infinities and NaNs as given. + """ + ctype, cmplx, creal, cimag, fabs, zero = _c_complex_names(field_dtype) + return ( + f"{{ {ctype} gb_ar = {creal} ({lhs}) ; {ctype} gb_ai = {cimag} ({lhs}) ; " + f"{ctype} gb_br = {creal} ({rhs}) ; {ctype} gb_bi = {cimag} ({rhs}) ; " + f"if (gb_br == {zero} && gb_bi == {zero}) " + f"{{ {target} = {cmplx} (gb_ar / {zero}, gb_ai / {zero}) ; }} " + f"else if ({fabs} (gb_br) >= {fabs} (gb_bi)) {{ " + f"{ctype} gb_r = gb_bi / gb_br ; {ctype} gb_t = gb_bi * gb_r ; " + f"{ctype} gb_d = gb_br + gb_t ; {ctype} gb_u = gb_ai * gb_r ; " + f"{ctype} gb_v = gb_ar * gb_r ; " + f"{target} = {cmplx} ((gb_ar + gb_u) / gb_d, (gb_ai - gb_v) / gb_d) ; }} " + f"else if (gb_bi == {zero}) {{ {target} = {cmplx} (NAN, NAN) ; }} " + f"else {{ {ctype} gb_r = gb_br / gb_bi ; {ctype} gb_t = gb_br * gb_r ; " + f"{ctype} gb_d = gb_t + gb_bi ; {ctype} gb_u = gb_ar * gb_r ; " + f"{ctype} gb_v = gb_ai * gb_r ; " + f"{target} = {cmplx} ((gb_u + gb_ai) / gb_d, (gb_v - gb_ar) / gb_d) ; }} }}" + ) + + +def _c_complex_times_stmt(target, lhs, rhs, field_dtype): + """Return a C block multiplying two complex fields into ``target``, as the Numba cfunc does. + + Numba and numpy multiply complex numbers by the textbook formula. C99 + Annex G adds a recovery step that turns some of its ``nan`` results into + infinities, and whether C's ``*`` on ``_Complex`` takes it depends on the + compiler: GCC under SuiteSparse's ``-fcx-limited-range`` does not, clang + does. So ``1j * (inf+infj)`` was ``nan+nanj`` on Linux and the cfunc but + ``-inf+infj`` from a C kernel on macOS. This block is the textbook formula + in real arithmetic, one product per statement so clang cannot contract it + into a fused multiply-add (see :func:`_c_complex_truediv_stmt`). + """ + ctype, cmplx, creal, cimag, _, _ = _c_complex_names(field_dtype) + return ( + f"{{ {ctype} gb_ar = {creal} ({lhs}) ; {ctype} gb_ai = {cimag} ({lhs}) ; " + f"{ctype} gb_br = {creal} ({rhs}) ; {ctype} gb_bi = {cimag} ({rhs}) ; " + f"{ctype} gb_rr = gb_ar * gb_br ; {ctype} gb_ii = gb_ai * gb_bi ; " + f"{ctype} gb_ri = gb_ar * gb_bi ; {ctype} gb_ir = gb_ai * gb_br ; " + f"{target} = {cmplx} (gb_rr - gb_ii, gb_ri + gb_ir) ; }}" + ) + + +def _c_complex_names(field_dtype): + """Return the C type, constructor and part functions for a complex field.""" + if field_dtype.itemsize == 8: + return "float", "GB_CMPLX32", "crealf", "cimagf", "fabsf", "0.0f" + return "double", "GB_CMPLX64", "creal", "cimag", "fabs", "0.0" -def _c_floordiv_expr(lhs, rhs, field_dtype): - """Return a C expression for Python-semantics floor division. +def _c_int_floordiv_expr(lhs, rhs, field_dtype): + """Return a C expression for integer floor division. - Python ``//`` is floor (rounds toward negative infinity); C ``/`` is - trunc toward zero for ints and true division for floats. The two only - agree for non-negative integer operands; for everything else the JIT - path silently disagreed with the Numba cfunc path before this helper. + Three ways C ``/`` differs from what the cfunc computes: - Float fields use ``floor()`` / ``floorf()`` from ````, which is - available in the JIT kernel via SuiteSparse's include chain - (``GraphBLAS.h`` -> ````). Signed integer fields use the - standard trunc-to-floor adjustment. Unsigned integers don't need - adjusting because both operands are non-negative. + - It truncates toward zero; Python ``//`` floors. Subtract one when the + operands have different signs and the remainder is non-zero. + - It raises SIGFPE (process death, not an exception) when the divisor is + zero. Numba's numpy error model returns 0, as does ``np.floor_divide``. + - It raises SIGFPE on ``INT_MIN / -1``, whose true quotient is not + representable. numpy wraps to ``INT_MIN``, which is what the cfunc + reaches through negation (see ``_expr_binary``). """ - if field_dtype is None: - # Caller didn't pass dtype info. The C ``/`` semantics match Python - # ``//`` for non-negative integer operands only. - return f"({lhs}) / ({rhs})" - kind = field_dtype.kind - if kind == "f": - if field_dtype.itemsize == 4: - return f"floorf((float)({lhs}) / (float)({rhs}))" - return f"floor((double)({lhs}) / (double)({rhs}))" - if kind in ("u", "b"): - return f"({lhs}) / ({rhs})" - # Signed integer: trunc-toward-zero is one greater than floor when the - # signs of ``a`` and ``b`` differ and the division has a non-zero - # remainder; subtract 1 in that case. - return f"(({lhs}) / ({rhs}) - ((({lhs}) % ({rhs}) != 0) && ((({lhs}) < 0) != (({rhs}) < 0))))" + if field_dtype.kind in ("u", "b"): + # Both operands are non-negative, so truncation already floors. + return f"(({rhs}) == 0 ? 0 : (({lhs}) / ({rhs})))" + type_min = _C_INT_MIN[field_dtype.itemsize] + floored = f"({lhs}) / ({rhs}) - ((({lhs}) % ({rhs}) != 0) && ((({lhs}) < 0) != (({rhs}) < 0)))" + return ( + f"(({rhs}) == 0 ? 0 : " + f"((({lhs}) == {type_min} && ({rhs}) == -1) ? {type_min} : ({floored})))" + ) + + +def _c_float_floordiv_stmt(target, lhs, rhs, field_dtype): + """Return a C block computing floating-point ``//`` into ``target``. + + ``floor(a / b)`` is not floor division. numpy and CPython both compute + the quotient from the remainder (``(a - fmod(a, b)) / b``) and snap it, + which is more accurate and handles infinities differently: ``1.0 // 0.1`` + is 9.0 but ``floor(1.0 / 0.1)`` is 10.0, and ``inf // 2.0`` is NaN but + ``floor(inf / 2.0)`` is ``inf``. This mirrors numpy's ``npy_divmod`` + (which CPython's ``float_divmod`` matches), so the C JIT kernel, the + Numba cfunc, and ``np.floor_divide`` all agree. + + The sign fix-up needs the remainder twice and the quotient twice, so + this emits a statement with temporaries rather than one expression. + """ + is_f32 = field_dtype.itemsize == 4 + ctype = "float" if is_f32 else "double" + sfx = "f" if is_f32 else "" + half = "0.5f" if is_f32 else "0.5" + one = "1.0f" if is_f32 else "1.0" + zero = "0.0f" if is_f32 else "0.0" + return ( + f"{{ {ctype} gb_a = ({lhs}) ; {ctype} gb_b = ({rhs}) ; {ctype} gb_q ; " + # A zero divisor is the one case numpy answers straight from the + # division: +-inf, or NaN for 0/0. + f"if (gb_b == {zero}) {{ gb_q = gb_a / gb_b ; }} " + f"else {{ {ctype} gb_m = fmod{sfx} (gb_a, gb_b) ; " + f"{ctype} gb_d = (gb_a - gb_m) / gb_b ; " + # NaN compares false both ways, so a NaN remainder skips the + # adjustment and carries through to the quotient. + f"if (gb_m != {zero} && ((gb_b < {zero}) != (gb_m < {zero}))) {{ gb_d -= {one} ; }} " + f"if (gb_d != {zero}) {{ gb_q = floor{sfx} (gb_d) ; " + f"if (gb_d - gb_q > {half}) {{ gb_q += {one} ; }} }} " + f"else {{ gb_q = copysign{sfx} ({zero}, gb_a / gb_b) ; }} }} " + f"{target} = gb_q ; }}" + ) def _c_expr_unary(py_op, operand, field_dtype=None): @@ -1288,13 +1783,11 @@ def _make_jit_c_definition(op_name, py_op, dtype, arity): """Generate a JIT C function definition for a UDT operator. Returns ``(c_name, c_defn)``, or ``None`` if the dtype can't be expressed - in C (unsupported field types, or names that collide with C reserved words), - or if the op isn't defined on the dtype's field types (e.g. ``min`` on a - complex field). + in C (unsupported field types, or names that collide with C reserved words). + An op not defined on the field types (``min`` on a complex field) never + gets here: ``_element_types`` rejects it when the op is typed. """ np_type = dtype.np_type - if not _op_supports_field_dtypes(op_name, np_type): - return None # Prefer the C name SuiteSparse already has for this type. ``GxB_JIT_C_NAME`` # on a ``GrB_Type`` is one-shot, so a later ``dtype.name`` rename does not # propagate to SS. If the op's signature referenced ``dtype.name`` it would @@ -1320,10 +1813,10 @@ def _make_jit_c_definition(op_name, py_op, dtype, arity): # to plain ``z->name``. leaves = list(_iter_record_leaves(np_type)) if arity == 2: - # Pass the leaf dtype to the binary expression builder so - # type-sensitive ops (currently floordiv) can emit correct C. + # Pass the leaf dtype to the binary statement builder so + # type-sensitive ops (min, max, and the divisions) emit correct C. assigns = " ".join( - f"z->{c} = {_c_expr_binary(py_op, f'x->{c}', f'y->{c}', leaf_dtype)} ;" + _c_assign_binary(py_op, f"z->{c}", f"x->{c}", f"y->{c}", leaf_dtype) for _py, c, leaf_dtype in leaves ) else: @@ -1337,10 +1830,12 @@ def _make_jit_c_definition(op_name, py_op, dtype, arity): # A loop, as in the Numba wrapper: one statement per element made the # C compiler take seconds on a 64 by 64 array, once per JIT cache. if arity == 2: - expr = _c_expr_binary(py_op, "x->v[i]", "y->v[i]", base_dtype) + # A statement, not an expression: floating-point ``//`` needs a + # block with temporaries (see ``_c_assign_binary``). + assign = _c_assign_binary(py_op, "z->v[i]", "x->v[i]", "y->v[i]", base_dtype) else: - expr = _c_expr_unary(py_op, "x->v[i]", base_dtype) - assigns = f"for (int64_t i = 0 ; i < {size} ; i++) {{ z->v[i] = {expr} ; }}" + assign = f"z->v[i] = {_c_expr_unary(py_op, 'x->v[i]', base_dtype)} ;" + assigns = f"for (int64_t i = 0 ; i < {size} ; i++) {{ {assign} }}" return c_name, f"void {c_name} ({params}) {{ {assigns} }}" diff --git a/graphblas/core/operator/utils.py b/graphblas/core/operator/utils.py index 2f9dc04e4..37da5cc5f 100644 --- a/graphblas/core/operator/utils.py +++ b/graphblas/core/operator/utils.py @@ -78,10 +78,11 @@ def get_typed_op(op, dtype, dtype2=None, *, is_left_scalar=False, is_right_scala return op[UINT64] if op._udt_types is None: raise KeyError(f"{op.name} does not work with ({dtype}, {dtype2})") - if kind == "binary" and isinstance(op, Monoid) and dtype != dtype2: + if kind == "binary" and isinstance(op, Monoid): # A Monoid takes one type, but used element-wise it is its - # BinaryOp, which takes the pair: ``v.ewise_add(w)`` works where - # ``v + w`` does. Indexing (``monoid.plus[X, Y]``) still raises. + # BinaryOp, which takes the pair: ``v.ewise_mult(w, monoid.plus)`` + # works where ``v * w`` does. Indexing (``monoid.plus[X, Y]``) + # still raises. return get_typed_op(op.binaryop, dtype, dtype2) return op._compile_udt(dtype, dtype2) # Generic case: try to unify the two dtypes @@ -235,9 +236,10 @@ def get_semiring(monoid, binaryop, name=None): monoid_type = monoid.type monoid = monoid.parent if isinstance(binaryop, BinaryOp): - binary_type = None + binary_type = binary_type2 = None else: binary_type = binaryop.type + binary_type2 = getattr(binaryop, "type2", binary_type) binaryop = binaryop.parent if monoid._anonymous or binaryop._anonymous: rv = Semiring.register_anonymous(monoid, binaryop, name=name) @@ -290,6 +292,9 @@ def get_semiring(monoid, binaryop, name=None): setattr(module, funcname, rv) if binary_type is not None: + if binary_type2 is not binary_type and (binary_type._is_udt or binary_type2._is_udt): + # A BinaryOp typed on two different UDTs gives a semiring typed on both. + return get_typed_op(rv, binary_type, binary_type2) return rv[binary_type] if monoid_type is not None: return rv[monoid_type] diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 14690d231..1e0cc194d 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -686,11 +686,14 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): ) op = get_typed_op(op, self.dtype, other.dtype, kind="binary") self._expect_op(op, ("BinaryOp", "Monoid"), within=method_name, argname="op") - if _ewise_add_converts_operands(op, self.dtype, other.dtype): + if _ewise_add_needs_cast(op, self.dtype, other.dtype): + if not _is_lifted_op(op.parent): + raise _ewise_add_cast_error(op, self.dtype, other.dtype) # GraphBLAS casts an unpaired entry to the result type, which a UDT # cannot be. Convert both to the result type first, which changes no # value, and apply the op in that one type: an int8 UDT plus 0.5 is - # float64, and an empty Scalar still adds as 0. + # float64, and an empty Scalar still adds as 0. A Scalar is one + # element, so the converted copies cost nothing. if self.dtype != op.return_type: self = _converted_to(self, op.return_type, name="") if other.dtype != op.return_type: @@ -856,8 +859,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): left_dtype = temp_op.type dtype = left_dtype if left_dtype._is_udt else None if type(left_default) is not Scalar: - if dtype is not None: - _check_literal_fits(dtype, left_default) + _check_literal_fits(dtype, left_default) try: left = Scalar.from_value( left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -877,8 +879,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): right_dtype = temp_op.type2 dtype = right_dtype if right_dtype._is_udt else None if type(right_default) is not Scalar: - if dtype is not None: - _check_literal_fits(dtype, right_default) + _check_literal_fits(dtype, right_default) try: right = Scalar.from_value( right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -1174,58 +1175,200 @@ def _literal_dtype(dtype, value, op): decides how a literal operand is typed beside a UDT; built-in dtypes always get ``None`` and are typed by ``unify``. - Beside a UDT, a tuple, list or array literal becomes an element of that - UDT, as it must for a user's op, but only when no value changes - (:func:`_check_literal_fits`). With a built-in op that lifts to UDTs, a - number is an operand in its own right, typed as numpy 2 types it (NEP 50): + Beside a UDT, a literal becomes an element of that UDT, as it must for a + user's op, when its type allows it (:func:`_check_literal_fits`). With a + built-in op that lifts to UDTs, a literal is an operand in its own right, + typed as numpy 2 types it (NEP 50), by :func:`_literal_type`: - - a numpy scalar or 0-d array keeps its dtype (strong); + - a numpy scalar or array keeps its dtype (strong); an array's type is the + array layout of its dtype (``udt_utils._array_literal_udt``); - a Python ``bool``, ``int``, ``float`` or ``complex`` is weak: in each element or leaf it takes that element's dtype when it fits the kind, so ``int8_udt + 1`` stays ``int8_udt`` and ``int8_udt += 1`` can be stored, while ``int8_udt * 2.5`` has float64 elements (see - ``udt_utils._weak_literal_udt``). Converting every literal to the UDT, - as before, made ``int8_udt * 2.5`` multiply by 2. + ``udt_utils._weak_literal_udt``), where converting it to the UDT would + multiply by 2; + - a tuple, list or dict of Python numbers, laid out like an element, is + weak in the same way, leaf by leaf: ``int8_udt + (1, 2, 3)`` stays + ``int8_udt`` and ``+ (0.5, 1.5, 2.5)`` has float64 elements. ``eq`` and ``ne`` take a number as it is, strong, since a comparison - needs no result type and ``int8_udt == 300`` should be False, not an error. + needs no result type and ``int8_udt == 300`` should be False, not an error; + a sequence likewise, so ``int8_udt == (0.5, 1, 2)`` compares the values. """ if not dtype._is_udt or getattr(op, "is_positional", False): return None - is_python_number = type(value) in {bool, int, float, complex} - is_numpy_number = isinstance(value, (np.bool_, np.number)) or ( - isinstance(value, np.ndarray) and value.ndim == 0 and value.dtype.kind in "biufc" - ) - if is_python_number or is_numpy_number: - from .operator.binary import BinaryOp - from .operator.udt_utils import BUILTIN_UDT_BINARY_OPS, _weak_literal_udt + if getattr(op, "_udt_types", True) is None: + # An op that does not take UDTs at all (``lt``) reports that itself. + return dtype + from .operator.binary import BinaryOp + from .operator.monoid import Monoid + from .operator.udt_utils import _LITERAL_KINDS, BUILTIN_UDT_BINARY_OPS, _as_python_number - if type(op) is BinaryOp and not op._anonymous: - if op.name in ("eq", "ne") or (is_numpy_number and op.name in BUILTIN_UDT_BINARY_OPS): + if isinstance(op, Monoid) and not op._anonymous: + # Element-wise, a Monoid is its BinaryOp (``get_typed_op`` types it so). + op = op.binaryop + if type(op) is BinaryOp and not op._anonymous: + is_comparison = op.name in ("eq", "ne") + if is_comparison or op.name in BUILTIN_UDT_BINARY_OPS: + if is_comparison and type(_as_python_number(value)) in _LITERAL_KINDS: return None - if op.name in BUILTIN_UDT_BINARY_OPS: - return _weak_literal_udt(dtype, value) - if getattr(op, "_udt_types", True) is not None: - # An op that does not take UDTs at all (``lt``) reports that itself. - _check_literal_fits(dtype, value) + literal_type = _literal_type(dtype, value, strong_numbers=is_comparison) + if literal_type is not None: + return literal_type + _check_literal_fits(dtype, value) return dtype -def _check_literal_fits(dtype, value): - """Raise if converting literal ``value`` into UDT ``dtype`` would change one of its values. - - A literal becomes an element of the UDT when nothing else types it: a tuple, - list or array beside a UDT, a literal beside a user-defined op, the - defaults of ``ewise_union``. numpy turns ``0.5`` into ``0`` in an integer - field without a word, and the op would then run on a value the caller never - wrote. A number may round to a narrower float, as numpy 2 rounds a Python - number beside float32, but may not lose an imaginary part or overflow to - infinity. +def _literal_type(dtype, value, *, strong_numbers=False): + """Return the type of literal ``value`` beside UDT ``dtype``, or ``None`` when it has none. + + A Python number (a ``Fraction`` or ``Decimal`` counts as the int or float + it is worth) or a tuple, list or dict of them is weak, leaf by leaf + (``udt_utils._weak_literal_udt``); with ``strong_numbers`` each leaf takes + the number's own dtype instead, as a comparison wants. A numpy value is + strong: its dtype, in the array layout of its shape for an array + (``udt_utils._array_literal_udt``); a sequence holding numpy values is the + array numpy makes of it. Anything else has no type here, and + ``Scalar.from_value`` reports it. """ - import numbers + from .operator.udt_utils import ( + _LITERAL_KINDS, + _array_literal_udt, + _as_python_number, + _weak_literal_udt, + ) + + value = _as_python_number(value) + np_type = dtype.np_type + if isinstance(value, dict) and np_type.names is not None: + value = _dict_to_record(np_type, value) + if type(value) in _LITERAL_KINDS or isinstance(value, (tuple, list)): + literal_type = _weak_literal_udt(dtype, value, weak=not strong_numbers) + if literal_type is not None or type(value) in _LITERAL_KINDS: + return literal_type + try: + value = np.asarray(value) + except (TypeError, ValueError): # ragged + return None + if isinstance(value, (np.bool_, np.number)): + return lookup_dtype(value.dtype) + if isinstance(value, (np.ndarray, np.void)): + if value.ndim == 0 and value.dtype.kind in "biufc": + return lookup_dtype(value.dtype) + return _array_literal_udt(dtype, value) + return None - if not isinstance(value, (numbers.Number, tuple, list, dict, np.generic, np.ndarray)): + +def _check_literal_fits(dtype, value): + """Raise unless literal ``value`` may become an element of UDT ``dtype``, by its type. + + A literal becomes an element of the UDT when nothing else types it: beside + a user-defined op, as an IndexUnaryOp thunk, or as an ``ewise_union`` + default, which GraphBLAS casts to the op's input type, and a UDT casts only + to itself. Its type decides, not its values (:func:`_literal_type`): a + Python number or a sequence of them is weak, so ``1`` and ``(1, 2, 3)`` + beside ``INT8[3]`` are ``INT8[3]`` and fit, ``0.5`` is ``FP64[3]`` and does + not, and ``300`` is ``OverflowError``; a numpy value is strong and fits when + each leaf casts safely into the UDT's (``np.can_cast``), so ``np.int8(1)`` + fits ``INT8[3]`` and ``np.int64(1)`` does not. ``Scalar.from_value`` then + converts it, which changes no value. A literal that does not fit by type but + converts exactly still converts, with a DeprecationWarning + (:func:`_refuse_unless_exact`). A ``dtype`` of ``None``, a built-in type, + takes any literal. + """ + if dtype is None: return + from .operator.udt_utils import _casts_safely_into + + literal_type = _literal_type(dtype, value) + if literal_type is not None and not _casts_safely_into(literal_type, dtype): + _refuse_unless_exact( + dtype, + value, + f"{value!r} does not fit {dtype}: it is typed as {literal_type}, and a literal " + "becomes an element of a UDT only when its type is that UDT or casts into it " + "safely, field by field (a Python number or sequence takes a field's dtype when " + "the field's kind holds it).", + "Write it in the element's types, such as False rather than 0 for a bool field " + "and 2 rather than 2.0 for an integer field.", + ) + return + # A weak float fits a float32 field by type, but a finite one too large for + # it would become infinity; numpy warns as it converts, on numpy 1 and 2. + import warnings + + np_type = dtype.np_type + if np_type.names is not None: + target = np.zeros(1, dtype=np_type) + if isinstance(value, dict): + value = _dict_to_record(np_type, value) + else: + target = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + try: + target[...] = value + except RuntimeWarning: + raise OverflowError( + f"{value!r} overflows {dtype}: a literal that becomes an element of a UDT " + "must not turn a finite number into infinity." + ) from None + except (TypeError, ValueError): + pass # Scalar.from_value reports these. + + +def _check_scalar_fits(dtype, scalar): + """Raise unless a typed Scalar casts safely into UDT ``dtype``, as a literal must. + + An ``ewise_union`` default must have the op's input type, so it is + converted; the same rule as for a literal (:func:`_check_literal_fits`). + """ + from .operator.udt_utils import _casts_safely_into + + if ( + dtype is not None + and scalar.dtype != dtype + and not scalar._is_empty + and not _casts_safely_into(scalar.dtype, dtype) + ): + _refuse_unless_exact( + dtype, + scalar.value, + f"a Scalar of type {scalar.dtype} does not fit {dtype}: it is used where only " + "that UDT can go, and casts into it only when every field casts safely.", + "Convert it first with .dup(dtype=...).", + ) + + +def _refuse_unless_exact(dtype, value, reason, advice): + """Raise ValueError ``reason``, or only warn if ``value`` converts into ``dtype`` exactly. + + MAINT 2026-10-07: a literal or Scalar that must become an element of a UDT + must fit it by type since #591. Before, any value was converted and what did + not fit was lost; one that converts exactly, such as ``0`` for a bool field + or ``2.0`` for an integer field, still converts, with a DeprecationWarning. + Eight months after the release that adds the warning (the deprecation policy + in docs/getting_started/faq.rst), always raise here and delete + ``_converts_exactly``, ``_object_layout`` and ``utils._warn_from_caller``. + """ + if not _converts_exactly(dtype, value): + raise ValueError(reason) + utils._warn_from_caller( + f"{reason} Its values convert exactly, so it is converted for now; this is " + f"deprecated and will raise ValueError in a future version. {advice}", + DeprecationWarning, + ) + + +def _converts_exactly(dtype, value): + """Whether converting ``value`` into UDT ``dtype`` keeps every number in it. + + A float may round to a narrower float, as numpy 2 rounds a Python number + beside float32, but may not become infinite or lose an imaginary part. + """ + import warnings + np_type = dtype.np_type try: if np_type.names is not None: @@ -1238,16 +1381,12 @@ def _check_literal_fits(dtype, value): given = np.empty(shape, dtype=object) converted = np.zeros(shape, dtype=base) given[:] = value - import warnings - with warnings.catch_warnings(): - # numpy warns for some lossy conversions (an out-of-range int on - # numpy 1, a complex array into floats); the comparison below - # reports each of them the same way on every numpy. + # numpy warns for some lossy conversions; the comparison below finds each. warnings.simplefilter("ignore") converted[:] = _view_if_same_layout(value, converted.dtype) - except (TypeError, ValueError): - return # Not something this can read; Scalar.from_value reports it. + except (TypeError, ValueError, OverflowError): + return False for given_leaf, converted_leaf in zip( _leaf_arrays(given), _leaf_arrays(converted), strict=True ): @@ -1257,29 +1396,21 @@ def _check_literal_fits(dtype, value): for x, y in zip(given_leaf.ravel().tolist(), converted_leaf.ravel().tolist(), strict=True): if kind in "fc": z = complex(x) - changed = (kind == "f" and z.imag != 0) or (np.isfinite(z) and not np.isfinite(y)) - else: - changed = x != y - if changed and kind in "iu" and type(x) is int: - # numpy 2 raises this for a Python int itself; numpy 1 wraps. - raise OverflowError( - f"Python integer {x} out of bounds for {converted_leaf.dtype}" - ) - if changed: - raise ValueError( - f"converting {value!r} to {dtype} would change {x!r} to {y!r}; a value is " - "converted into a UDT only when none of it changes." - ) + if (kind == "f" and z.imag != 0) or (np.isfinite(z) and not np.isfinite(y)): + return False + elif x != y: + return False + return True -def _check_scalar_fits(dtype, scalar): - """Raise if a typed Scalar would change a value when converted into UDT ``dtype``. - - An ``ewise_union`` default must have the op's input type, so it is - converted; the same rule as for a literal (:func:`_check_literal_fits`). - """ - if dtype is not None and scalar.dtype != dtype and not scalar._is_empty: - _check_literal_fits(dtype, scalar.value) +def _object_layout(np_type): + """Return ``np_type`` with every leaf an object, keeping field names, nesting and shapes.""" + if np_type.names is not None: + return np.dtype([(name, _object_layout(np_type.fields[name][0])) for name in np_type.names]) + if np_type.subdtype is not None: + base, shape = np_type.subdtype + return np.dtype((_object_layout(base), shape)) + return np.dtype(object) def _is_lifted_op(op): @@ -1293,16 +1424,6 @@ def _is_lifted_op(op): return type(op) is BinaryOp and not op._anonymous and op.name in BUILTIN_UDT_BINARY_OPS -def _object_layout(np_type): - """Return ``np_type`` with every leaf an object, keeping field names, nesting and shapes.""" - if np_type.names is not None: - return np.dtype([(name, _object_layout(np_type.fields[name][0])) for name in np_type.names]) - if np_type.subdtype is not None: - base, shape = np_type.subdtype - return np.dtype((_object_layout(base), shape)) - return np.dtype(object) - - def _leaf_arrays(arr): """Yield the array of each leaf of a structured array, in field order.""" if arr.dtype.names is None: @@ -1312,49 +1433,48 @@ def _leaf_arrays(arr): yield from _leaf_arrays(arr[name]) -def _ewise_add_converts_operands(op, left_dtype, right_dtype): - """Whether ``ewise_add`` must convert its operands to ``op``'s result type before the op. +def _ewise_add_needs_cast(op, left_dtype, right_dtype): + """Whether ``ewise_add`` with ``op`` would have GraphBLAS cast to or from a UDT. eWiseAdd copies an entry present in only one input into the result, cast to the op's output type, and GraphBLAS checks that cast before it looks at the data. A UDT casts only to itself, so such a pair is a bare - GrB_DOMAIN_MISMATCH even when every entry has a partner. Each lifted - arithmetic op gives the same values on operands already converted to its - result type, since that type is what it promotes their elements to, so - those convert first (``_ewise_add_in_result_type``). Any other op raises - ``DomainMismatch`` here. + GrB_DOMAIN_MISMATCH even when every entry has a partner. """ out = op.return_type if not (out._is_udt or left_dtype._is_udt or right_dtype._is_udt): return False - if left_dtype == out and right_dtype == out: - return False - from .operator.udt_utils import BUILTIN_UDT_BINARY_OPS + return left_dtype != out or right_dtype != out + - parent = op.parent - if op.opclass == "BinaryOp" and not parent._anonymous and parent.name in BUILTIN_UDT_BINARY_OPS: - return True +def _ewise_add_cast_error(op, left_dtype, right_dtype): + """Return the ``DomainMismatch`` for an ``ewise_add`` that would need a UDT cast.""" from ..exceptions import DomainMismatch - raise DomainMismatch( - f"ewise_add cannot use {parent!r} on {left_dtype} and {right_dtype}: an entry present " - f"in only one input is cast to the result type {out}, and UDTs do not cast. Use " - "ewise_union, which takes a default for each side, or ewise_mult." + return DomainMismatch( + f"ewise_add cannot use {op.parent!r} on {left_dtype} and {right_dtype}: an entry " + f"present in only one input is copied into the result as {op.return_type}, and " + "GraphBLAS cannot cast a UDT, so this would need converted copies of the operands. " + "Use ewise_union, which takes a default for each side (infix + does, with zeros), or " + "ewise_mult; or convert an operand first with .dup(dtype=...)." ) -def _ewise_add_in_result_type(updater, left, right, op): - """Recipe for ``ewise_add`` when ``_ewise_add_converts_operands`` says to convert. +def _plus_zero(dtype): + """Return a Scalar of ``dtype`` that ``plus`` adds to any value without changing it. - Each operand is converted to the result type (:func:`_converted_to`), and - the op then runs on the two operands in that one type. + Its float and complex leaves are ``-0.0``, because ``-0.0 + x`` is ``x`` + for every float, where ``0.0 + -0.0`` is ``0.0``. """ - out = op.return_type - if left.dtype != out: - left = _converted_to(left, out, name="left_temp", **updater.opts) - if right.dtype != out: - right = _converted_to(right, out, name="right_temp", **updater.opts) - updater << left.ewise_add(right, op.parent) + np_type = dtype.np_type + if np_type.subdtype is not None: + zero = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) + else: + zero = np.zeros(1, dtype=np_type) + for leaf in _leaf_arrays(zero): + if leaf.dtype.kind in "fc": + leaf[...] = complex(-0.0, -0.0) if leaf.dtype.kind == "c" else -0.0 + return Scalar.from_value(zero if np_type.subdtype is not None else zero[0], dtype, name="") def _converted_to(operand, dtype, **new_kwargs): @@ -1364,19 +1484,8 @@ def _converted_to(operand, dtype, **new_kwargs): zero on the left, the result is ``dtype`` itself even when the operand has the same elements in another layout (a packed and an aligned record), since a left operand that holds the promoted elements is the result type. - The float and complex leaves of that zero are ``-0.0``, because - ``-0.0 + x`` is ``x`` for every float, where ``0.0 + -0.0`` is ``0.0``. """ - np_type = dtype.np_type - if np_type.names is not None: - zero = np.zeros(1, dtype=np_type) - else: - zero = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) - for leaf in _leaf_arrays(zero): - if leaf.dtype.kind in "fc": - leaf[...] = complex(-0.0, -0.0) if leaf.dtype.kind == "c" else -0.0 - zero = Scalar.from_value(zero[0] if np_type.names is not None else zero, dtype, name="") - rv = operand.apply(binary.plus, left=zero).new(**new_kwargs) + rv = operand.apply(binary.plus, left=_plus_zero(dtype)).new(**new_kwargs) if rv.dtype != dtype: # pragma: no cover (safety) from ..exceptions import DomainMismatch @@ -1385,19 +1494,27 @@ def _converted_to(operand, dtype, **new_kwargs): def _cast_for_store(value, dtype): - """Return a Scalar, Vector or Matrix ``value`` as ``dtype``, to be assigned into an object. + """Return Scalar, Vector or Matrix ``value``, or its expression, to assign as ``dtype``. GraphBLAS's assign and setElement take a UDT value as raw bytes, so one of another type would be read as the wrong type; ``base._store_cast_op`` - casts it first or raises. + casts it or raises. Assign cannot cast as it writes, so a value of + another UDT type needs a converted copy: one element for a Scalar, which + is made here, but the whole value for a Vector or Matrix, which raises. """ from .base import _store_cast_op if (cast_op := _store_cast_op(value.dtype, dtype)) is None: return value + if output_type(value) is not Scalar: + from ..exceptions import DomainMismatch + + raise DomainMismatch( + f"cannot assign a value of type {value.dtype} into an object of type {dtype}: " + "assign cannot cast a UDT as it writes, so this would need a converted copy of " + "the whole value. Make it explicitly with .dup(dtype=...) first." + ) if not isinstance(value, BaseType): - # An expression, such as ``x[[0, 1]] << y + 0.5``: compute it in its own - # type, then cast it. value = value.new(name="") return value.apply(cast_op).new(name=value.name) diff --git a/graphblas/core/utils.py b/graphblas/core/utils.py index 72f7219bd..86ec61a8a 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -49,6 +49,30 @@ def inner(func_wo_doc): _output_types.update((k, k) for k in set(np.sctypeDict.values())) +def _warn_from_caller(message, category): + """Warn at the first caller outside python-graphblas, so default filters can show it. + + Python shows a DeprecationWarning by default only when it is attributed to + ``__main__``, and a check that warns may sit at any depth below the public + method. python-graphblas's own tests count as callers. + """ + import sys + import warnings + from pathlib import Path + + package = Path(__file__).parent.parent + tests = package / "tests" + frame = sys._getframe(1) + stacklevel = 2 + while frame is not None: + filename = Path(frame.f_code.co_filename) + if not filename.is_relative_to(package) or filename.is_relative_to(tests): + break + frame = frame.f_back + stacklevel += 1 + warnings.warn(message, category, stacklevel=stacklevel) + + def output_type(val): try: return _output_types[type(val)] diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index ba48dce80..6c377e535 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -28,8 +28,8 @@ _cast_for_store, _check_literal_fits, _check_scalar_fits, - _ewise_add_converts_operands, - _ewise_add_in_result_type, + _ewise_add_cast_error, + _ewise_add_needs_cast, _literal_dtype, _scalar_index, ) @@ -599,7 +599,11 @@ def dup(self, dtype=None, *, clear=False, mask=None, name=None, **opts): dtype = self.dtype rv = Vector(dtype, size=self._size, name=name) if not clear: - rv(mask=mask, **opts)[...] = self + if rv.dtype._is_udt or self.dtype._is_udt: + # ``<<``, not assign: a UDT of another type casts in one pass. + rv(mask=mask, **opts) << self + else: + rv(mask=mask, **opts)[...] = self else: if opts: # Ignore opts for now @@ -1040,20 +1044,13 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): f"to columns of Matrix in {method_name}. Matrix.nrows (={other._nrows}) " f"must equal Vector.size (={self._size})." ) - if _ewise_add_converts_operands(op, self.dtype, other.dtype): - # GraphBLAS casts an unpaired entry to the result type, which a UDT - # cannot be. Convert both operands to the result type first and - # apply the op in that one type: an int8 UDT plus a float32 one is - # float32, with every entry of either side. - args = [self, other, _ewise_add_in_result_type, (self, other, op)] - if other.ndim == 2: - return MatrixExpression( - method_name, None, args, nrows=other._nrows, ncols=other._ncols, op=op + if _ewise_add_needs_cast(op, self.dtype, other.dtype): + # GraphBLAS would report the UDT cast before the sizes. + if other.ndim == 1 and self._size != other._size: + raise DimensionMismatch( + f"Vector sizes must match in {method_name}: {self._size} and {other._size}" ) - expr = VectorExpression(method_name, None, args, op=op) - if self._size != other._size: - expr.new(name="") # incompatible shape; raise now - return expr + raise _ewise_add_cast_error(op, self.dtype, other.dtype) if other.ndim == 2: return MatrixExpression( @@ -1228,8 +1225,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): left_dtype = temp_op.type dtype = left_dtype if left_dtype._is_udt else None if type(left_default) is not Scalar: - if dtype is not None: - _check_literal_fits(dtype, left_default) + _check_literal_fits(dtype, left_default) try: left = Scalar.from_value( left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -1249,8 +1245,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): right_dtype = temp_op.type2 dtype = right_dtype if right_dtype._is_udt else None if type(right_default) is not Scalar: - if dtype is not None: - _check_literal_fits(dtype, right_default) + _check_literal_fits(dtype, right_default) try: right = Scalar.from_value( right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -1626,7 +1621,8 @@ def select(self, op, thunk=None): if thunk is None: thunk = False # most basic form of 0 when unifying dtypes if type(thunk) is not Scalar: - dtype = self.dtype if (self.dtype._is_udt and not op.is_positional) else None + # A thunk beside a UDT becomes an element of it, so it must fit, as in apply. + dtype = _literal_dtype(self.dtype, thunk, op) try: thunk = Scalar.from_value(thunk, dtype, is_cscalar=None, name="") except TypeError: @@ -1875,6 +1871,15 @@ def _extract_element( dtype = self.dtype else: dtype = lookup_dtype(dtype) + if dtype is not self.dtype and (dtype._is_udt or self.dtype._is_udt): + # GraphBLAS copies a UDT element as raw bytes of this vector's type, + # past the end of a smaller one, so extract it in this type and cast + # it as a store does (or raise). + value = self._extract_element(resolved_indexes, None, opts, is_cscalar=is_cscalar) + if result is None: + result = Scalar(dtype, is_cscalar=is_cscalar, name=name) + result << value + return result idx = resolved_indexes.indices[0] if result is None: result = Scalar(dtype, is_cscalar=is_cscalar, name=name) @@ -1905,8 +1910,6 @@ def _prep_for_extract(self, resolved_indexes, mask=None, is_submask=False): def _assign_element(self, resolved_indexes, value): idx = resolved_indexes.indices[0] - if type(value) is Scalar: - value = _cast_for_store(value, self.dtype) if type(value) is not Scalar: dtype = self.dtype if self.dtype._is_udt else None try: @@ -1919,6 +1922,10 @@ def _assign_element(self, resolved_indexes, value): argname="value", extra_message="Literal scalars also accepted.", ) + if value.dtype._is_udt or self.dtype._is_udt: + # After the conversion: an autocomputed element of another UDT type + # arrives here as a Scalar too. + value = _cast_for_store(value, self.dtype) if value._is_cscalar: if value._empty: call("GrB_Vector_removeElement", [self, idx.index]) diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index ea8a29f26..fcb2be985 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -3696,11 +3696,12 @@ def test_udt(): assert s == t t = A.reduce_scalar(monoid.any, allow_empty=False).new() assert s == t - A << binary.first(1, A) + # A literal fills every field by type; an int does not fit the bool field, True does. + A << binary.first(True, A) assert A.isequal(ones) A << binary.first(s, A) assert A.isequal(zeros) - A << binary.second(A, 1) + A << binary.second(A, True) assert A.isequal(ones) A << binary.second(A, s) assert A.isequal(zeros) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index b6a5f5424..566ce6e87 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1266,8 +1266,8 @@ def _udt_index(val, idx, _, thunk): # pragma: no cover (numba) result = v.apply(_udt_index, (False, 3)).new() expected = Vector.from_coo([0, 1, 2], [3, -3, -3]) assert result.isequal(expected) - # A number fills every field, and 3 in the bool field would be True. - with pytest.raises(ValueError, match="would change 3 to True"): + # A number fills every field, and an int is not a bool: 3 does not fit. + with pytest.raises(ValueError, match="3 does not fit"): v.apply(_udt_index, 3) def _udt_first(x, y): @@ -1277,7 +1277,11 @@ def _udt_first(x, y): assert udt in udt_first assert operator.get_typed_op(udt_first, udt) is udt_first[udt] assert udt_first(v & w).new().isequal(v) - assert udt_first(v, 1).new().isequal(v) + # A literal fills every field, by type: an int does not fit the bool field, + # True does. 1 still converts, being exact, but that is deprecated. + assert udt_first(v, True).new().isequal(v) + with pytest.warns(DeprecationWarning, match="1 does not fit TestUDT"): + assert udt_first(v, 1).new().isequal(v) assert udt_first[udt, dtypes.INT64].return_type == udt assert udt_first[dtypes.INT64, udt].return_type == dtypes.INT64 assert udt_first[udt, dtypes.BOOL].return_type == udt @@ -1354,10 +1358,11 @@ def _this_or_that(val, idx, _, thunk): # pragma: no cover (numba) sel = SelectOp.register_anonymous(_this_or_that, is_udt=True) sel[udt] assert udt in sel - result = v.select(sel, 0).new() + # The thunk becomes an element of the record, whose bool field takes False, not 0. + result = v.select(sel, False).new() assert result.nvals == 0 assert result.dtype == v.dtype - result = w.select(sel, 0).new() + result = w.select(sel, False).new() assert result.nvals == 3 assert result.isequal(w) @@ -1429,6 +1434,435 @@ def _ne_thunk(x, i, j, thunk): # pragma: no cover (numba) assert v.select(sel, 5).new().isequal(Vector.from_coo([0, 2], [1, 9], size=3)) +_jit_can_compile_cache = [] + + +def _jit_can_compile(): + """True when SuiteSparse has a C compiler it can actually use. + + Without one it falls back to the Numba cfunc and says nothing, so the + ``jit`` parameter would run the cfunc and report a pass. + + ``jit_compiler_is_usable`` alone is not enough: it only checks that the + configured compiler path exists on disk, and a runner can have the file + yet fail every compile (broken toolchain, missing headers). + ``_enable_jit_for_udt`` is the same call the UDT auto-lift path makes, + and it does a real compile once per process, so its answer is the + honest one here; ``test_ssjit`` keys its skips on the same signal. The + fixture calls this before it mutates the control, and the cache keeps + later per-test mutations from flipping it. + """ + if not _jit_can_compile_cache: + from graphblas.core.ss.jit_config import _enable_jit_for_udt + + _jit_can_compile_cache.append(gb.ss.jit_compiler_is_usable() and _enable_jit_for_udt()) + return _jit_can_compile_cache[0] + + +@pytest.fixture(params=["jit", "cfunc"]) +def udt_op_path(request): + """Pin SuiteSparse to one execution path for built-in UDT operators. + + Each auto-lifted UDT op carries both a C JIT definition and a Numba + cfunc, and SuiteSparse chooses between them per call depending on + whether a C compiler is available. A machine with one and a machine + without therefore run different code, so results have to hold on both. + """ + path = request.param + if backend != "suitesparse" or "jit_c_control" not in gb.ss.config: + if path == "jit": + pytest.skip("no SuiteSparse C JIT on this backend") + yield path + return + previous = gb.ss.config["jit_c_control"] + if path == "jit": + if not _jit_can_compile(): + pytest.skip("C JIT compilation not available (probe failed or compiler missing)") + # Set it rather than assume it. SuiteSparse demotes ``on`` to ``load`` + # after a failed compile, and a demoted control routes to the cfunc + # silently, so this parameter would pass while running the other path. + gb.ss.config["jit_c_control"] = "on" + else: + gb.ss.config["jit_c_control"] = "off" + try: + yield path + finally: + # Read before restoring: a demotion during the test is the signal that + # the kernel never compiled, which no assertion in the test can see. + demoted = path == "jit" and gb.ss.config["jit_c_control"] != "on" + gb.ss.config["jit_c_control"] = previous + if demoted: + pytest.fail("SuiteSparse demoted jit_c_control; the JIT path did not run") + + +def _udt_vectors(udt, xs, ys=None): + """Build one or two dense UDT vectors whose leaves all hold the given values. + + Values that repeat down the whole vector make it iso-valued, and + SuiteSparse answers those from a single element without reaching for a + C JIT kernel, so callers pass varied data. + """ + names = udt.np_type.names + out = [] + for vals in (xs, ys): + if vals is None: + continue + v = Vector(udt, size=len(vals)) + for i, val in enumerate(vals): + v[i] = tuple(val for _ in names) if names else np.full(udt.np_type.subdtype[1], val) + out.append(v) + return out + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_mixed_record_dtypes_use_each_operands_own_dtype(udt_op_path): + """Two records sharing field names but not field types promote to the wider one. + + ``_check_udt_pair`` matches record operands on field names only, so their + leaf dtypes can differ. Offering the return-type resolver just the left + operand left it no choice but that record, so an int record over a float + record came back as the int record: ``7 / 2.0`` landed as 3 and ``6 / 0.0`` + as INT64_MAX. Swapping the operands changed the answer for the same pair. + """ + int_udt = dtypes.register_anonymous(np.dtype([("mxd_a", np.int64)], align=True), "_MixedRecInt") + float_udt = dtypes.register_anonymous( + np.dtype([("mxd_a", np.float64)], align=True), "_MixedRecFloat" + ) + v = Vector(int_udt, size=2) + v[0] = (6,) + v[1] = (7,) + w = Vector(float_udt, size=2) + w[0] = (0.0,) + w[1] = (2.0,) + + result = v.ewise_mult(w, binary.truediv).new() + assert result.dtype == float_udt, "result should promote to the float record" + assert result[0].new().value["mxd_a"] == float("inf") + assert result[1].new().value["mxd_a"] == 3.5 + + # The same pair the other way round must agree, which it did not when the + # resolver only ever saw the left operand. + swapped = w.ewise_mult(v, binary.truediv).new() + assert swapped.dtype == float_udt + assert swapped[1].new().value["mxd_a"] == 2.0 / 7.0 + + +def _bitwise_eq(got, want): + """Compare two floats by bit pattern, treating any two NaNs as equal. + + Bit patterns rather than ``==`` because ``-0.0 == 0.0``, and the sign of + a zero is exactly what a min/max tie-break decides. NaNs are exempted + because ``fmin`` may hand back either operand's NaN payload. + """ + if np.isnan(got) and np.isnan(want): + return True + return got.tobytes() == want.tobytes() + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +@pytest.mark.parametrize("np_dtype", [np.float64, np.float32]) +def test_udt_min_max_answer_what_the_builtin_dtype_answers(udt_op_path, np_dtype): + """``binary.min[udt]`` must give what ``binary.min[FP64]`` gives, bit for bit. + + An operator that means one thing on FP64 and another on a record of + FP64 is not one operator. SuiteSparse's ``GrB_MIN_FP64`` is C99 ``fmin``, + so that is what the UDT kernels have to be, and this compares them + directly against the built-in rather than against a convention chosen on + the Python side. The grid is every ordered pair drawn from NaN, both + infinities, both zeros and two ordinary values, so it covers a NaN on + either side, two NaNs, and a signed-zero tie either way round. + + The signed-zero tie itself is compared by value only. C99 leaves + ``fmin(-0.0, 0.0)`` unspecified and the built-in answers differently per + platform (left operand on macOS x86, right operand on Linux x86, IEEE + minNum on arm64 and Windows), so bit-for-bit agreement on that one pair + is not something any implementation can promise. Everything else, + including which zero a mixed zero/nonzero pair keeps, stays bit-exact. + + What this catches, in the two spellings it replaces: Python's builtin + ``min``, which the generated code reached through the exec namespace, + ordered NaN by position, and ``np.fmin`` under Numba gets the NaN rule + right but keeps the left operand on a signed-zero tie, so it drifts from + the C JIT kernel on ``min(0.0, -0.0)``. Both execution paths are checked + because SuiteSparse picks between them without telling anyone. + """ + nan, inf = float("nan"), float("inf") + values = [nan, inf, -inf, -0.0, 0.0, 1.5, -2.5] + xs = [x for x in values for _ in values] + ys = list(values) * len(values) + + udt = dtypes.register_anonymous( + np.dtype([("mmb_a", np_dtype)], align=True), f"_MinMaxBuiltin{np.dtype(np_dtype).name}" + ) + v, w = _udt_vectors(udt, xs, ys) + ref_v = Vector.from_dense(np.array(xs, dtype=np_dtype)) + ref_w = Vector.from_dense(np.array(ys, dtype=np_dtype)) + + for gb_op in (binary.min, binary.max): + expected = gb_op(ref_v & ref_w).new().to_dense() + result = gb_op(v & w).new() + for i, (x, y) in enumerate(zip(xs, ys, strict=True)): + got = result[i].new().value[0] + if x == 0 and y == 0 and np.signbit(x) != np.signbit(y): + # The one unspecified cell of the grid: either signed zero is + # a correct answer from either implementation, so only agree + # that both produced a zero. + msg = ( + f"{udt_op_path} {gb_op.name}({x}, {y}) on {udt.name}: " + f"got {got!r}, built-in {np.dtype(np_dtype).name} gives {expected[i]!r}" + ) + assert got == 0, msg + assert expected[i] == 0, msg + continue + assert _bitwise_eq(got, expected[i]), ( + f"{udt_op_path} {gb_op.name}({x}, {y}) on {udt.name}: " + f"got {got!r}, built-in {np.dtype(np_dtype).name} gives {expected[i]!r}" + ) + + # A NaN anywhere in the input must not change where a reduce lands. Under + # the Python-builtin semantics this same multiset reduced to 1.0 or to nan + # depending on which index the NaN sat at. + for data in ([1.0, 2.0, 3.0, nan], [nan, 1.0, 2.0, 3.0], [1.0, nan, 3.0, 2.0]): + (u,) = _udt_vectors(udt, data) + assert u.reduce(monoid.min[udt]).new().value[0] == 1.0, f"{udt_op_path} {data}" + assert u.reduce(monoid.max[udt]).new().value[0] == 3.0, f"{udt_op_path} {data}" + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_truediv_divides_in_floating_point(udt_op_path): + """``binary.truediv`` on integer fields divides in float64 and keeps the quotient. + + Regression: the C JIT kernel emitted C ``/``, which is integer division. + ``10**18 / 3`` came out as 333333333333333333 under the C JIT and + 333333333333333312 (float64, like numpy) through the cfunc, so the same + program gave different answers depending on whether a C compiler was + installed. The result fields are now float64, as ``truediv`` on two INT64 + vectors is FP64, so the quotient is not truncated either. + """ + udt = dtypes.register_anonymous( + np.dtype([("tdv_i", np.int64), ("tdv_j", np.int64)], align=True), "_TrueDivIntUDT" + ) + xs = [10**18, 10**18 + 1, 7, 22] + ys = [3, 3, 2, 7] + expected = np.array(xs, np.int64) / np.array(ys, np.int64) + v, w = _udt_vectors(udt, xs, ys) + result = binary.truediv(v & w).new() + assert [result.dtype.np_type[name] for name in ("tdv_i", "tdv_j")] == [np.dtype(np.float64)] * 2 + got = [result[i].new().value[0] for i in range(len(xs))] + assert got == list(expected), udt_op_path + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_floordiv_matches_numpy_on_floats(udt_op_path): + """``binary.floordiv`` on float fields is not ``floor(a / b)``. + + Regression: the C JIT kernel computed ``floor(a / b)``, which rounds + differently from the remainder-based algorithm numpy and CPython use and + treats infinities as ordinary values. ``1.0 // 0.1`` came out as 10.0 + instead of 9.0, and ``inf // 2.0`` as ``inf`` instead of NaN, while the + cfunc agreed with numpy all along. + """ + nan = float("nan") + inf = float("inf") + # ``floor(a / b)`` disagrees with numpy on the first four pairs: inf and + # -inf where numpy gives NaN, 10.0 rather than 9.0 for 1.0 // 0.1, and + # -0.0 rather than -1.0 for -2.0 // inf. + xs = [inf, -inf, 1.0, -2.0, 2.0, -2.0, 0.0, nan, -7.0, 7.0, -0.0, 7.5] + ys = [2.0, 2.0, 0.1, inf, 0.0, 0.0, 0.0, 2.0, 2.0, -2.0, 4.0, 2.5] + for np_dtype, name in ((np.float64, "_FloorDivF64UDT"), (np.float32, "_FloorDivF32UDT")): + udt = dtypes.register_anonymous( + np.dtype([("fdv_a", np_dtype), ("fdv_b", np_dtype)], align=True), name + ) + with np.errstate(divide="ignore", invalid="ignore"): + expected = np.floor_divide(np.array(xs, np_dtype), np.array(ys, np_dtype)) + v, w = _udt_vectors(udt, xs, ys) + result = binary.floordiv(v & w).new() + got = np.array([result[i].new().value[0] for i in range(len(xs))], np_dtype) + np.testing.assert_array_equal(got, expected, err_msg=f"{udt_op_path} {np_dtype.__name__}") + # ``assert_array_equal`` reads -0.0 and 0.0 as equal, so the sign of a + # zero quotient needs its own assertion. It is the whole job of the + # ``copysign`` branch the C JIT kernel emits for an exact-zero result. + np.testing.assert_array_equal( + np.signbit(got), + np.signbit(expected), + err_msg=f"{udt_op_path} {np_dtype.__name__} sign of zero", + ) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_integer_division_by_zero_is_defined(udt_op_path): + """Integer division by zero must return a value rather than trap. + + The C JIT kernel divided in integers, where a zero divisor is undefined + behaviour: on x86-64 ``idiv`` raises #DE, which is SIGFPE and process + death rather than an exception. AArch64's ``sdiv`` returns 0 and does not + trap, so this cannot be exhibited on an arm64 machine. The same trap + fires on ``INT_MIN / -1``, whose quotient is not representable. + + ``truediv`` now gives float64 fields, as on built-in integer vectors, so a + zero divisor gives an infinity there. ``floordiv`` stays in integers, where + the values are a choice: a zero divisor gives 0, as ``np.floor_divide`` + does, and ``INT_MIN // -1`` wraps to ``INT_MIN``, as numpy does. + """ + signed = dtypes.register_anonymous( + np.dtype([("dvz_a", np.int32), ("dvz_b", np.int8)], align=True), "_DivZeroSignedUDT" + ) + unsigned = dtypes.register_anonymous( + np.dtype([("dvz_c", np.uint32), ("dvz_d", np.uint64)], align=True), "_DivZeroUnsignedUDT" + ) + inf = float("inf") + v, w = _udt_vectors(signed, [7, -7, 100, -128], [0, 0, 3, -1]) + result = binary.truediv(v & w).new() + got = [result[i].new().value[0] for i in range(4)] + assert got == [inf, -inf, 100 / 3, 128.0], udt_op_path + result = binary.floordiv(v & w).new() + got = [result[i].new().value[0] for i in range(4)] + assert got[:3] == [0, 0, 33], udt_op_path + # ``-128 // -1`` is the second trapping case; numpy wraps it to INT8_MIN. + result = binary.floordiv(v & w).new() + assert result[3].new().value[1] == np.iinfo(np.int8).min, udt_op_path + + v, w = _udt_vectors(unsigned, [7, 9, 100, 5], [0, 0, 3, 2]) + result = binary.truediv(v & w).new() + got = [result[i].new().value[0] for i in range(4)] + assert got == [inf, inf, 100 / 3, 2.5], udt_op_path + result = binary.floordiv(v & w).new() + got = [result[i].new().value[0] for i in range(4)] + assert got == [0, 0, 33, 2], udt_op_path + + # Floor division still floors for signed operands of mixed sign. + v, w = _udt_vectors(signed, [-7, 7, -9, 11], [2, -2, 2, 3]) + result = binary.floordiv(v & w).new() + got = [result[i].new().value[0] for i in range(4)] + assert got == [-4, -4, -5, 3], udt_op_path + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("not dtypes._supports_complex") +def test_udt_complex_truediv_by_zero(udt_op_path): + """``binary.truediv`` on a complex field survives a zero divisor. + + Numba's complex division raises ``ZeroDivisionError`` unconditionally, + outside the error model's control, so the cfunc left the element unwritten + while the C JIT kernel returned numpy's infinities. Reading back the + abandoned element gave uninitialized memory, or the previous element's + answer, either of which looks like a plausible value. + """ + udt = dtypes.register_anonymous( + np.dtype([("cxz_a", np.complex128)], align=True), "_ComplexDivZeroUDT" + ) + xs = [3 + 4j, 0j, 1 + 1j, 2 - 2j] + ys = [0j, 0j, 2 + 0j, 0j] + v, w = _udt_vectors(udt, xs, ys) + result = binary.truediv(v & w).new() + got = np.array([result[i].new().value[0] for i in range(len(xs))]) + with np.errstate(divide="ignore", invalid="ignore"): + expected = np.array(xs) / np.array(ys) + np.testing.assert_array_equal(got, expected, err_msg=udt_op_path) + + +@pytest.mark.skipif("not supports_udfs") +# 136-byte UDT, which SS < 9 rejects; see test_udt_large_array. +@pytest.mark.skipif( + "ss_version_major < 9", + reason="SuiteSparse < 9 rejects a 136-byte UDT on builds without VLA support", +) +def test_udt_float_truediv_by_zero_is_infinite(udt_op_path): + """A zero divisor on a float field gives numpy's infinity, not a lost element. + + Unlike the integer case, nothing here is guarded: the generated code + divides and lets IEEE produce the infinity. That only holds because the + generated wrapper is compiled under Numba's numpy error model, which + nothing else in the suite pins down. + """ + udt = dtypes.register_anonymous(np.dtype((np.float64, (17,))), "_FloatDivZeroArr17") + xs = [1.0, -1.0, 0.0, 6.0] + ys = [0.0, 0.0, 0.0, 3.0] + v, w = _udt_vectors(udt, xs, ys) + result = binary.truediv(v & w).new() + got = np.array([result[i].new().value[0] for i in range(len(xs))]) + with np.errstate(divide="ignore", invalid="ignore"): + expected = np.array(xs) / np.array(ys) + np.testing.assert_array_equal(got, expected, err_msg=udt_op_path) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_ops_match_record_ops(udt_op_path): + """The array-UDT codegen carries the same division and NaN fixes as records. + + Records and flat arrays go through separate branches in both the Numba + and the C JIT generators, so each fix has to land in both. + """ + nan = float("nan") + inf = float("inf") + float_udt = dtypes.register_anonymous(np.dtype((np.float64, (13,))), "_ArrOpsF64") + xs = [inf, 1.0, -2.0, nan, -7.0, 2.0] + ys = [2.0, 0.1, inf, 2.0, 2.0, 1.0] + v, w = _udt_vectors(float_udt, xs, ys) + # The reference computations touch inf and nan, and numpy raises the FP + # invalid flag for them on some platforms (Linux and Windows, via fmod) + # but not others; pyproject promotes the RuntimeWarning to an error. + with np.errstate(divide="ignore", invalid="ignore"): + expected_floordiv = np.floor_divide(np.array(xs), np.array(ys)) + expected_min = np.fmin(np.array(xs), np.array(ys)) + np.testing.assert_array_equal( + [binary.floordiv(v & w).new()[i].new().value[0] for i in range(len(xs))], + expected_floordiv, + err_msg=udt_op_path, + ) + # ``np.fmin``, not ``np.minimum``: ``binary.min`` is SuiteSparse's + # ``GrB_MIN_FP64``, which ignores a NaN operand rather than propagating it. + np.testing.assert_array_equal( + [binary.min(v & w).new()[i].new().value[0] for i in range(len(xs))], + expected_min, + err_msg=udt_op_path, + ) + + int_udt = dtypes.register_anonymous(np.dtype((np.int64, (6,))), "_ArrOpsI64") + ixs = [10**18, 7, -7, 100, -9, 5] + iys = [3, 0, 0, 3, 2, 2] + v, w = _udt_vectors(int_udt, ixs, iys) + result = binary.truediv(v & w).new() + got = [result[i].new().value[0] for i in range(len(ixs))] + with np.errstate(divide="ignore"): + assert got == list(np.array(ixs) / np.array(iys)), udt_op_path + result = binary.floordiv(v & w).new() + got = [result[i].new().value[0] for i in range(len(ixs))] + assert got == [333333333333333333, 0, 0, 33, -5, 2], udt_op_path + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_multidim_array_ops_match_numpy(udt_op_path): + """Built-in ops on a multi-dimensional array UDT agree with numpy on both paths. + + The C JIT typedef flattens any rank to ``double v [N]`` and the Numba + wrapper walks the same flat run, so a 2-D UDT covers codegen that the 1-D + cases reach only by accident of both being contiguous. + """ + udt = dtypes.register_anonymous(np.dtype((np.float64, (3, 2))), "_ArrOps2D") + xs = [1.0, -7.0, float("inf"), 2.0] + ys = [0.1, 2.0, 2.0, 0.0] + v, w = _udt_vectors(udt, xs, ys) + for gb_op, reference in ( + (binary.floordiv, np.floor_divide), + (binary.truediv, np.true_divide), + # ``fmin`` rather than ``minimum``: these inputs carry no NaN, so the + # two agree here, but ``binary.min`` is the NaN-ignoring one. + (binary.min, np.fmin), + ): + result = gb_op(v & w).new() + element = result[0].new().value + assert element.shape == (3, 2) + got = np.array([result[i].new().value[1, 1] for i in range(len(xs))]) + with np.errstate(divide="ignore", invalid="ignore"): + expected = reference(np.array(xs), np.array(ys)) + np.testing.assert_array_equal(got, expected, err_msg=f"{udt_op_path} {gb_op.name}") + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_tuple_return_binaryop(record_udt): @@ -3510,9 +3944,12 @@ def test_udt_array_pairs_broadcast_as_numpy(x_type, y_type): np.testing.assert_array_equal(value, expected) assert not binary.eq(a & b).new()[0].new().value assert binary.ne(a & b).new()[0].new().value - # ewise_add converts an entry without a partner to the result type, so - # it broadcasts to the result's shape. - added = a.ewise_add(b, binary.minus).new() + # ewise_add would copy an entry without a partner as the result type, a + # cast GraphBLAS cannot make for a UDT; ewise_union computes it instead, + # so it broadcasts to the result's shape. + with pytest.raises(DomainMismatch, match="ewise_add cannot use binary.minus"): + a.ewise_add(b, binary.minus) + added = a.ewise_union(b, binary.minus, 0, 0).new() assert added.dtype == result_type np.testing.assert_array_equal(added[0].new().value, expected) lone = 1 if a is v else 2 @@ -3636,10 +4073,13 @@ def test_udt_array_pair_promotes_the_base(): assert binary.ne(a & b).new().isequal(Vector.from_coo([0, 1], [True, False])) # A semiring's monoid is typed by its multiplier's result. assert v.inner(w, semiring.plus_times).new().dtype == floats - # GraphBLAS would cast an entry of ``v`` alone to the result type, which - # UDTs do not do; ``plus`` converts both sides first instead. - assert v.ewise_add(w, binary.plus).new().isequal((v + w).new()) + # GraphBLAS would cast an entry of ``v`` alone to the result type, which a + # UDT cannot be; ewise_union computes that entry with the op instead, and + # infix ``+`` uses it, with 0 for the missing value. + with pytest.raises(DomainMismatch, match="ewise_add cannot use binary.plus"): + v.ewise_add(w) assert v.ewise_union(w, binary.plus, 0, 0.0).new().dtype == floats + assert (v + w).new().isequal(v.ewise_union(w, binary.plus, 0, 0.0).new()) @pytest.mark.skipif("not supports_udfs") @@ -3779,8 +4219,8 @@ def test_udt_user_op_named_like_a_lifted_op(): ``"plus"`` or ``"max"`` was compiled as the built-in, its function ignored, and on a built-in dtype it raised KeyError. A literal beside a UDT is converted to the UDT for it, since nothing says how a UDF combines a scalar - with an element, but only when that changes no value: ``0.5`` became an - element of zeros. + with an element, but only when its type fits: ``0.5`` became an element of + zeros. """ udt = dtypes.register_anonymous(np.dtype((np.int64, (2, 3))), "_NamedLikeLifted") @@ -3796,15 +4236,19 @@ def difference(x, y): v[0] = np.ones((2, 3), dtype=np.int64) for udf_op, right, expected in [ (plus_like, 3, 13), - (plus_like, 2.0, 12), + (plus_like, np.int8(2), 12), (max_like, 3, -2), ]: result = v.apply(udf_op, right=right).new() assert result.dtype == udt np.testing.assert_array_equal(result[0].new().value, np.full((2, 3), expected)) assert udf_op[udt].jit_c_source is None - with pytest.raises(ValueError, match="would change 0.5 to 0"): + # A float is not an int, whatever its value, as numpy refuses ``ints += 2.0``; + # 2.0 still converts, being exact, but that is deprecated. + with pytest.raises(ValueError, match="0.5 does not fit"): v.apply(plus_like, right=0.5) + with pytest.warns(DeprecationWarning, match="2.0 does not fit"): + v.apply(plus_like, right=2.0) assert plus_like(Vector.from_coo([0], [5]), 2).new().isequal(Vector.from_coo([0], [52])) @@ -3853,19 +4297,19 @@ def record(i_dtype, f_dtype): @pytest.mark.skipif("not supports_udfs") -def test_udt_infix_plus_on_mixed_types_is_the_union_with_zero(): - """``ewise_add`` with a lifted op on UDT operands of different types runs without a cast. - - GraphBLAS casts an entry present in only one input to the result type, - which a UDT cannot be: a UDT Scalar plus ``0.5`` (infix ``+`` is - ``ewise_add``), an int8 array UDT vector plus a float32 one, and - ``truediv`` on two int UDT vectors of the same type (float64 result) all - raised a bare GrB_DOMAIN_MISMATCH. Each lifted op gives the same values on - operands already converted to its result type, so ``ewise_add`` converts - them (subtracting zero, which keeps ``-0.0``) and applies the op there. +def test_udt_ewise_add_needs_operands_of_the_result_type(): + """``ewise_add`` on UDT operands of another type than the result's: Scalars only. + + GraphBLAS copies an entry present in only one input into the result, cast + to the result type, which a UDT cannot be. A Scalar is one element, so its + operands are converted to the result type first (which changes no value, + not even ``-0.0``). Vectors and Matrices would need converted copies of + whole operands, so they raise, and ``ewise_union`` computes the same sum + with the op instead. """ int8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_UniPlusI8") fp32s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_UniPlusF32") + f64s = dtypes.lookup_dtype(np.dtype((np.float64, (3,)))) s = gb.Scalar.from_value([1, 2, 3], dtype=int8s) for expr in [s + 0.5, 0.5 + s]: result = expr.new() @@ -3873,79 +4317,92 @@ def test_udt_infix_plus_on_mixed_types_is_the_union_with_zero(): np.testing.assert_array_equal(result.value, [1.5, 2.5, 3.5]) result = (gb.Scalar(int8s) + 0.5).new() # an empty Scalar adds as 0 np.testing.assert_array_equal(result.value, [0.5, 0.5, 0.5]) + neg = gb.Scalar.from_value([-0.0, 0.0, -1.5], dtype=fp32s) + result = (neg + gb.Scalar(f64s)).new() + assert result.dtype == f64s + np.testing.assert_array_equal(np.signbit(result.value), [True, False, True]) v = Vector(int8s, size=3) v[0] = [1, 2, 3] v[1] = [4, 5, 6] w = Vector(fp32s, size=3) w[1] = [0.5, 0.5, 0.5] w[2] = [7, 8, 9] - for a, b in [(v, w), (w, v)]: - result = (a + b).new() - assert result.dtype == fp32s - assert result.nvals == 3 - np.testing.assert_array_equal(result[0].new().value, [1, 2, 3]) - np.testing.assert_array_equal(result[1].new().value, [4.5, 5.5, 6.5]) - np.testing.assert_array_equal(result[2].new().value, [7, 8, 9]) A = Matrix(int8s, nrows=1, ncols=2) A[0, 0] = [1, 2, 3] B = Matrix(fp32s, nrows=1, ncols=2) B[0, 1] = [0.5, 0.5, 0.5] + for expr in [lambda: v.ewise_add(w), lambda: w.ewise_add(v), lambda: binary.plus(v | w)]: + with pytest.raises(DomainMismatch, match="ewise_add cannot use .* converted copies"): + expr() + # Infix ``+`` takes the union with each operand's zero for a missing value + # instead, as ``-`` takes 0. + for left, right in [(v, w), (w, v)]: + result = (left + right).new() + assert result.dtype == fp32s + assert result.isequal(left.ewise_union(right, binary.plus, False, False).new()) + assert result.to_coo()[1].tolist() == [[1, 2, 3], [4.5, 5.5, 6.5], [7, 8, 9]] result = (A + B).new() assert result.dtype == fp32s - assert result.nvals == 2 - # A broadcast of the wrong size raises up front, as for a same-type pair, - # not once both operands have been converted. - with pytest.raises(DimensionMismatch, match="Matrix.nrows"): - v.ewise_add(B, binary.plus) - with pytest.raises(DimensionMismatch, match="Matrix.ncols"): - B.ewise_add(v, binary.plus) - # A same-type pair still goes through ewise_add itself. + assert result.to_coo()[2].tolist() == [[1, 2, 3], [0.5, 0.5, 0.5]] + assert (A + B.T.T).new().isequal(result) + result = (A + Vector.from_coo([1], [[0.5, 0.5, 0.5]], dtype=fp32s, size=2)).new() + assert result.to_coo()[2].tolist() == [[1, 2, 3], [0.5, 0.5, 0.5]] + # That zero is -0.0 in a float field, so an entry of one operand alone keeps + # its value, as ewise_add would copy it; beside an int field, 0 turns -0.0 + # into 0.0. + zeros = Vector(fp32s, size=3) + zeros[2] = [-0.0, 0.0, -1.5] + for left, right in [(zeros, Vector(f64s, size=3)), (Vector(f64s, size=3), zeros)]: + result = (left + right).new() + assert result.dtype == f64s + np.testing.assert_array_equal(np.signbit(result[2].new().value), [True, False, True]) + result = (v + zeros).new() + np.testing.assert_array_equal(np.signbit(result[2].new().value), [False, False, True]) + x = v.dup() + x += w # an accumulating store: cast, then added, as for built-in types + assert x.dtype == int8s + assert x.to_coo()[1].tolist() == [[1, 2, 3], [4, 5, 6], [7, 8, 9]] + # The same type in and out still goes through ewise_add itself... assert (v + v).new().dtype == int8s - # The method does the same, for each lifted op; an unpaired entry is kept as is. - assert v.ewise_add(w, binary.plus).new().isequal((v + w).new()) - for binop, both in [(binary.minus, [3.5, 4.5, 5.5]), (binary.max, [4, 5, 6])]: - result = v.ewise_add(w, binop).new() - assert result.dtype == fp32s - np.testing.assert_array_equal(result[0].new().value, [1, 2, 3]) - np.testing.assert_array_equal(result[1].new().value, both) - np.testing.assert_array_equal(result[2].new().value, [7, 8, 9]) - u = Vector(int8s, size=3) - u[1] = [2, 2, 2] - u[2] = [-3, 0, 4] - result = v.ewise_add(u, binary.truediv).new() - assert result.dtype.np_type == np.dtype((np.float64, (3,))) - assert result.to_coo()[1].tolist() == [[1, 2, 3], [2, 2.5, 3], [-3, 0, 4]] - # Converting keeps the sign of a zero, as GraphBLAS's cast does. - f64s = dtypes.lookup_dtype(np.dtype((np.float64, (3,)))) - neg = Vector(fp32s, size=2) - neg[0] = [-0.0, 0.0, -1.5] - pos = Vector(f64s, size=2) - pos[1] = [1, 1, 1] - result = (neg + pos).new() - assert result.dtype == f64s - np.testing.assert_array_equal(np.signbit(result[0].new().value), [True, False, True]) - # Any other op cannot run in the result type, so it is a clear error up front. + # ...but not int truediv, whose result is float. + with pytest.raises(DomainMismatch, match="binary.truediv"): + v.ewise_add(v, binary.truediv) udf = BinaryOp.register_anonymous(lambda x, y: x, "_uni_first", is_udt=True) with pytest.raises(DomainMismatch, match="ewise_add cannot use .* on _UniPlusI8 and"): v.ewise_add(Vector(f64s, size=3), udf) - # A result of another type is cast to the object's type before it is - # accumulated, as for built-in types. + # A size or broadcast error is reported first, as for a same-type pair. + with pytest.raises(DimensionMismatch): + v.ewise_add(Vector(fp32s, size=2), binary.plus) + with pytest.raises(DimensionMismatch, match="Matrix.nrows"): + v.ewise_add(B, binary.plus) + with pytest.raises(DimensionMismatch, match="Matrix.ncols"): + B.ewise_add(v, binary.plus) + # ewise_union computes an entry alone with the op, so it needs no cast. + result = v.ewise_union(w, binary.plus, 0, 0).new() + assert result.dtype == fp32s + assert result.to_coo()[1].tolist() == [[1, 2, 3], [4.5, 5.5, 6.5], [7, 8, 9]] + # Accumulating it into the int8 object casts as it stores, as built-ins do. x = v.dup() - x(accum=binary.plus) << x + w + x(accum=binary.plus) << x.ewise_union(w, binary.plus, 0, 0) assert x.dtype == int8s assert x.to_coo()[1].tolist() == [[2, 4, 6], [8, 10, 12], [7, 8, 9]] @pytest.mark.skipif("not supports_udfs") def test_udt_literal_converts_only_when_exact(): - """A literal that has to become an element of a UDT must keep every value. + """A literal that has to become an element of a UDT must be of that type. - Tuples, lists and arrays beside a UDT, literals beside a user-defined op, - and ``ewise_union`` defaults are converted into the UDT, which numpy does + A literal beside a user-defined op, an IndexUnaryOp thunk and an + ``ewise_union`` default are converted into the UDT, which numpy does silently: ``(0.5, 1.5, 2.5)`` added ``(0, 1, 2)`` to an int8 array UDT, and - a default of ``0.5`` stood in as ``0``. A Monoid in ``apply`` is typed as - its BinaryOp, so its literal promotes instead of being converted. + a default of ``0.5`` stood in as ``0``. The literal's type decides, as + numpy's same_kind rule decides ``ints += x``: a Python number or sequence + is weak, so ``1`` and ``(1, 2, 3)`` are int8 and ``0.5`` and ``2.0`` are + float64; a numpy value is strong and must cast safely. A Monoid in + ``apply`` is typed as its BinaryOp, so its literal promotes instead. """ + import re + arr = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_ExactArr") rec = dtypes.register_anonymous( np.dtype([("ex_i", np.int16), ("ex_f", np.float32)], align=True), "_ExactRec" @@ -3956,34 +4413,49 @@ def test_udt_literal_converts_only_when_exact(): w[1] = [4, 5, 6] r = Vector(rec, size=1) r[0] = (1, 1.5) - # Exact values convert, including a float of a narrower float field. - plus = binary.plus - assert v.apply(plus, right=(1, 2, 3)).new()[0].new().value.tolist() == [2, 4, 6] - assert v.apply(plus, right=[2.0, 2.0, 2.0]).new().dtype == arr - assert (v + np.ones(3, dtype=np.int64)).new().dtype == arr - expected = (2, np.float32(1.5) + np.float32(0.1)) - assert r.apply(plus, right=(1, 0.1)).new()[0].new().value.tolist() == expected - assert r.apply(binary.plus, right={"ex_i": 2, "ex_f": 0.5}).new()[0].new().value.tolist() == ( - 3, - 2.0, - ) - for expr in [ - lambda: v.apply(plus, right=(0.5, 1.5, 2.5)), - lambda: v.apply(binary.times, left=np.array([0.5, 1, 1])), - lambda: r.apply(plus, left=(0.5, 1)), - lambda: v.ewise_union(w, binary.plus, 0.5, 0), - lambda: v.ewise_union(w, binary.plus, 0, 0.5), - lambda: gb.Scalar.from_value([1, 2, 3], arr).ewise_union( - gb.Scalar(arr), binary.plus, 0, 0.5 + udf = BinaryOp.register_anonymous(lambda x, y: x + y, "_exact_plus", is_udt=True) + second = BinaryOp.register_anonymous(lambda x, y: y, "_exact_second", is_udt=True) + # Literals of the UDT's type convert, a float rounding into a float32 field. + assert v.apply(udf, right=(1, 2, 3)).new()[0].new().value.tolist() == [2, 4, 6] + assert v.apply(udf, right=[True, 2, 3]).new().dtype == arr + assert v.apply(udf, right=np.ones(3, dtype=np.int8)).new().dtype == arr + assert v.apply(udf, right=np.int8(1)).new().dtype == arr + assert r.apply(second, right=(1, 0.1)).new()[0].new().value.tolist() == (1, np.float32(0.1)) + assert r.apply(second, right={"ex_i": 2, "ex_f": 0.5}).new()[0].new().value.tolist() == (2, 0.5) + assert r.apply(second, right=np.int8(2)).new()[0].new().value.tolist() == (2, 2.0) + # A sequence holding numpy values is strong, as the array numpy makes of it. + assert binary.plus(v, (np.int8(1), 2, 3)).new().dtype.np_type == np.dtype((np.int64, (3,))) + # The error names the literal's type; a layout's name is whatever UDT + # registered it first, so only the built-in scalar types are pinned here. + for expr, typed_as in [ + (lambda: v.apply(udf, right=(0.5, 1.5, 2.5)), "_ExactArr"), + (lambda: v.apply(udf, left=np.array([0.5, 1, 1])), "_ExactArr"), + (lambda: v.apply(udf, right=(np.float64(0.5), 1, 1)), "_ExactArr"), + (lambda: r.apply(second, left=(0.5, 1)), "_ExactRec"), + (lambda: v.ewise_union(w, udf, 0.5, 0), "_ExactArr"), + (lambda: v.ewise_union(w, udf, 0, 0.5), "_ExactArr"), + ( + lambda: gb.Scalar.from_value([1, 2, 3], arr).ewise_union(gb.Scalar(arr), udf, 0, 0.5), + "_ExactArr", ), - lambda: v == (0.5, 2, 3), ]: - with pytest.raises(ValueError, match="would change 0.5 to 0"): + with pytest.raises(ValueError, match=re.escape(f"does not fit {typed_as}")): + expr() + # One that does not fit by type but converts exactly still converts, as it + # did on main, with a DeprecationWarning. + for expr, typed_as in [ + (lambda: v.apply(udf, right=[2.0, 2.0, 2.0]), "_ExactArr"), + (lambda: v.apply(udf, right=np.ones(3, dtype=np.int64)), "_ExactArr"), + (lambda: v.apply(udf, right=np.int64(1)), "_ExactArr: it is typed as INT64,"), + (lambda: v.apply(udf, right=(np.int8(1), 2, 3)), "_ExactArr"), + (lambda: r.apply(second, right=np.int32(2)), "_ExactRec: it is typed as INT32,"), + ]: + with pytest.warns(DeprecationWarning, match=re.escape(f"does not fit {typed_as}")): expr() with pytest.raises(OverflowError, match="300 out of bounds for int8"): - v.apply(plus, right=(1, 2, 300)) + v.apply(udf, right=(1, 2, 300)) # Defaults that fit still work, and fill in for the missing side. - result = v.ewise_union(w, binary.plus, 10, 20).new() + result = v.ewise_union(w, udf, 10, 20).new() assert result.to_coo()[1].tolist() == [[21, 22, 23], [14, 15, 16]] # apply with a Monoid promotes as its BinaryOp does, on either side. for expr in [v.apply(monoid.plus, right=0.5), v.apply(monoid.plus, left=0.5)]: @@ -3993,20 +4465,218 @@ def test_udt_literal_converts_only_when_exact(): assert v.apply(monoid.max, right=2).new().dtype == arr +# _SeqNest's numpy repr is over 128 chars on numpy 1; see test_udt_eq_nested_record_with_nan_leaf. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +@pytest.mark.skipif("not supports_udfs") +def test_udt_sequence_and_array_literals_type_like_numbers(): + """A tuple, list or array beside a UDT under a lifted op is typed, not converted into the UDT. + + A tuple or list of Python numbers is weak leaf by leaf, as a Python number + is: ``int8_udt + (1, 2, 3)`` stays int8 and ``+ (0.5, 1.5, 2.5)`` is + float64, where converting it raised (and before that added ``(0, 1, 2)``). + A numpy array is strong, as its own array type. ``eq`` and ``ne`` compare + the values. A literal beside a user-defined op is still converted + (``test_udt_literal_converts_only_when_exact``). + """ + int8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_SeqI8") + fp32s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_SeqF32") + nested = dtypes.register_anonymous( + np.dtype( + [("sq_a", np.int16), ("sq_in", [("sq_b", np.int8), ("sq_c", np.float32)])], + align=True, + ), + "_SeqNest", + ) + f64s = np.dtype((np.float64, (3,))) + v = Vector(int8s, size=2) + v[0] = [1, 2, 3] + v[1] = [4, 5, 6] + # Weak: the leaf's dtype when its kind holds the numbers, of the widest kind among them. + for literal in [(1, 2, 3), [1, 2, 3], (True, 2, 3), (1,)]: + result = (v + literal).new() + assert result.dtype == int8s + np.testing.assert_array_equal(result[0].new().value, np.add([1, 2, 3], literal)) + halves = (0.5, 1.5, 2.5) + for expr in [v + halves, halves + v, v.apply(binary.plus, right=list(halves))]: + result = expr.new() + assert result.dtype.np_type == f64s + np.testing.assert_array_equal(result[0].new().value, [1.5, 3.5, 5.5]) + result = binary.plus(v, (1, 0.5, 3)).new() # one float anywhere makes the leaf float + assert result.dtype.np_type == f64s + np.testing.assert_array_equal(result[0].new().value, [2, 2.5, 6]) + result = binary.times(v, (1j, 1, 1)).new() + assert result.dtype.np_type == np.dtype((np.complex128, (3,))) + np.testing.assert_array_equal(result[0].new().value, [1j, 2, 3]) + with pytest.raises(OverflowError, match="300 out of bounds for int8"): + binary.plus(v, (1, 2, 300)) + v += (1, 1, 1) + np.testing.assert_array_equal(v[0].new().value, [2, 3, 4]) + # The float64 result is cast as it is stored, as a built-in int8 vector's would be. + v += (0.5, 1, 1) + np.testing.assert_array_equal(v[0].new().value, [2, 4, 5]) + # A float list beside float32 stays float32, as 0.5 does. + x = Vector(fp32s, size=1) + x[0] = [1, 2, 3] + assert binary.plus(x, [0.5, 0.5, 0.5]).new().dtype == fp32s + # Strong: an array is its own dtype, in its shape, or broadcast to the UDT's. + for literal, expected in [ + (np.array([0.5, 1, 2]), f64s), + (np.array([1, 2, 3]), np.dtype((np.int64, (3,)))), + (np.array([1], np.int8), int8s.np_type), + (np.array([0.5]), f64s), + (np.ones((1, 3)), np.dtype((np.float64, (1, 3)))), + ]: + result = (v + literal).new() + assert result.dtype.np_type == expected + np.testing.assert_array_equal(result[0].new().value, np.add([2, 4, 5], literal)) + # Records: by position, nested for a nested record, a dict by name, or a + # 0-d structured array of another record type. + n = Vector(nested, size=1) + n[0] = (1, (2, 3.5)) + assert binary.plus(n, (1, (1, 1))).new().dtype == nested + result = binary.plus(n, (1, (0.5, 1))).new() + assert result.dtype.np_type == np.dtype( + [("sq_a", np.int16), ("sq_in", [("sq_b", np.float64), ("sq_c", np.float32)])], + align=True, + ) + assert result[0].new().value.tolist() == (2, (2.5, 4.5)) + assert (n + {"sq_a": 0.5, "sq_in": {"sq_b": 1, "sq_c": 1}}).new()[0].new().value.tolist() == ( + 1.5, + (3, 4.5), + ) + strong = np.array( + (1, (0.5, 1)), + dtype=[("sq_a", np.int8), ("sq_in", [("sq_b", np.float64), ("sq_c", np.int8)])], + ) + assert (n + strong).new()[0].new().value.tolist() == (2, (2.5, 4.5)) + # eq and ne take the literal as it is, like numpy, instead of raising. + assert (v == (0.5, 4, 5)).new().isequal(Vector.from_coo([0, 1], [False, False])) + assert (v == (2, 4, 5)).new().isequal(Vector.from_coo([0, 1], [True, False])) + assert (v != (2, 4, 500)).new().isequal(Vector.from_coo([0, 1], [True, True])) + assert (v == np.array([2.0, 4, 5])).new()[0].new().value + # Matrix and Scalar take the same path. + A = Matrix(int8s, nrows=1, ncols=1) + A[0, 0] = [1, 2, 3] + ones = [1, 1, 1] + assert (A + halves).new().dtype.np_type == f64s + assert (A + ones).new().dtype == int8s + s = gb.Scalar.from_value([1, 2, 3], int8s) + assert (s + halves).new().dtype.np_type == f64s + assert (s * ones).new().dtype == int8s + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_ewise_union_defaults_are_typed_not_converted(): + """An ``ewise_union`` default beside a UDT is typed as a literal and must fit the operand. + + GraphBLAS casts a default to the op's input type, which for a UDT operand + is its own type, and converting the operands instead would copy them. So + a default is typed by the literal rules and accepted only when that type + fits: ``1`` and ``(1, 1, 1)`` fit an int8 UDT, ``0.5``, ``np.int64(1)`` and + a float64 Scalar do not. A Scalar (one element) of a type that casts + safely is converted. + """ + import re + + int8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_UnionI8") + fp32s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_UnionF32") + v = Vector(int8s, size=3) + v[0] = [0, 1, 2] + v[1] = [4, 5, 6] + w = Vector(int8s, size=3) + w[1] = [1, 1, 1] + w[2] = [7, 8, 9] + plus = binary.plus + expected = [[1, 2, 3], [5, 6, 7], [7, 8, 9]] + for left, right in [ + (0, 1), + ((0, 0, 0), [1, 1, 1]), + (np.int8(0), np.array([1, 1, 1], np.int8)), + (gb.Scalar.from_value(0, "INT8"), gb.Scalar.from_value([1, 1, 1], int8s)), + (False, True), + ]: + for expr in [ + v.ewise_union(w, plus, left, right), + v.ewise_union(w, monoid.plus, left, right), + ]: + result = expr.new() + assert result.dtype == int8s + assert result.to_coo()[1].tolist() == expected + + def union_forms(d): + return [ + lambda: v.ewise_union(w, plus, d, 0), + lambda: v.ewise_union(w, plus, 0, d), + lambda: plus(v | w, left_default=d, right_default=0), + lambda: v.ewise_union(w, monoid.plus, d, 0), + ] + + for default, message in [ + (0.5, "0.5 does not fit _UnionI8"), + ((0.5, 0.5, 0.5), "(0.5, 0.5, 0.5) does not fit _UnionI8"), + (np.float32(0.5), "does not fit _UnionI8: it is typed as FP32,"), + (gb.Scalar.from_value(0.5), "Scalar of type FP64 does not fit _UnionI8"), + ]: + for expr in union_forms(default): + with pytest.raises(ValueError, match=re.escape(message)): + expr() + # Exact values of another type still convert, as on main, with a + # DeprecationWarning. + for default, message in [ + (2.0, "2.0 does not fit _UnionI8"), + (np.int64(1), "does not fit _UnionI8: it is typed as INT64,"), + (np.array([1, 1, 1]), "array([1, 1, 1]) does not fit _UnionI8"), + (gb.Scalar.from_value(1), "Scalar of type INT64 does not fit _UnionI8"), + ]: + for expr in union_forms(default): + with pytest.warns(DeprecationWarning, match=re.escape(message)): + expr() + with pytest.raises(OverflowError, match="300 out of bounds for int8"): + v.ewise_union(w, plus, 300, 0) + # Each side is typed beside its own operand, so a mixed pair takes a + # default of each type, and an int8 Scalar casts safely into float32. + x = Vector(fp32s, size=3) + x[2] = [0.5, 0.5, 0.5] + result = v.ewise_union(x, plus, 0, 0.5).new() + assert result.dtype == fp32s + assert result.to_coo()[1].tolist() == [[0.5, 1.5, 2.5], [4.5, 5.5, 6.5], [0.5, 0.5, 0.5]] + result = v.ewise_union(x, plus, 0, gb.Scalar.from_value([1, 1, 1], int8s)).new() + assert result.to_coo()[1].tolist() == [[1, 2, 3], [5, 6, 7], [0.5, 0.5, 0.5]] + with pytest.raises(ValueError, match="0.5 does not fit _UnionI8"): + v.ewise_union(x, plus, 0.5, 0) + # Matrix and Scalar check the same way. + A = Matrix(int8s, nrows=2, ncols=1) + A[0, 0] = [0, 1, 2] + B = Matrix(int8s, nrows=2, ncols=1) + B[1, 0] = [7, 8, 9] + assert A.ewise_union(B, plus, 1, 1).new().to_coo()[2].tolist() == [[1, 2, 3], [8, 9, 10]] + s = gb.Scalar.from_value([1, 2, 3], int8s) + assert s.ewise_union(gb.Scalar(int8s), plus, 0, 1).new().value.tolist() == [2, 3, 4] + for expr in [ + lambda: A.ewise_union(B, plus, 0.5, 0), + lambda: A.ewise_union(v, plus, 0, 0.5), + lambda: s.ewise_union(gb.Scalar(int8s), plus, 0, 0.5), + ]: + with pytest.raises(ValueError, match="0.5 does not fit _UnionI8"): + expr() + + @pytest.mark.skipif("not supports_udfs") # SS < 9 names a UDT by its numpy repr and warns when that is over 128 chars, # as the records with a field per dtype pair are. @pytest.mark.filterwarnings("ignore:UDT repr is too large") -def test_udt_store_casts_each_element_as_builtin_types_do(monkeypatch): +def test_udt_store_casts_each_element_as_builtin_types_do(): """A UDT result stored in an object of another UDT type is cast element by element. GraphBLAS cannot cast a UDT, so storing across UDT types raised, and setting one element from a Scalar of another UDT type read its bytes as the object's type. Each array element or record leaf is now cast as GraphBLAS - casts built-in types on store (compared below), for every way of storing: - ``<<``, a mask, an accumulator, ``replace``, ``.new(dtype=)``, Scalars and - assignment. Layouts that do not correspond, and a built-in type on either - side, raise ``DomainMismatch``. + casts built-in types on store (compared below), without a temporary copy: + an object stores through the cast op (``<<``, a mask, an accumulator, + ``replace``, ``dup``), an element-wise expression of a lifted op casts as + it computes, and a Scalar or one element is converted. Anything that would + need a converted copy of a whole object raises, as do layouts that do not + correspond and a built-in type on either side. """ # One record casts every pair of leaf dtypes, so it compiles once. pairs = [ @@ -4064,7 +4734,7 @@ def test_udt_store_casts_each_element_as_builtin_types_do(monkeypatch): w = Vector.from_coo([0, 1, 2], [ones] * 3, dtype=i8s) w(mask.S, accum=binary.plus) << v assert w.to_coo()[1].tolist() == [[[2, 0], [-128, 1], [-127, 3]], ones, ones] - w(mask.S) << v * 1 # an expression, computed into a temporary first + w(mask.S) << v * 1 # an expression of a lifted op casts as it computes assert w.to_coo()[1].tolist() == [first, ones] w(mask.S, replace=True) << v assert w.to_coo()[1].tolist() == [first] @@ -4085,8 +4755,25 @@ def test_udt_store_casts_each_element_as_builtin_types_do(monkeypatch): C = Matrix(i8s, nrows=2, ncols=2) C << A.T assert C.to_coo()[2].tolist() == [first] - C[0, :] = v[[0, 1]].new() + # A whole Vector of another type would need a converted copy to assign. + with pytest.raises(DomainMismatch, match="converted copy of the whole value"): + C[0, :] = v[[0, 1]].new() + C[0, :] = v[[0, 1]].new().dup(dtype=i8s) assert C[0, 0].new().value.tolist() == first + # So would a computation that cannot write another type as it computes. + with pytest.raises(DomainMismatch, match="would need a temporary copy"): + C << A.mxm(A, semiring.plus_times) + with pytest.raises(DomainMismatch, match="would need a temporary copy"): + w << v.ewise_add(v, binary.plus) + with pytest.raises(DomainMismatch, match="extraction cannot"): + w << v[[0, 1, 2]] + # The computed result, stored as an existing object, casts in one pass. + C << A.mxm(A, semiring.plus_times).new() + # power raises before it computes any product; n of 0 and 1 store as above. + with pytest.raises(DomainMismatch, match="mxm with plus_times cannot"): + C << A.power(3) + C << A.power(1) + assert C.to_coo()[2].tolist() == [first] # Leading axes of length 1 do not change the layout. lead = dtypes.register_anonymous(np.dtype((np.float32, (1, 3, 2))), "_StoreCastLead") u = Vector(lead, size=3) @@ -4099,6 +4786,9 @@ def test_udt_store_casts_each_element_as_builtin_types_do(monkeypatch): renamed = np.dtype([(f"sr{i}", dst) for i, (_src, dst) in enumerate(pairs)]) with pytest.raises(DomainMismatch, match="same field names"): Vector(dtypes.register_anonymous(renamed), size=7) << x + raw = Vector(dtypes.register_anonymous(np.dtype("S48"), "_StoreCastBytes"), size=3) + with pytest.raises(DomainMismatch, match="only record and array UDTs cast"): + raw << v for expr in [ lambda: Vector(f64s, size=3) << Vector.from_coo([0], [1.5], size=3), lambda: Vector(dtypes.FP64, size=3) << v, @@ -4106,21 +4796,311 @@ def test_udt_store_casts_each_element_as_builtin_types_do(monkeypatch): ]: with pytest.raises(DomainMismatch, match="a UDT casts only to another UDT"): expr() - # The switch back to hard: every UDT store across types raises. - monkeypatch.setattr(gb.core.base, "_udt_store_casts", False) - for expr in [lambda: w << v, lambda: w.__setitem__(0, v[0].new())]: - with pytest.raises(DomainMismatch, match="UDTs do not cast"): - expr() + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_extracted_element_of_another_type_is_cast(): + """One element extracted into another UDT type is cast, not copied as raw bytes. + + GraphBLAS's extractElement for a UDT copies the bytes of the object's type + into the destination, so ``s << v[0]`` with ``s`` of another UDT read + garbage, and wrote past the end of a smaller element (also on main). + """ + f64s = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_ExtractF64") + i16s = dtypes.register_anonymous(np.dtype((np.int16, (2,))), "_ExtractI16") + rec = dtypes.register_anonymous(np.dtype([("ex_a", np.int8), ("ex_b", np.int8)]), "_ExtractRec") + v = Vector(f64s, size=3) + v[0] = [1.7, -300.2] + A = Matrix(f64s, nrows=2, ncols=3) + A[1, 2] = [2.5, -1e30] + s = gb.Scalar(i16s) + s << v[0] + assert s.value.tolist() == [1, -300] + assert v[0].new(dtype=i16s).value.tolist() == [1, -300] + assert A.T[2, 1].new(dtype=i16s).value.tolist() == [2, -32768] # saturates + w = Vector(i16s, size=3) + C = Matrix(i16s, nrows=2, ncols=2) + with gb.config.set(autocompute=True): + w[1] = v[0] + C[1, 0] = A[1, 2] + assert w[1].new().value.tolist() == [1, -300] + assert C[1, 0].new().value.tolist() == [2, -32768] + # A larger destination, and a missing element. + x = Vector(i16s, size=2) + x[0] = [7, -8] + assert x[0].new(dtype=f64s).value.tolist() == [7.0, -8.0] + assert x[1].new(dtype=f64s).is_empty + for dtype in [rec, dtypes.FP64]: + with pytest.raises(DomainMismatch, match="cannot store"): + v[0].new(dtype=dtype) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_fused_store_matches_compute_then_cast(): + """Storing an element-wise result into another UDT type casts as the op computes. + + The op computes each element as it would for its own result type and casts + it into the object's type as it writes it, so the values are those of + computing the result and then casting it (``.new().dup(dtype=)``), with no + temporary of the result's type: special values, records with array + fields, literals on either side, unary ops, masks and accumulators. + """ + i16s = dtypes.register_anonymous(np.dtype((np.int16, (2, 2))), "_FuseI16") + f32s = dtypes.register_anonymous(np.dtype((np.float32, (2, 2))), "_FuseF32") + rec_i = dtypes.register_anonymous( + np.dtype([("fz_a", np.int8), ("fz_b", np.uint16, (3,))]), "_FuseRecI" + ) + rec_f = dtypes.register_anonymous( + np.dtype([("fz_a", np.float64), ("fz_b", np.float64, (3,))]), "_FuseRecF" + ) + i = Vector(i16s, size=4) + i[0] = [[1, -2], [32767, -32768]] + i[2] = [[7, 0], [-7, 3]] + f = Vector(f32s, size=4) + f[0] = [[np.nan, -np.inf], [1e10, -0.5]] + f[1] = [[2.5, -2.5], [65535.9, 1.5]] + f[2] = [[0.25, 3.75], [-1.0, 2.0]] + r = Vector(rec_f, size=4) + r[0] = (300.7, [-1.5, 70000.2, np.nan]) + r[1] = (-129.9, [2.5, 3.5, 1e300]) + cases = [ + (lambda: i * 2.5, i16s), + (lambda: 2.5 - i, i16s), + (lambda: i / 4, i16s), + (lambda: i.ewise_mult(f, binary.times), i16s), + (lambda: f.ewise_mult(i, binary.minus), i16s), + (lambda: i.ewise_union(f, binary.plus, 0, 0), i16s), + (lambda: f.apply(unary.ainv), i16s), + (lambda: f.apply(monoid.max, right=3), i16s), + (lambda: f * f, i16s), + (lambda: f.ewise_mult(f, monoid.plus), i16s), + (lambda: r * 1.5, rec_i), + (lambda: r.apply(unary.abs), rec_i), + ] + mask = Vector.from_coo([0, 2], [True, True], size=4) + for make, out in cases: + expected = make().new().dup(dtype=out) + fused = Vector(out, size=expected.size) + fused << make() + assert fused.isequal(expected), make + # A mask and an accumulator act on the cast values, as they would on + # values that were already of the object's type. + start = expected.dup() + explicit = start.dup() + explicit(mask.S, accum=binary.plus) << make().new().dup(dtype=out) + start(mask.S, accum=binary.plus) << make() + assert start.isequal(explicit), make + A = Matrix(f32s, nrows=2, ncols=2) + A[0, 1] = [[1.5, -2.5], [np.nan, 1e9]] + C = Matrix(i16s, nrows=2, ncols=2) + C << A.T.apply(binary.times, right=2) + assert C.isequal(A.T.apply(binary.times, right=2).new().dup(dtype=i16s)) + + +def _udt_leaf_arrays(vector, size): + """The values of a record UDT vector's entries 0 to ``size - 1``, one array per leaf.""" + values = [vector[k].new().value for k in range(size)] + return { + name: np.array([value[name] for value in values]) for name in vector.dtype.np_type.names + } + + +def _assert_same_floats(got, expected, msg): + """Equal, NaN equal to NaN, and with the same sign on every zero (of each complex part).""" + got, expected = np.asarray(got), np.asarray(expected) + np.testing.assert_array_equal(got, expected, err_msg=msg) + for part in (np.real, np.imag) if got.dtype.kind == "c" else (np.asarray,): + np.testing.assert_array_equal( + np.signbit(part(got)), np.signbit(part(expected)), err_msg=msg + ) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_udt_record_array_fields_compute_element_by_element(): + """An array-valued record field computes each element as a scalar field does. + + As a whole-array expression it ran through Numba's ufuncs: ``min`` and + ``max`` did not compile (and the scalar NaN rule, applied to an array, left + garbage), a zero complex divisor raised inside the cfunc and left the + element unwritten, and ``int8 / int8`` ran in float32. Each element now + takes the scalar expression, so a field agrees with numpy's ``fmin``, + ``fmax``, ``floor_divide`` and ``true_divide`` element for element, with a + record or a number on the other side. ``INT64_MIN // -1`` is the one + integer quotient Numba gets wrong by itself (0, where numpy wraps). + """ + nan, inf = float("nan"), float("inf") + rec = dtypes.register_anonymous( + np.dtype( + [ + ("aef_s", np.float64), + ("aef_f", np.float64, (2, 3)), + ("aef_i", np.int8, (4,)), + ("aef_l", np.int64, (3,)), + ("aef_m", np.int64), + ], + align=True, + ), + "_ArrElemFields", + ) + fx = np.array([[nan, 1.0, -0.0], [inf, 1.0, 7.5]]) + fy = np.array([[1.0, nan, 2.0], [2.0, 0.1, -2.5]]) + ix = np.array([-128, 7, -7, 100], np.int8) + iy = np.array([-1, 0, 2, 3], np.int8) + int64_min = np.iinfo(np.int64).min + lx = np.array([int64_min, -7, 9], np.int64) + ly = np.array([-1, 2, 0], np.int64) + v = Vector(rec, size=1) + v[0] = (nan, fx, ix, lx, int64_min) + w = Vector(rec, size=1) + w[0] = (1.0, fy, iy, ly, -1) + # The references overflow (``-128 // -1``) and divide by zero on purpose. + with np.errstate(all="ignore"): + for gb_op, reference in [ + (binary.min, np.fmin), + (binary.max, np.fmax), + (binary.floordiv, np.floor_divide), + (binary.truediv, np.true_divide), + ]: + got = gb_op(v & w).new()[0].new().value + _assert_same_floats(got["aef_s"], reference(nan, 1.0), gb_op.name) + _assert_same_floats(got["aef_f"], reference(fx, fy), gb_op.name) + np.testing.assert_array_equal(got["aef_i"], reference(ix, iy), err_msg=gb_op.name) + np.testing.assert_array_equal(got["aef_l"], reference(lx, ly), err_msg=gb_op.name) + np.testing.assert_array_equal( + got["aef_m"], reference(np.int64(int64_min), np.int64(-1)), err_msg=gb_op.name + ) + # A number on the other side reaches every element of the field. + for other in [np.float64(0.5), np.float64(nan)]: + got = v.apply(gb_op, right=other).new()[0].new().value + _assert_same_floats(got["aef_f"], reference(fx, other), f"{gb_op.name} {other}") + got = v.apply(gb_op, right=np.int8(-1)).new()[0].new().value + np.testing.assert_array_equal( + got["aef_i"], reference(ix, np.int8(-1)), err_msg=gb_op.name + ) + np.testing.assert_array_equal( + got["aef_l"], reference(lx, np.int8(-1)), err_msg=gb_op.name + ) + + cplx = dtypes.register_anonymous( + np.dtype([("aef_c", np.complex128, (3,))], align=True), "_ArrElemComplex" + ) + cx = np.array([3 + 4j, 1j, 2 - 2j]) + cy = np.array([0j, 2, 0j]) + v = Vector(cplx, size=1) + v[0] = (cx,) + w = Vector(cplx, size=1) + w[0] = (cy,) + with np.errstate(divide="ignore", invalid="ignore"): + expected = cx / cy + np.testing.assert_array_equal(binary.truediv(v & w).new()[0].new().value["aef_c"], expected) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_min_max_and_division_match_numpy(udt_op_path): + """Array UDT elements take ``fmin``, ``fmax`` and numpy's divisions on both paths. + + NaN sits on each side in turn, and ``INT64_MIN // -1`` is the integer + quotient Numba gets wrong by itself. + """ + nan, inf = float("nan"), float("inf") + f64 = dtypes.register_anonymous(np.dtype((np.float64, (2, 3))), "_ArrNumpyF64") + i64 = dtypes.register_anonymous(np.dtype((np.int64, (3,))), "_ArrNumpyI64") + fx = np.array([[nan, 1.0, -7.5], [inf, 1.0, 0.0]]) + fy = np.array([[1.0, nan, 2.0], [2.0, 0.1, 0.0]]) + int64_min = np.iinfo(np.int64).min + ix = np.array([int64_min, -7, 9], np.int64) + iy = np.array([-1, 2, 0], np.int64) + # Two entries each, with the values swapped in the second, so neither + # vector is iso-valued (SuiteSparse answers those without a kernel). + v = Vector(f64, size=2) + w = Vector(f64, size=2) + v[0], w[0], v[1], w[1] = fx, fy, fy, fx + a = Vector(i64, size=2) + b = Vector(i64, size=2) + a[0], b[0], a[1], b[1] = ix, iy, ix[::-1], iy[::-1] + # The references divide by zero and overflow on purpose. + with np.errstate(all="ignore"): + for gb_op, reference in [ + (binary.min, np.fmin), + (binary.max, np.fmax), + (binary.floordiv, np.floor_divide), + (binary.truediv, np.true_divide), + ]: + got = gb_op(v & w).new() + for k, (x, y) in enumerate([(fx, fy), (fy, fx)]): + msg = f"{udt_op_path} {gb_op.name} [{k}]" + _assert_same_floats(got[k].new().value, reference(x, y), msg) + got = gb_op(a & b).new() + for k, (x, y) in enumerate([(ix, iy), (ix[::-1], iy[::-1])]): + msg = f"{udt_op_path} {gb_op.name} [{k}]" + np.testing.assert_array_equal(got[k].new().value, reference(x, y), err_msg=msg) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_fused_store_keeps_min_max_and_division_semantics(udt_op_path): + """A fused store computes ``min``, ``max`` and the divisions as the op does, then casts. + + The fused op is cfunc only, so with the C JIT on, the plain op runs the C + kernel and the two paths are compared here directly: NaN in ``min`` and + ``max``, zero divisors, ``INT_MIN // -1``, a complex zero divisor, and + array-valued record fields. + """ + nan, inf = float("nan"), float("inf") + f64 = dtypes.register_anonymous(np.dtype((np.float64, (2, 2))), "_FuseSemF64") + f32 = dtypes.register_anonymous(np.dtype((np.float32, (2, 2))), "_FuseSemF32") + i8s = dtypes.register_anonymous(np.dtype((np.int8, (4,))), "_FuseSemI8") + i32s = dtypes.register_anonymous(np.dtype((np.int32, (4,))), "_FuseSemI32") + rec64 = dtypes.register_anonymous( + np.dtype([("fss_a", np.float64), ("fss_b", np.float64, (3,))], align=True), "_FuseSemR64" + ) + rec32 = dtypes.register_anonymous( + np.dtype([("fss_a", np.float32), ("fss_b", np.float32, (3,))], align=True), "_FuseSemR32" + ) + c128 = dtypes.register_anonymous( + np.dtype([("fss_c", np.complex128)], align=True), "_FuseSemC128" + ) + c64 = dtypes.register_anonymous(np.dtype([("fss_c", np.complex64)], align=True), "_FuseSemC64") + f, g = _udt_vectors(f64, [nan, 1.0, -7.0, inf, 0.0], [1.0, nan, 2.0, 2.0, 0.0]) + i, j = _udt_vectors(i8s, [-128, 7, -7, 100, 5], [-1, 0, 2, 3, -2]) + r = Vector(rec64, size=3) + s = Vector(rec64, size=3) + for k, (x, y) in enumerate([(nan, 1.0), (1.0, nan), (-7.5, 2.0)]): + r[k] = (x, [x, y, x * 2]) + s[k] = (y, [y, x, 0.0]) + c = Vector(c128, size=2) + d = Vector(c128, size=2) + c[0], d[0] = (3 + 4j,), (0j,) + c[1], d[1] = (1j,), (2 + 0j,) + cases = [(f, g, f32, op) for op in ("min", "max", "floordiv", "truediv")] + cases += [(i, j, i32s, "floordiv"), (r, s, rec32, "min"), (r, s, rec32, "max")] + cases += [(r, s, rec32, "floordiv"), (c, d, c64, "truediv")] + for x, y, out, name in cases: + gb_op = getattr(binary, name) + expected = gb_op(x & y).new().dup(dtype=out) + fused = Vector(out, size=x.size) + fused << gb_op(x & y) + assert fused.nvals == expected.nvals, (udt_op_path, name, out) + for k in range(x.size): + got, want = fused[k].new().value, expected[k].new().value + msg = f"{udt_op_path} {name} into {out} at {k}" + if out.np_type.names is None: + _assert_same_floats(got, want, msg) + else: + for field in out.np_type.names: + _assert_same_floats(got[field], want[field], msg) @pytest.mark.skipif("not supports_udfs") @pytest.mark.filterwarnings("ignore:UDT repr is too large") def test_udt_records_differing_only_in_layout_combine(): - """A packed and an aligned record with the same leaves add, either way round. + """A packed and an aligned record with the same leaves combine, either way round. The result is the left operand's type, since it holds every leaf. Converting - the other operand to it for ``ewise_add`` subtracted a zero, which came back - as the other operand's own type, so the conversion recursed forever. + the other operand to it for a Scalar's ``ewise_add`` subtracted a zero, + which came back as the other operand's own type, so the conversion recursed + forever; Vectors raised RecursionError the same way before ``ewise_add`` + refused operands of another type than the result's. """ fields = [("lay_f", np.float64), ("lay_i", np.int16), ("lay_n", [("lay_k", np.int32)])] packed = dtypes.register_anonymous(np.dtype(fields), "_LayPacked") @@ -4131,20 +5111,22 @@ def test_udt_records_differing_only_in_layout_combine(): y[0] = (-0.0, 1, (1,)) y[2] = (10.0, 20, (30,)) for a, b in [(x, y), (y, x)]: - for expr in [a + b, a.ewise_add(b, binary.max), a.ewise_add(b)]: - result = expr.new() - assert result.dtype == a.dtype - assert result.nvals == 2 - assert result[2].new().value.tolist() == (10.0, 20, (30,)) - assert (a + b).new()[0].new().value.tolist() == (1.5, 3, (4,)) + result = a.ewise_union(b, binary.plus, 0, 0).new() + assert result.dtype == a.dtype + assert result.to_coo()[1].tolist() == [(1.5, 3, (4,)), (10.0, 20, (30,))] + assert (a * b).new().dtype == a.dtype + with pytest.raises(DomainMismatch, match="ewise_add cannot use"): + a.ewise_add(b) + assert (a + b).new().isequal(result) s = gb.Scalar.from_value((1.5, 2, (3,)), dtype=packed) t = gb.Scalar.from_value((1.0, 1, (1,)), dtype=aligned) assert (s + t).new().value.tolist() == (2.5, 3, (4,)) # Converting keeps a negative zero where there is no partner. - z = Vector(aligned, size=1) - z[0] = (-0.0, 0, (0,)) - result = (Vector(packed, size=1) + z).new() - assert np.signbit(result[0].new().value["lay_f"]) + z = gb.Scalar.from_value((-0.0, 0, (0,)), dtype=aligned) + assert np.signbit((gb.Scalar(packed) + z).new().value["lay_f"]) + # Either layout stores into the other through the cast op. + x << y + assert x.to_coo()[1].tolist() == [(-0.0, 1, (1,)), (10.0, 20, (30,))] @pytest.mark.skipif("not supports_udfs") @@ -4194,11 +5176,15 @@ def test_udt_typed_scalar_operands_and_defaults(): B = Matrix(i8s, nrows=1, ncols=2) B[0, 1] = [4, 5, 6] for a, b in [(v, w), (A, B)]: - with pytest.raises(ValueError, match="would change 0.5 to 0"): + # A default must have the operand's type, and FP64 does not cast safely to + # int8; 2.0 still converts, being exact, but that is deprecated. + with pytest.raises(ValueError, match="Scalar of type FP64 does not fit _TsI8"): a.ewise_union(b, binary.plus, gb.Scalar.from_value(0.5), 0) - with pytest.raises(ValueError, match="would change 0.5 to 0"): + with pytest.warns(DeprecationWarning, match="Scalar of type FP64 does not fit _TsI8"): + a.ewise_union(b, binary.plus, gb.Scalar.from_value(2.0), 0) + with pytest.raises(ValueError, match="Scalar of type _TsF8 does not fit _TsI8"): a.ewise_union(b, binary.plus, 0, sf) - assert a.ewise_union(b, binary.plus, gb.Scalar.from_value(2.0), 0).new().dtype == i8s + assert a.ewise_union(b, binary.plus, gb.Scalar.from_value(2, "INT8"), 0).new().dtype == i8s # Matrix.apply with a Monoid promotes as Vector.apply does. result = A.apply(monoid.plus, right=0.5).new() assert result.dtype.np_type == np.dtype((np.float64, (3,))) @@ -4225,11 +5211,12 @@ def test_udt_errors_name_the_real_problem(): @pytest.mark.skipif("not supports_udfs") def test_udt_literal_rules_edge_cases(): - """Edges of the weak rule and of the exactness check, each pinned on its own. + """Edges of the weak rule, each pinned on its own. numpy 2 raises OverflowError itself when ``300`` is converted to int8, which hid that the weak rule's own range check matters on numpy 1, where the - conversion wraps. + conversion wraps. A ``Fraction`` or ``Decimal`` is typed as the int or + float it is worth; numpy would have truncated either into an int field. """ from decimal import Decimal from fractions import Fraction @@ -4246,13 +5233,18 @@ def test_udt_literal_rules_edge_cases(): v = Vector(i8s, size=1) v[0] = [1, 2, 3] for value in [Fraction(1, 2), Decimal("0.5")]: - with pytest.raises(ValueError, match="would change"): - v + value - assert (v + Fraction(2)).new()[0].new().value.tolist() == [3, 4, 5] + result = (v + value).new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + assert result[0].new().value.tolist() == [1.5, 2.5, 3.5] + for value in [Fraction(2), Decimal(2)]: + result = (v + value).new() + assert result.dtype == i8s + assert result[0].new().value.tolist() == [3, 4, 5] + udf = BinaryOp.register_anonymous(lambda x, y: x + y, "_edge_plus", is_udt=True) + with pytest.raises(ValueError, match=r"Fraction\(1, 2\) does not fit _EdgeI8"): + v.apply(udf, right=Fraction(1, 2)) c = Vector(c8s, size=1) c[0] = [1, 2, 3] - with pytest.raises(ValueError, match="would change 1e[+]300 to"): - c.apply(binary.plus, right=(1e300, 0, 0)) assert c.apply(binary.plus, right=(1 + 2j, 0, 0)).new()[0].new().value[0] == 2 + 2j @@ -4283,11 +5275,15 @@ def test_udt_store_casts_reach_every_store_path(): w = Vector(f8s, size=3) w[0] = [0.5, 0.5, 0.5] w[1] = [1.5, 1.5, 1.5] + # Assigning part of an object from a whole object of another type would + # need a converted copy, an expression's (extract) as well; so would + # ewise_add, which copies an entry without a partner as it is. x = v.dup() - x[[0, 1]] << w[[0, 1]] + for expr in [lambda: x[[0, 1]] << w[[0, 1]], lambda: x[[0, 2]] << (v + 0.5).new()[[0, 2]]]: + with pytest.raises(DomainMismatch, match="converted copy of the whole value"): + expr() + x[[0, 1]] << w[[0, 1]].new().dup(dtype=i8s) assert x.to_coo()[1].tolist() == [[0, 0, 0], [1, 1, 1], [4, 5, 6]] - x[[0, 2]] << (v + 0.5).new()[[0, 2]] - assert x.to_coo()[1].tolist() == [[1, 2, 3], [1, 1, 1], [4, 5, 6]] s = gb.Scalar.from_value([1, 2, 3], dtype=i8s) s += 0.5 np.testing.assert_array_equal(s.value, [1, 2, 3]) @@ -4300,6 +5296,139 @@ def test_udt_store_casts_reach_every_store_path(): s.dup(dtype="FP64") +@pytest.mark.skipif("not supports_udfs") +def test_udt_literal_fits_in_select_bool_arrays_and_overflow(): + """Three edges of the literal rules found by review. + + A user IndexUnaryOp's thunk in ``select`` skipped the must-fit check (1.5 + became 1), a bool numpy array beside an array UDT failed to look up a + ``BOOL[3]`` type (main converted it), and a float literal too large for a + float32 field became infinity where it must become an element. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_EdgeSelI8") + f4s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_EdgeSelF4") + + def ge_first(x, i, j, t): # pragma: no cover (numba) + return x[0] >= t[0] + + ge = IndexUnaryOp.register_anonymous(ge_first, "_edge_sel_ge", is_udt=True) + v = Vector(i8s, size=3) + v[0] = [1, 2, 3] + v[1] = [2, 3, 4] + assert v.select(ge, 2).new().to_coo()[0].tolist() == [1] + for thunk in [1.5, (1.5, 0, 0), np.float64(2.5)]: + with pytest.raises(ValueError, match="does not fit _EdgeSelI8"): + v.select(ge, thunk) + A = Matrix(i8s, nrows=1, ncols=2) + A[0, 0] = [1, 2, 3] + with pytest.raises(ValueError, match="does not fit _EdgeSelI8"): + A.select(ge, 1.5) + # A bool array converts into any numeric element, as it did, and bools in a + # sequence compare as 0 and 1 (a BOOL[3] UDT cannot be registered). + result = v.apply(binary.plus, right=np.array([True, False, True])).new() + assert result.dtype == i8s + assert result[0].new().value.tolist() == [2, 2, 4] + ones = Vector(i8s, size=2) + ones[0] = [1, 0, 1] + assert (ones == (True, False, True)).new()[0].new().value + assert not (ones != [True, False, True]).new()[0].new().value + # A float too large for a float32 field raises where it must become an element. + f = Vector(f4s, size=2) + f[0] = [1, 2, 3] + with pytest.raises(OverflowError, match="overflows _EdgeSelF4"): + f.ewise_union(f, binary.plus, 1e300, 0) + with pytest.raises(OverflowError, match="overflows _EdgeSelF4"): + f.apply(binary.register_anonymous(lambda x, y: x, "_edge_sel_first", is_udt=True), 1e300) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_exact_literal_of_another_type_is_deprecated(): + """A must-fit value that fits only by value converts, with a DeprecationWarning. + + On main any literal, thunk or default was converted into the UDT, and what + did not fit was lost. Now it must fit by type; one whose values convert + exactly, such as 0 for a bool field, still converts for the deprecation + period, and the warning points at the caller's line. + """ + rec = dtypes.register_anonymous( + np.dtype([("dp_b", np.bool_), ("dp_f", np.float64), ("dp_i", np.int8)]), "_DeprecRec" + ) + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_DeprecI8") + r = Vector(rec, size=2) + r[0] = (True, 1.5, 2) + s = Vector(rec, size=2) + s[1] = (False, 2.5, 3) + expected = r.ewise_union(s, binary.plus, False, False).new() + for default in [0, (0, 0.0, 0), {"dp_b": 0, "dp_f": 0.0, "dp_i": 0}]: + with pytest.warns(DeprecationWarning, match="does not fit _DeprecRec") as record: + result = r.ewise_union(s, binary.plus, default, False).new() + assert result.isequal(expected) + assert {w.filename for w in record} == {__file__} + with pytest.raises(ValueError, match="3 does not fit _DeprecRec"): + r.ewise_union(s, binary.plus, 3, False) + + def ge_first(x, i, j, t): # pragma: no cover (numba) + return x[0] >= t[0] + + ge = IndexUnaryOp.register_anonymous(ge_first, "_deprec_ge", is_udt=True) + v = Vector(i8s, size=3) + v[0] = [1, 2, 3] + v[1] = [2, 3, 4] + A = Matrix(i8s, nrows=1, ncols=2) + A[0, 1] = [2, 3, 4] + for x in [v, A]: + with pytest.warns(DeprecationWarning, match="2.0 does not fit _DeprecI8"): + assert x.select(ge, 2.0).new().isequal(x.select(ge, 2).new()) + for thunk in [1.5, float("nan"), np.int64(300)]: + with pytest.raises(ValueError, match="does not fit _DeprecI8"): + x.select(ge, thunk) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_monoid_literal_is_typed_as_its_binaryop(): + """A literal beside a Monoid is typed as beside its BinaryOp, everywhere. + + ``apply`` already used the BinaryOp, but a Scalar's ``ewise_add`` and + ``ewise_mult`` with a Monoid converted the literal into the UDT (so ``0.5`` + had to fit), where the same call with ``binary.plus`` promoted. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_MonoLitI8") + s = gb.Scalar.from_value([1, 2, 3], dtype=i8s) + for expr in [s.ewise_add(0.5, monoid.plus), s.ewise_add(0.5, binary.plus)]: + result = expr.new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + np.testing.assert_array_equal(result.value, [1.5, 2.5, 3.5]) + assert s.ewise_mult(2, monoid.times).new().dtype == i8s + # ewise_union defaults still have to fit the op's input types. + v = Vector(i8s, size=2) + v[0] = [1, 2, 3] + with pytest.raises(ValueError, match="does not fit _MonoLitI8"): + v.ewise_union(v, monoid.plus, 0.5, 0) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_outer_on_mixed_types(): + """``outer`` builds its semiring from the op, typed on both input types. + + ``get_semiring`` typed the semiring by the op's first input type only, so + ``outer`` on two different UDTs, or a UDT and a built-in dtype, was a bare + GrB_DOMAIN_MISMATCH. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_OuterI8") + f4s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_OuterF4") + v = Vector(i8s, size=2) + v[0] = [1, 2, 3] + u = Vector(f4s, size=2) + u[1] = [0.5, 0.5, 0.5] + for a, b in [(v, u), (u, v)]: + result = a.outer(b).new() + assert result.dtype == f4s + assert result.to_coo()[2].tolist() == [[0.5, 1.0, 1.5]] + result = v.outer(Vector.from_coo([0], [2.0], size=2)).new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + assert result.to_coo()[2].tolist() == [[2.0, 4.0, 6.0]] + + @pytest.mark.skipif("not supports_udfs") def test_udt_jit_kernel_only_for_a_same_type_result(): """A same-type pair gets the arithmetic JIT kernel only when its result is that type too. @@ -4369,11 +5498,15 @@ def test_udt_pair_is_never_unified(): v[0] = (1.0, 2) w = Vector(Y, size=1) w[0] = (1, 2.0) - for expr in [v.ewise_add(w, monoid.plus), v.ewise_add(w), monoid.plus(v | w)]: + for expr in [v.ewise_mult(w, monoid.plus), v.ewise_union(w, monoid.plus, 0, 0)]: result = expr.new() assert result.dtype.np_type == np.dtype([("uni_a", np.float64), ("uni_b", np.float64)]) assert result[0].new().value.tolist() == (2.0, 4.0) assert v.ewise_mult(w, monoid.max).new()[0].new().value.tolist() == (1.0, 2.0) + # ewise_add, the Monoid's default, would need converted copies of the operands. + for expr in [lambda: v.ewise_add(w), lambda: monoid.plus(v | w)]: + with pytest.raises(DomainMismatch, match="ewise_add cannot use binary.plus"): + expr() with pytest.raises(TypeError, match="Monoid inputs must be the same dtype"): monoid.plus[X, Y] diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 024b4117a..228d69b91 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -1104,41 +1104,182 @@ def test_floordiv_udt_jit_matches_python_semantics(): @pytest.mark.skipif("not supports_udfs") -def test_min_max_udt_jit_propagates_nan(): - """``binary.min``/``max`` on a float UDT must propagate NaN like Python/numba. - - Regression: the JIT codegen used to emit ``(a < b ? a : b)``, which - silently swallows NaN to the right-hand side. Python ``min(a, b)`` (and - numba's ``min``) returns ``a`` when neither comparison is true (NaN - involved), so ``min(NaN, 1.0) == NaN`` and ``min(1.0, NaN) == 1.0``. - The fix swaps the ternary to ``(b < a ? b : a)``. +def test_min_max_udt_jit_ignores_nan_and_breaks_zero_ties(): + """The C JIT kernel for ``binary.min`` on a float UDT answers as the cfunc does. + + ``GrB_MIN_FP64`` is C99 ``fmin``, which ignores a NaN operand from either + side, and so does ``binary.min`` on a UDT field. C99 leaves the tie + between ``-0.0`` and ``0.0`` to the C library (glibc's x86-64 ``fmin`` + returns an operand by position), so the kernel spells the comparison out, + as the cfunc does, rather than calling ``fmin``: a tie gives ``-0.0`` for + ``min`` and ``0.0`` for ``max`` on every platform. Earlier spellings + decided NaN by position, in opposite directions: ``(a < b ? a : b)`` + dropped a NaN on the right, and ``(b < a ? b : a)`` one on the left. + + Integer fields keep the plain comparison: they have no NaN or signed zero. """ - if _IS_SSGB7: - pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + if not _has_jit_set: + pytest.skip("jit_c_source introspection requires SuiteSparse:GraphBLAS >= 9") _require_jit_on() # Field names unique to this test; see floordiv test for the cache rationale. udt = dtypes.register_anonymous( - np.dtype([("nan_a", np.float64), ("nan_b", np.float64)]), "_NanJitMM" + np.dtype([("nan_a", np.float64), ("nan_b", np.float32), ("nan_c", np.int32)]), "_NanJitMM" ) + csrc = binary.min[udt].jit_c_source + assert "fmin" not in csrc, csrc + assert "signbit(x->nan_a)" in csrc, csrc + assert "signbit(x->nan_b)" in csrc, csrc + assert "((x->nan_c) < (y->nan_c) ? (x->nan_c) : (y->nan_c))" in csrc, csrc + assert "!signbit(x->nan_a)" in binary.max[udt].jit_c_source + N = 100 v = gb.Vector(udt, N) u = gb.Vector(udt, N) nan = float("nan") for i in range(N): # field nan_a: NaN on the left; field nan_b: NaN on the right at odd indices. - v[i] = (nan, 2.0 + i) - u[i] = (1.0 + i, nan if i % 2 else 3.0 + i) + v[i] = (nan, 2.0 + i, i) + u[i] = (1.0 + i, nan if i % 2 else 3.0 + i, 2 * i) w = v.ewise_mult(u, binary.min).new() - assert np.isnan(w[0].new().value[0]) # min(NaN, 1.0) -> NaN - assert w[1].new().value[1] == 3.0 # min(3.0, NaN) -> 3.0 (NaN swallowed) + assert w[0].new().value[0] == 1.0 # min(NaN, 1.0) -> 1.0 + assert w[1].new().value[1] == 3.0 # min(3.0, NaN) -> 3.0 assert w[2].new().value[1] == 4.0 # min(4.0, 5.0) -> 4.0 (normal case) + assert w[3].new().value[2] == 3 # integer field is unaffected w = v.ewise_mult(u, binary.max).new() - assert np.isnan(w[0].new().value[0]) # max(NaN, 1.0) -> NaN - assert w[1].new().value[1] == 3.0 # max(3.0, NaN) -> 3.0 (NaN swallowed) + assert w[0].new().value[0] == 1.0 # max(NaN, 1.0) -> 1.0 + assert w[1].new().value[1] == 3.0 # max(3.0, NaN) -> 3.0 assert w[2].new().value[1] == 5.0 # max(4.0, 5.0) -> 5.0 (normal case) + assert w[3].new().value[2] == 6 # integer field is unaffected + + # Both operands NaN is the one case where a NaN survives, for min and max + # alike, and it is the only case ``fmin`` has no non-NaN answer for. + nan_only = gb.Vector(udt, N) + for i in range(N): + nan_only[i] = (nan, np.float32(i), i) + w = nan_only.ewise_mult(nan_only, binary.min).new() + assert np.isnan(w[0].new().value[0]) + + # Signed-zero ties, in both orders and both fields, on the C kernel and the cfunc. + left = gb.Vector(udt, N) + right = gb.Vector(udt, N) + for i in range(N): + a, b = (-0.0, 0.0) if i % 2 else (0.0, -0.0) + left[i] = (a, a, i) + right[i] = (b, b, i) + for mode in ["on", "off"]: + with _jit_mode(mode): + lo = left.ewise_mult(right, binary.min).new() + hi = left.ewise_mult(right, binary.max).new() + for i in range(2): + assert np.signbit(lo[i].new().value[0]), mode + assert np.signbit(lo[i].new().value[1]), mode + assert not np.signbit(hi[i].new().value[0]), mode + assert not np.signbit(hi[i].new().value[1]), mode + + +@pytest.mark.skipif("not supports_udfs") +def test_complex_truediv_udt_jit_matches_cfunc_at_extremes(): + """Complex division in the C JIT kernel is the cfunc's, bit for bit, at any magnitude. + + SuiteSparse compiles C JIT kernels with ``-fcx-limited-range`` under GCC, + which turns C's ``_Complex`` division into the naive formula: complex64 + ``1e20 / 1e20`` came out ``nan`` and ``(1+1j) / (1e160+1e160j)`` came out + ``0`` on Linux. The kernel divides with real arithmetic instead, by the + algorithm Numba (and CPython) use, so no flag changes it. + """ + if not _has_jit_set: + pytest.skip("jit_c_source introspection requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + if numba is None: + pytest.skip("numba required for the cfunc baseline") + + udt = dtypes.register_anonymous( + np.dtype([("cxd_s", np.complex64), ("cxd_d", np.complex128)], align=True), "_CxDivJit" + ) + csrc = binary.truediv[udt].jit_c_source + assert "crealf" in csrc, csrc + assert "(x->cxd_d) / (y->cxd_d)" not in csrc, csrc + # (complex64 pair, complex128 pair); the first four are finite and the C + # library's division got them wrong under the limited-range flag. + pairs = [ + ((1e20, 1e20), (1e200, 1e200)), + ((1e-25 + 1e-25j, 1e-25), (1e200 + 1e200j, 1e200 + 1e200j)), + ((1 + 1j, 1e30 + 1e30j), (1 + 1j, 1e160 + 1e160j)), + ((3e30 + 4e30j, 1e30 - 2e30j), (1e-170, 1e-170 + 1e-170j)), + ((0j, -2 + 0.5j), (0j, -2 + 0.5j)), + ((-2 + 0.5j, -2 + 0.5j), (-2 + 0.5j, -2 + 0.5j)), + ((3 + 4j, complex(0.0, -0.0)), (3 + 4j, complex(0.0, -0.0))), + ((complex(np.inf, np.nan), 1 + 1j), (complex(np.inf, np.nan), 1 + 1j)), + ] + N = 64 + v = gb.Vector(udt, N) + u = gb.Vector(udt, N) + for i in range(N): + (a_s, b_s), (a_d, b_d) = pairs[i % len(pairs)] + v[i] = (a_s, a_d) + u[i] = (b_s, b_d) + with _jit_mode("on"): + jit = v.ewise_mult(u, binary.truediv).new() + with _jit_mode("off"): + cfunc = v.ewise_mult(u, binary.truediv).new() + assert jit.to_dense().tobytes() == cfunc.to_dense().tobytes() + result = jit.to_dense() + for i, ((a_s, b_s), (a_d, b_d)) in enumerate(pairs[:4]): + assert np.isclose(result[i]["cxd_s"], np.complex64(a_s) / np.complex64(b_s), rtol=1e-6) + assert np.isclose(result[i]["cxd_d"], np.complex128(a_d) / np.complex128(b_d), rtol=1e-15) + + +@pytest.mark.skipif("not supports_udfs") +def test_complex_times_udt_jit_matches_cfunc_with_infinities(): + """Complex multiplication in the C JIT kernel is the cfunc's, bit for bit. + + Numba and numpy multiply by the textbook formula; clang's ``_Complex`` + multiply adds C99 Annex G's recovery, so ``1j * (inf+infj)`` came out + ``-inf+infj`` from a C kernel on macOS and ``nan+nanj`` from the cfunc + (GCC's ``-fcx-limited-range`` drops the recovery, so Linux agreed). The + kernel multiplies with real arithmetic instead, for record fields and + array elements alike. + """ + if not _has_jit_set: + pytest.skip("jit_c_source introspection requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + if numba is None: + pytest.skip("numba required for the cfunc baseline") + + rec = dtypes.register_anonymous( + np.dtype([("cxt_s", np.complex64), ("cxt_d", np.complex128)], align=True), "_CxTimesJit" + ) + arr = dtypes.register_anonymous(np.dtype((np.complex128, (2,))), "_CxTimesJitArr") + assert "(x->cxt_d) * (y->cxt_d)" not in binary.times[rec].jit_c_source + inf, nan = np.inf, np.nan + pairs = [ + (1j, complex(inf, inf)), + (-2.5j, complex(inf, -inf)), + (complex(inf, 0.0), complex(0.0, 1.0)), + (complex(nan, inf), 2 + 1j), + (complex(1e30, 1e30), complex(1e30, -1e30)), + (3 + 4j, -1.5 + 0.25j), + (complex(-0.0, 0.0), complex(0.0, -0.0)), + ] + N = 3 * len(pairs) + for udt, element in [(rec, lambda z: (z, z)), (arr, lambda z: [z, -z])]: + v = gb.Vector(udt, N) + u = gb.Vector(udt, N) + for i in range(N): + a, b = pairs[i % len(pairs)] + v[i] = element(a) + u[i] = element(b) + with _jit_mode("on"): + jit = v.ewise_mult(u, binary.times).new() + with _jit_mode("off"): + cfunc = v.ewise_mult(u, binary.times).new() + assert jit.to_dense().tobytes() == cfunc.to_dense().tobytes() + result = jit[0].new().value # 1j * (inf+infj), numpy's nan+nanj + assert np.isnan(result[0].real) + assert np.isnan(result[0].imag) @pytest.mark.skipif("not supports_udfs") diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 52c382742..24c1889f9 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -2097,11 +2097,12 @@ def test_udt(): assert s == t t = v.reduce(monoid.any, allow_empty=False).new() assert s == t - v << binary.first(1, v) + # A literal fills every field by type; an int does not fit the bool field, True does. + v << binary.first(True, v) assert v.isequal(ones) v << binary.first(s, v) assert v.isequal(zeros) - v << binary.second(v, (1, 1)) + v << binary.second(v, (True, 1)) assert v.isequal(ones) v << binary.second(v, s) assert v.isequal(zeros) From 66effa146fe731eaba6cd4b0995aad5ca0e149fe Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 8 Oct 2026 09:36:25 -0500 Subject: [PATCH 06/54] Type Python numbers beside built-in dtypes as numpy 2 does (NEP 50) A Python number beside a Vector, Matrix or Scalar of a built-in dtype was typed through numpy, with a 0-d array of value 0 standing in for it (dtypes.unify). numpy 1 judged that array by its value, so the collection's dtype won and a large literal wrapped (int8_vec * 1000 multiplied by -24, int8_vec < 300 compared with 44); numpy 2 (NEP 50) made 0-d arrays strong, so every int was INT64 and every float FP64 (int8_vec + 1 was INT64, uint64_vec + 1 FP64, losing precision). The same script gave other dtypes, and on numpy 1 wrong answers, depending on the installed numpy. A Python bool, int, float or complex operand is now weak, by the table UDT leaves already use: it takes the other operand's dtype when that dtype's kind holds it, else INT64, FP64 or FC64, on numpy 1 and 2 alike. So int8_vec + 1 is INT8 and wraps as numpy does, int8_vec * 0.5 is FP64, fp32_vec * 0.5 is FP32, int8_vec + 300 raises OverflowError, and fp32_vec + 1e300 (or + 2**200) is inf with a RuntimeWarning, as in numpy 2. A strong operand keeps the old widening: int8_vec + np.int64(1). This covers apply (infix ops, binary.plus(v, 1), Monoids), Scalar ewise ops, ewise_union defaults, value thunks of select and IndexUnaryOps, and UDFs. A comparison keeps no result type, so an int out of range compares exactly (int8_vec < 300 is all True), and so does an op whose result is another type anyway: int8_vec / 300 is FP64, as numpy divides integers in float64, for UDT fields too. A to_dense fill value that does not fit widens, an int or a float (fp32_vec.to_dense(fill_value=1e300) is float64), rather than raising or becoming inf. numpy scalars, Scalars and typed ops stay strong, and a thunk that is not compared with the values, such as a row index, keeps its own type. unify no longer stands a scalar in as a value-0 array. eq and ne beside a UDT now type a Python number or sequence the same way, leaf by leaf, with the same exact fallback: fp32_udt == 0.1 compares in float32, as fp32_vec == 0.1 and numpy do. --- docs/user_guide/types.rst | 25 +++++ docs/user_guide/udt.rst | 10 +- graphblas/core/dtypes.py | 13 +-- graphblas/core/matrix.py | 14 ++- graphblas/core/operator/udt_utils.py | 59 +++++++++--- graphblas/core/scalar.py | 136 +++++++++++++++++++++++---- graphblas/core/vector.py | 14 ++- graphblas/tests/test_op.py | 33 +++++++ graphblas/tests/test_vector.py | 88 +++++++++++++++++ 9 files changed, 346 insertions(+), 46 deletions(-) diff --git a/docs/user_guide/types.rst b/docs/user_guide/types.rst index aaeba1c6b..f6e6c6ea3 100644 --- a/docs/user_guide/types.rst +++ b/docs/user_guide/types.rst @@ -30,6 +30,31 @@ When a data type is needed in an API call, a string or numpy or numba dtype may instead of the actual data type object. Additionally, the Python builtin ``bool``, ``int``, and ``float`` may be used. ``int`` indicates INT64 and ``float`` indicates FP64. +Python numbers in operations +---------------------------- + +A Python ``bool``, ``int``, ``float`` or ``complex`` used as an operand, such as +the ``1`` in ``v + 1`` or ``v.apply(binary.times, right=0.5)``, is *weak*, as in +numpy 2 (NEP 50). It takes the other operand's data type when that type's kind +can hold it, and its own default type (BOOL, INT64, FP64 or FC64) when not. So +with an INT8 vector ``v``, ``v + 1`` is INT8 and ``v * 0.5`` is FP64, while +``v + 300`` raises ``OverflowError`` rather than wrapping, and ``v * 2`` wraps +in INT8 as it does in numpy. An FP32 vector times ``0.5`` is FP32, plus ``1j`` +is FC32, and plus ``1e300`` is infinity, with a ``RuntimeWarning`` as in numpy. numpy scalars, 0-d arrays and ``Scalar`` objects keep their own +type: ``v + np.int64(1)`` is INT64. python-graphblas applies this rule itself, +so it does not depend on the installed version of numpy. User-defined types +follow the same rule, field by field (see :doc:`udt`). + +A comparison has no result type to keep, so an integer outside the operand's +range compares exactly, as in numpy: ``v < 300`` is True for every element. +Nor does ``v / 300`` raise, since its result is FP64 whatever the integer. An +FP32 vector ``== 0.1`` compares in FP32. The defaults of ``ewise_union`` +follow the rule, and so does the ``fill_value`` of ``to_dense``, except that a +value that does not fit widens the result instead of raising. A literal given +to a typed operator, such as ``binary.plus["INT64"]``, takes that operator's +type, and a ``select`` or ``IndexUnaryOp`` thunk that is not compared with the +values, such as a row index, keeps its own type. + User-defined Types ------------------ diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 61eb97900..87d0bc6ba 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -171,12 +171,16 @@ dtype, so it never narrows. Below, ``v`` is a Vector of an ``INT8[3]`` UDT. - Weak, field by field: a field's dtype when its kind holds the number, else INT64, FP64 or FC64 (numpy 2's NEP 50 rule, on numpy 1 too, and extended to sequences). An int out of an integer field's range raises - ``OverflowError``. + ``OverflowError``, except for ``truediv``, whose fields are floats + whatever the int (``v / 300`` is ``FP64[3]``). - ``v + 1`` is ``INT8[3]``, ``v * 2.5`` and ``v + (0.5, 1, 2)`` are ``FP64[3]``, ``v + 300`` raises * - ``eq`` and ``ne`` - - Compare the values as given and return ``BOOL``. - - ``v == 300`` is ``False`` + - Type a literal as above, as built-in comparisons do, except that an int + out of a field's range takes a type that holds it, so the comparison is + exact. Return ``BOOL``. + - ``v == 300`` is ``False``; an ``FP32[3]`` UDT ``== 0.1`` compares in + FP32 * - A literal that must become an element: beside a user-defined op, an ``IndexUnaryOp`` thunk, an ``ewise_union`` default - Its type by the rules above must fit the UDT: the UDT itself, or a diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 93b3d4d8a..777fbc2a6 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -3,7 +3,7 @@ from ast import literal_eval import numpy as np -from numpy import promote_types, result_type +from numpy import promote_types from .. import backend, dtypes from ..core import NULL, _has_numba, ffi, lib @@ -566,14 +566,15 @@ def unify(type1, type2, *, is_left_scalar=False, is_right_scalar=False): unify(INT8, UINT16) -> INT32 unify(BOOL, UINT16) -> UINT16 unify(FP32, INT32) -> FP64 + + ``is_left_scalar`` and ``is_right_scalar`` no longer change the result. A + Python number is typed before it gets here (``scalar._literal_dtype``), and + a typed scalar promotes as any operand does, as numpy 2 promotes it. They + used to stand the scalar in as a 0-d array of value 0, which numpy 1 typed + by value, so ``int8_vec + 300`` wrapped there. """ if type1 is type2: return type1 - if is_left_scalar: - if not is_right_scalar: - return lookup_dtype(result_type(np.array(0, type1.np_type), type2.np_type)) - elif is_right_scalar: - return lookup_dtype(result_type(type1.np_type, np.array(0, type2.np_type))) return lookup_dtype(promote_types(type1.np_type, type2.np_type)) diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 0a0992b7f..a3477459e 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -33,6 +33,7 @@ _ewise_add_needs_cast, _literal_dtype, _scalar_index, + _weak_builtin_literal_dtype, ) from .utils import ( _CArray, @@ -1519,7 +1520,8 @@ def to_dense(self, fill_value=None, dtype=None, **opts): # dtype of fill_value can upcast the dtype if type(fill_value) is not Scalar: try: - fill_value = Scalar.from_value(fill_value, is_cscalar=None, name="") + dtype = _weak_builtin_literal_dtype(self.dtype, fill_value) + fill_value = Scalar.from_value(fill_value, dtype, is_cscalar=None, name="") except TypeError: fill_value = self._expect_type( fill_value, @@ -2130,7 +2132,10 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): left_dtype = temp_op.type dtype = left_dtype if left_dtype._is_udt else None if type(left_default) is not Scalar: - _check_literal_fits(dtype, left_default) + if dtype is not None: + _check_literal_fits(dtype, left_default) + else: + dtype = _literal_dtype(left_dtype, left_default, op) try: left = Scalar.from_value( left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -2150,7 +2155,10 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): right_dtype = temp_op.type2 dtype = right_dtype if right_dtype._is_udt else None if type(right_default) is not Scalar: - _check_literal_fits(dtype, right_default) + if dtype is not None: + _check_literal_fits(dtype, right_default) + else: + dtype = _literal_dtype(right_dtype, right_default, op) try: right = Scalar.from_value( right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index 87805ecbf..abcbab340 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -579,6 +579,12 @@ def _promoted_record(records, result_leaves): } +@lru_cache +def _int_bounds(elem): + info = np.iinfo(elem) + return info.min, info.max + + def _weak_literal_element(value, kind, elem): """Return the dtype a Python literal of ``kind`` takes beside an element of ``elem``. @@ -590,13 +596,42 @@ def _weak_literal_element(value, kind, elem): """ if elem.kind in _WEAK_LITERAL_FITS[kind]: if kind == "i" and elem.kind in "iu": - info = np.iinfo(elem) - if not info.min <= value <= info.max: + lo, hi = _int_bounds(elem) + if not lo <= value <= hi: raise OverflowError(f"Python integer {value} out of bounds for {elem}") return elem if kind == "c" and elem.kind == "f": return np.result_type(elem, np.complex64) - return _WEAK_LITERAL_DEFAULT[kind] + rv = _WEAK_LITERAL_DEFAULT[kind] + if kind == "i": # beside a bool, an int is INT64, and must fit it + lo, hi = _int_bounds(rv) + if not lo <= value <= hi: + raise OverflowError(f"Python integer {value} out of bounds for {rv}") + return rv + + +def _weak_literal_leaf(values, kind, elem, *, exact=False): + """Return the one dtype Python numbers ``values`` of widest kind ``kind`` take beside ``elem``. + + Each is typed by :func:`_weak_literal_element`, so an int out of an integer + ``elem``'s range raises ``OverflowError``. With ``exact``, for a comparison, + which keeps no result type, such an int takes int64, uint64 or float64, + whichever holds the values, so it compares exactly, as numpy compares it: + ``int8_vec < 300`` is all True. Built-in dtypes and UDT leaves both use this. + """ + if exact and kind == "i" and elem.kind in "biu": + # The integer type an int would take: ``elem``, or INT64 beside a bool. + lo, hi = _int_bounds(elem if elem.kind != "b" else _WEAK_LITERAL_DEFAULT["i"]) + for x in values: + if not lo <= x <= hi: + if all(-(2**63) <= x < 2**63 for x in values): + return np.dtype(np.int64) + if all(0 <= x < 2**64 for x in values): + return np.dtype(np.uint64) + return np.dtype(np.float64) + for x in values: + elem_type = _weak_literal_element(x, kind, elem) + return elem_type _LITERAL_KINDS = {bool: "b", int: "i", float: "f", complex: "c"} @@ -642,7 +677,7 @@ def _literal_leaf_values(value, shape): return values if all(type(x) in _LITERAL_KINDS for x in values) else None -def _weak_literal_udt(dtype, value, *, weak=True): +def _weak_literal_udt(dtype, value, *, exact=False): """Return the UDT a Python number, or a tuple or list of them, becomes beside UDT ``dtype``. The literal is read as an element of ``dtype``: a number stands in every @@ -654,9 +689,10 @@ def _weak_literal_udt(dtype, value, *, weak=True): ``int8_udt + (1, 2, 3)`` stay ``int8_udt``, while ``int8_udt * 2.5`` and ``int8_udt + (0.5, 1.5, 2.5)`` give float64 elements. A leaf that is not a number (bytes, datetime) takes the literal's default dtype and fails to - pair. With ``weak=False`` every leaf takes the numbers' default dtype - (INT64, FP64, FC64), as ``eq`` and ``ne`` take a number as it is. A sequence - not laid out like an element gives ``None``. + pair. With ``exact``, for ``eq`` and ``ne``, an int out of a leaf's range + takes a type that holds it instead of raising (:func:`_weak_literal_leaf`), + so ``int8_udt == 300`` is False. A sequence not laid out like an element + gives ``None``. """ from ..dtypes import lookup_dtype @@ -666,15 +702,10 @@ def leaf_dtype(leaf, part): if values is None: raise _LiteralDoesNotFit kind = max((_LITERAL_KINDS[type(x)] for x in values), key="bifc".index) - if weak and base.kind in "biufc": - # One dtype for the leaf; an int out of the element's range raises. - for x in values: - elem = _weak_literal_element(x, kind, base) + if base.kind in "biufc": + elem = _weak_literal_leaf(values, kind, base, exact=exact) else: elem = _WEAK_LITERAL_DEFAULT[kind] - if elem.kind == "b" and shape and base.kind in "biufc": - # A bool array cannot be a UDT, and a bool compares as 0 or 1. - elem = base return np.dtype((elem, shape)) if shape else elem def record_leaves(np_type, part): diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 1e0cc194d..20bdff5a9 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -10,7 +10,8 @@ from .base import BaseExpression, BaseType, call from .dtypes import _view_if_same_layout from .expr import AmbiguousAssignOrExtract -from .operator import get_typed_op +from .operator import IndexUnaryOp, SelectOp, TypedOpBase, binary_from_string, get_typed_op +from .operator.udt_utils import _LITERAL_KINDS, _weak_literal_leaf from .utils import _Pointer, output_type, wrapdoc if _supports_udfs: @@ -859,7 +860,10 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): left_dtype = temp_op.type dtype = left_dtype if left_dtype._is_udt else None if type(left_default) is not Scalar: - _check_literal_fits(dtype, left_default) + if dtype is not None: + _check_literal_fits(dtype, left_default) + else: + dtype = _literal_dtype(left_dtype, left_default, op) try: left = Scalar.from_value( left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -879,7 +883,10 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): right_dtype = temp_op.type2 dtype = right_dtype if right_dtype._is_udt else None if type(right_default) is not Scalar: - _check_literal_fits(dtype, right_default) + if dtype is not None: + _check_literal_fits(dtype, right_default) + else: + dtype = _literal_dtype(right_dtype, right_default, op) try: right = Scalar.from_value( right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -1172,8 +1179,8 @@ def _literal_dtype(dtype, value, op): """Return the dtype to build ``value`` with, as the other operand of ``op`` beside ``dtype``. ``None`` means the value's own dtype. This is where python-graphblas - decides how a literal operand is typed beside a UDT; built-in dtypes - always get ``None`` and are typed by ``unify``. + decides how a literal operand is typed; beside a built-in dtype, see + :func:`_weak_builtin_literal_dtype`. Beside a UDT, a literal becomes an element of that UDT, as it must for a user's op, when its type allows it (:func:`_check_literal_fits`). With a @@ -1192,18 +1199,21 @@ def _literal_dtype(dtype, value, op): weak in the same way, leaf by leaf: ``int8_udt + (1, 2, 3)`` stays ``int8_udt`` and ``+ (0.5, 1.5, 2.5)`` has float64 elements. - ``eq`` and ``ne`` take a number as it is, strong, since a comparison - needs no result type and ``int8_udt == 300`` should be False, not an error; - a sequence likewise, so ``int8_udt == (0.5, 1, 2)`` compares the values. + ``eq`` and ``ne`` type a literal the same way, as numpy compares, but a + comparison keeps no result type, so an int out of a leaf's range takes a + type that holds it: ``int8_udt == 300`` is False, not an error, and + ``fp32_udt == 0.1`` compares in float32, as ``fp32_vec == 0.1`` does. """ - if not dtype._is_udt or getattr(op, "is_positional", False): + if getattr(op, "is_positional", False): return None + if not dtype._is_udt: + return _weak_builtin_literal_dtype(dtype, value, op) if getattr(op, "_udt_types", True) is None: # An op that does not take UDTs at all (``lt``) reports that itself. return dtype from .operator.binary import BinaryOp from .operator.monoid import Monoid - from .operator.udt_utils import _LITERAL_KINDS, BUILTIN_UDT_BINARY_OPS, _as_python_number + from .operator.udt_utils import BUILTIN_UDT_BINARY_OPS if isinstance(op, Monoid) and not op._anonymous: # Element-wise, a Monoid is its BinaryOp (``get_typed_op`` types it so). @@ -1211,22 +1221,114 @@ def _literal_dtype(dtype, value, op): if type(op) is BinaryOp and not op._anonymous: is_comparison = op.name in ("eq", "ne") if is_comparison or op.name in BUILTIN_UDT_BINARY_OPS: - if is_comparison and type(_as_python_number(value)) in _LITERAL_KINDS: - return None - literal_type = _literal_type(dtype, value, strong_numbers=is_comparison) + # truediv gives float leaves whatever the int, so, as a comparison, it + # types an int out of an integer leaf's range exactly instead of raising. + exact = is_comparison or op.name == "truediv" + literal_type = _literal_type(dtype, value, exact=exact) if literal_type is not None: return literal_type _check_literal_fits(dtype, value) return dtype -def _literal_type(dtype, value, *, strong_numbers=False): +# 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)]) + + +# The types a Python float or complex can overflow. +_SINGLE_PRECISION = {np.dtype(np.float32), np.dtype(np.complex64)} + + +def _weak_builtin_literal_dtype(dtype, value, op=None): + """Return the dtype of Python number ``value`` beside built-in ``dtype``, or ``None``. + + A Python ``bool``, ``int``, ``float`` or ``complex`` is weak, as numpy 2 + types it (NEP 50), by the table UDT leaves use + (``udt_utils._weak_literal_leaf``): ``int8_vec + 1`` is INT8, + ``fp32_vec * 0.5`` is FP32, ``int8_vec * 0.5`` is FP64, and + ``int8_vec + 300`` raises ``OverflowError``. python-graphblas applies the + table itself, so numpy 1 gives the same answers. ``None`` keeps the value's + own dtype (strong): for anything but a Python number, for a typed op, which + fixes the operand types itself, and for the thunk of an index op that does + not compare it with the values, such as a row index. + + A comparison keeps no result type, so an int beyond ``dtype``'s range + takes a type that holds it and compares exactly, as numpy compares it: + ``int8_vec < 300`` is all True. ``op`` of ``None``, for a fill value, does + the same, so the value is not lost. + """ + kind = _LITERAL_KINDS.get(type(value)) + if kind is None: + return None + if isinstance(op, str): + try: + op = binary_from_string(op) + except Exception: + return None + exact = op is None + if op is not None: + if isinstance(op, TypedOpBase): + return None + exact = getattr(op, "name", None) in _COMPARISONS and not op._anonymous + if isinstance(op, (IndexUnaryOp, SelectOp)) and not exact: + return None + try: + np_type = _weak_literal_leaf([value], kind, dtype.np_type, exact=exact) + except OverflowError: + if exact or not _gives_another_type(op, dtype): + raise + # An op whose result is not the operand's type has no narrow result to + # keep, so an int out of range is typed exactly, as for a comparison: + # int8_vec / 300 is FP64, as numpy divides integers in float64. + np_type = _weak_literal_leaf([value], kind, dtype.np_type, exact=True) + if np_type == dtype.np_type: + rv = dtype + else: + try: + rv = lookup_dtype(np_type) + except ValueError: # A complex literal where the backend has no complex types + return None + if kind != "b" and (op is None or not exact) and rv.np_type in _SINGLE_PRECISION: + import struct + + parts = (value.real, value.imag) if kind == "c" else (float(value),) + try: + for part in parts: + struct.pack(" 0 +def test_python_number_is_weak_beside_builtin_dtypes(): + """A Python number takes a built-in collection's dtype when that kind holds it (NEP 50). + + How a literal was typed used to depend on the installed numpy: on numpy 2 a + Python int was INT64 and a float FP64 beside anything, and on numpy 1 the + collection's dtype won by value, so ``int8_vec + 300`` wrapped and + ``int8_vec < 300`` compared with 44. python-graphblas now applies numpy 2's + rule itself, the table UDT elements use, on numpy 1 and 2 alike. numpy + scalars and Scalars keep their dtype. + """ + i8 = Vector.from_coo([0, 1], [100, -100], dtype=dtypes.INT8) + u8 = Vector.from_coo([0, 1], [0, 200], dtype=dtypes.UINT8) + f32 = Vector.from_coo([0, 1], [0.1, 1.5], dtype=dtypes.FP32) + u64 = Vector.from_coo([0, 1], [0, 2**63], dtype=dtypes.UINT64) + bools = Vector.from_coo([0, 1], [True, False]) + for expr, dtype, values in [ + (i8 + 1, dtypes.INT8, [101, -99]), + (1 + i8, dtypes.INT8, [101, -99]), + (i8 + True, dtypes.INT8, [101, -99]), + (binary.plus(i8, 1), dtypes.INT8, [101, -99]), + (i8.apply(binary.minus, left=1), dtypes.INT8, [-99, 101]), + (i8.apply(monoid.max, right=1), dtypes.INT8, [100, 1]), + (i8.apply("+", right=1), dtypes.INT8, [101, -99]), + (1 - u8, dtypes.UINT8, [1, 57]), # wraps, as numpy 2 does + (i8 * 0.5, dtypes.FP64, [50, -50]), + (u64 + 1, dtypes.UINT64, [1, 2**63 + 1]), + (bools + 1, dtypes.INT64, [2, 1]), + (f32 * 2, dtypes.FP32, [np.float32(0.2), 3]), + # Strong: numpy scalars, Scalars, and the literal of a typed op + (i8 + np.int64(1), dtypes.INT64, [101, -99]), + (i8 + Scalar.from_value(1), dtypes.INT64, [101, -99]), + (binary.plus["INT64"](i8, 1), dtypes.INT64, [101, -99]), + (binary.plus["FP64"](i8, 300), dtypes.FP64, [400, 200]), + ]: + assert expr.dtype == dtype + assert expr.new().to_coo()[1].tolist() == values + assert (f32 * 0.5).dtype == dtypes.FP32 + if dtypes._supports_complex: + assert (f32 + 1j).dtype == dtypes.FC32 + s8 = Scalar.from_value(100, dtype=dtypes.INT8) + assert (s8 + 1).new().dtype == (1 + s8).new().dtype == dtypes.INT8 + assert (Scalar.from_value(0.5, dtype=dtypes.FP32) * 0.5).new().dtype == dtypes.FP32 + A = Matrix.from_coo([0], [0], [100], dtype=dtypes.INT8) + assert (A * 2).dtype == dtypes.INT8 + for expr in [lambda: i8 + 300, lambda: u8 + -1, lambda: i8 // 300, lambda: s8 * 300]: + with pytest.raises(OverflowError, match="out of bounds"): + expr() + # truediv gives FP64 whatever the int, so an int out of range does not raise: + # numpy divides integers in float64 too. + assert (i8 / 300).dtype == (300 / i8).dtype == dtypes.FP64 + assert (i8 / 300).new().to_coo()[1].tolist() == [100 / 300, -100 / 300] + # A float too large for FP32 becomes infinity with a warning, as in numpy 2. + with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: + result = (f32 + 1e300).new() + assert result.dtype == dtypes.FP32 + assert np.isinf(result.to_coo()[1]).all() + assert {w.filename for w in record} == {__file__} + with pytest.warns(RuntimeWarning, match="overflow encountered in cast"): + assert np.isinf((f32 + 2**200).new().to_coo()[1]).all() # an int too + # A comparison keeps no result type, so an int out of range compares exactly. + assert (i8 < 300).new().to_coo()[1].tolist() == [True, True] + assert (i8 == 300).new().to_coo()[1].tolist() == [False, False] + assert (u8 > -1).new().to_coo()[1].tolist() == [True, True] + assert i8.apply(indexunary.valuelt, 300).new().to_coo()[1].tolist() == [True, True] + assert (f32 == 0.1).new().to_coo()[1].tolist() == [True, False] # compared as float32 + # ewise_union defaults and fill values follow the rule too. + w8 = Vector.from_coo([1], [1], dtype=dtypes.INT8, size=2) + assert i8.ewise_union(w8, binary.plus, 0, 0).dtype == dtypes.INT8 + assert i8.ewise_union(w8, binary.plus, 0.5, 0).dtype == dtypes.FP64 + sparse = Vector.from_coo([0], [5], dtype=dtypes.INT8, size=2) + assert sparse.to_dense(fill_value=0).dtype == np.int8 + assert sparse.to_dense(fill_value=300).tolist() == [5, 300] # widened, not lost + sparse32 = Vector.from_coo([0], [5], dtype=dtypes.FP32, size=2) + assert sparse32.to_dense(fill_value=0.5).dtype == np.float32 + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert sparse32.to_dense(fill_value=1e300).tolist() == [5, 1e300] # widened + assert sparse32.to_dense(fill_value=2**200).tolist() == [5, 2.0**200] + # A thunk that is not compared with the values keeps its own type. + assert i8.select(select.rowle, 2**40).new().nvals == 2 + if supports_udfs: + above = indexunary.register_anonymous(lambda x, i, j, thunk: x > thunk) + assert i8.apply(above, 300).new().to_coo()[1].tolist() == [False, False] + # Typed scalars promote as numpy 2 promotes them, on numpy 1 too. + assert dtypes.unify(dtypes.INT8, dtypes.INT64, is_right_scalar=True) == dtypes.INT64 + + def test_inner(v): R = Matrix(v.dtype, nrows=1, ncols=v.size) # row vector C = Matrix(v.dtype, nrows=v.size, ncols=1) # column vector From 8e208816ff4202b76519e51532b34d2969243706 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 8 Oct 2026 12:53:58 -0500 Subject: [PATCH 07/54] Compare select value thunks exactly; isequal compares arrays as np.array_equal SuiteSparse:GraphBLAS's select with a value op (valueeq, valuelt, ...) typed for another type than the input's reads the thunk's bytes as the input's type (8.0.2 through 10.6.0; apply with the same op is right; fix proposed upstream in DrTimothyAldenDavis/GraphBLAS#463). On an INT8 vector, select(valueeq, 1.0) kept nothing and select(valuelt, 300) compared with 44, on main too. select now compares in the input's own type with a thunk that keeps the same entries (x < 2.5 is x < 3 for integers). isequal on array UDTs compares as np.array_equal: elements must have the same shape, apart from leading axes of length 1, so an FP64[3] of [1, 1, 1] is no longer isequal to an FP64[1] of [1]. == still broadcasts, and a Scalar's ==, which returns a bool, now computes binary.eq, as == on a Vector does (it called isequal). A literal beside a UDT Scalar is compared as == types it, so [0, 0, 0] == 0.5 is False (0.5 was truncated to 0), and Scalar(FP32 0.1) == 0.1 compares in FP32, as numpy does. A float too large for FP32 in a comparison (fp32_vec == 1e300) compares with infinity, as numpy 2 does, and now warns at the caller's line, as arithmetic does. A Python int from 2**63 to 2**64 - 1 now goes into a UINT64 Scalar, Vector or Matrix element (Scalar.from_value(2**63, "UINT64") and uint64_vec[0] = 2**63 raised OverflowError, on main too). --- docs/user_guide/types.rst | 3 +- docs/user_guide/udt.rst | 8 +- graphblas/core/base.py | 4 +- graphblas/core/matrix.py | 18 ++++- graphblas/core/operator/select.py | 118 ++++++++++++++++++++++++++++++ graphblas/core/scalar.py | 118 ++++++++++++++++++++++++++++-- graphblas/core/vector.py | 16 +++- graphblas/tests/test_op.py | 60 +++++++++++++++ graphblas/tests/test_scalar.py | 47 ++++++++++++ graphblas/tests/test_vector.py | 53 ++++++++++++++ 10 files changed, 427 insertions(+), 18 deletions(-) diff --git a/docs/user_guide/types.rst b/docs/user_guide/types.rst index f6e6c6ea3..e6fd91756 100644 --- a/docs/user_guide/types.rst +++ b/docs/user_guide/types.rst @@ -48,7 +48,8 @@ follow the same rule, field by field (see :doc:`udt`). A comparison has no result type to keep, so an integer outside the operand's range compares exactly, as in numpy: ``v < 300`` is True for every element. Nor does ``v / 300`` raise, since its result is FP64 whatever the integer. An -FP32 vector ``== 0.1`` compares in FP32. The defaults of ``ewise_union`` +FP32 vector ``== 0.1`` compares in FP32, and ``== 1e300`` compares with +infinity, with a ``RuntimeWarning`` as in numpy. The defaults of ``ewise_union`` follow the rule, and so does the ``fill_value`` of ``to_dense``, except that a value that does not fit widens the result instead of raising. A literal given to a typed operator, such as ``binary.plus["INT64"]``, takes that operator's diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 87d0bc6ba..e7aab686b 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -178,9 +178,13 @@ dtype, so it never narrows. Below, ``v`` is a Vector of an ``INT8[3]`` UDT. * - ``eq`` and ``ne`` - Type a literal as above, as built-in comparisons do, except that an int out of a field's range takes a type that holds it, so the comparison is - exact. Return ``BOOL``. + exact. Return ``BOOL``. Array elements broadcast, as numpy's ``==`` + does, and so does ``==`` on a Scalar, which returns a bool. ``isequal`` + compares as ``np.array_equal``: elements must have the same shape, + apart from leading axes of length 1. - ``v == 300`` is ``False``; an ``FP32[3]`` UDT ``== 0.1`` compares in - FP32 + FP32; a Scalar ``s`` of ``[1, 1, 1]`` has ``s == 1`` but not + ``s.isequal(1)`` * - A literal that must become an element: beside a user-defined op, an ``IndexUnaryOp`` thunk, an ``ewise_union`` default - Its type by the rules above must fit the UDT: the UDT itself, or a diff --git a/graphblas/core/base.py b/graphblas/core/base.py index 8dac7eda5..796771595 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -509,12 +509,12 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * # w << (v & v) expr = expr._to_expr() else: - from .scalar import Scalar + from .scalar import Scalar, _literal_store_dtype if type(expr) is Scalar: scalar = expr else: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_store_dtype(self.dtype, expr) try: scalar = Scalar.from_value(expr, dtype, is_cscalar=None, name="") except TypeError: diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index a3477459e..92998e9a0 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -19,6 +19,7 @@ get_typed_op, op_from_string, ) +from .operator.select import _value_select_in_own_type from .scalar import ( _COMPLETE, _MATERIALIZE, @@ -29,9 +30,11 @@ _cast_for_store, _check_literal_fits, _check_scalar_fits, + _element_shapes_differ, _ewise_add_cast_error, _ewise_add_needs_cast, _literal_dtype, + _literal_store_dtype, _scalar_index, _weak_builtin_literal_dtype, ) @@ -401,6 +404,11 @@ def isequal(self, other, *, check_dtype=False, **opts): ------- bool + Notes + ----- + Elements of array UDTs must have the same shape, apart from leading axes + of length 1, as in ``np.array_equal``; ``==`` broadcasts instead. + See Also -------- :meth:`isclose` : For equality check of floating point dtypes @@ -411,6 +419,8 @@ def isequal(self, other, *, check_dtype=False, **opts): ) if check_dtype and self.dtype != other.dtype: return False + if _element_shapes_differ(self.dtype, other.dtype): + return False if self._nrows != other._nrows: return False if self._ncols != other._ncols: @@ -2641,6 +2651,7 @@ def select(self, op, thunk=None): if thunk is None: thunk = False # most basic form of 0 when unifying dtypes + given = thunk if type(thunk) is not Scalar: # A thunk beside a UDT becomes an element of it, so it must fit, as in apply. dtype = _literal_dtype(self.dtype, thunk, op) @@ -2656,6 +2667,7 @@ def select(self, op, thunk=None): op=op, ) op = get_typed_op(op, self.dtype, thunk.dtype, is_right_scalar=True, kind="select") + op, thunk = _value_select_in_own_type(op, self.dtype, thunk, given) self._expect_op(op, ("SelectOp", "IndexUnaryOp"), within=method_name, argname="op") if thunk._is_cscalar: if thunk.dtype._is_udt: @@ -2995,7 +3007,7 @@ def setdiag(self, values, k=0, *, mask=None, accum=None, **opts): elif type(values) is Scalar: is_scalar = True else: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_store_dtype(self.dtype, values) try: # Try to make it a Scalar values = Scalar.from_value(values, dtype, is_cscalar=None, name="") @@ -3141,7 +3153,7 @@ def _prep_for_extract(self, resolved_indexes): def _assign_element(self, resolved_indexes, value): rowidx, colidx = resolved_indexes.indices if type(value) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_store_dtype(self.dtype, value) try: value = Scalar.from_value(value, dtype, is_cscalar=None, name="") except TypeError: @@ -3370,7 +3382,7 @@ def _prep_for_assign(self, resolved_indexes, value, mask, is_submask, replace, o ) else: if type(value) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_store_dtype(self.dtype, value) try: value = Scalar.from_value(value, dtype, is_cscalar=None, name="") except (TypeError, ValueError): diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 3cc506d84..4d12bcbd6 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -10,6 +10,124 @@ from .base import _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper +_VALUE_OPS = {"valueeq", "valuene", "valuelt", "valuele", "valuegt", "valuege"} + + +def _value_select_in_own_type(op, dtype, thunk, given): + """Return ``op`` and ``thunk`` for ``select`` on built-in ``dtype``, as a value op of ``dtype``. + + SuiteSparse:GraphBLAS (8.0 to 10.5, at least) runs ``select`` with a value + op typed for another type than the input's by reading the thunk's bytes as + the input's type, where the spec casts both to the op's type: on an INT8 + vector, ``valueeq`` with ``1.0`` compared with 0, and ``valuelt`` with + ``300`` compared with 44. ``apply`` with the same op is right. So the + comparison is made in the input's own type, with a thunk that selects the + same elements: for integers, ``x < 2.5`` is ``x < 3``, and a thunk no + element can equal, or that every element passes, selects nothing or + everything; for floats, the thunk is rounded toward the side that keeps + the comparison exact. ``given`` is the thunk as the caller gave it, whose + value is exact where ``thunk`` may have rounded it (an int beyond 64 bits + is FP64). Other ops and types are returned unchanged. + """ + # MAINT 2026-10-07: works around the SuiteSparse:GraphBLAS bug above (seen in + # 8.0.2 through 10.6.0; fix proposed in DrTimothyAldenDavis/GraphBLAS#463). + # Once every supported version has the fix, pass the typed op and thunk + # through and delete this function. + if op.type is dtype or dtype._is_udt or op.parent.name not in _VALUE_OPS: + return op, thunk + import math + + import numpy as np + + from ... import indexunary + from ..scalar import Scalar + + parent = op.parent + name = parent.name + if ( + parent is not getattr(select, name) + and parent is not getattr(indexunary, name) + or dtype.np_type.kind not in "biuf" + or thunk._is_empty + ): + return op, thunk + value = given.value if type(given) is Scalar else given + if isinstance(value, (np.generic, np.ndarray)): + value = value.item() + if isinstance(value, complex): + if value.imag != 0: + if name not in {"valueeq", "valuene"}: # pragma: no cover (no complex order) + return op, thunk + value = math.nan # equal to nothing, as a complex number is to a real one + else: + value = value.real + np_type = dtype.np_type + if np_type.kind == "f": + nan, ftype = np_type.type(np.nan), np_type.type + if value != value: # NaN: only ``!=`` holds, for every element + new_name, new_value = name, nan + else: + try: + near = ftype(value) + except OverflowError: # an int too large for any float + near = ftype(math.copysign(math.inf, value)) + above = near if float(near) >= value else np.nextafter(near, ftype(math.inf)) + below = near if float(near) <= value else np.nextafter(near, ftype(-math.inf)) + new_name = name + if name in {"valueeq", "valuene"}: + # No element equals NaN: == selects nothing and != everything. + new_value = near if float(near) == value else nan + else: + new_value = above if name in {"valuelt", "valuege"} else below + else: + if np_type.kind == "b": + lo, hi = 0, 1 + else: + info = np.iinfo(np_type) + lo, hi = int(info.min), int(info.max) + keep_all, keep_none = ("valuege", lo), ("valuelt", lo) + if value != value: + new_name, new_value = keep_all if name == "valuene" else keep_none + elif name in {"valueeq", "valuene"}: + exact = value == math.floor(value) if math.isfinite(value) else False + if exact and lo <= value <= hi: + new_name, new_value = name, int(value) + else: + new_name, new_value = keep_all if name == "valuene" else keep_none + else: + # Integers: x < t is x < ceil(t), x <= t is x <= floor(t), and so on. + if math.isinf(value): + bound = value + elif name in {"valuelt", "valuege"}: + bound = math.ceil(value) + else: + bound = math.floor(value) + passes_all = { + "valuelt": bound > hi, + "valuele": bound >= hi, + "valuegt": bound < lo, + "valuege": bound <= lo, + }[name] + passes_none = { + "valuelt": bound <= lo, + "valuele": bound < lo, + "valuegt": bound >= hi, + "valuege": bound > hi, + }[name] + if passes_all: + new_name, new_value = keep_all + elif passes_none: + new_name, new_value = keep_none + else: + new_name, new_value = name, int(bound) + if np_type.kind == "b": + new_value = bool(new_value) + module = select if isinstance(parent, SelectOp) else indexunary + new_op = getattr(module, new_name)[dtype] + new_thunk = Scalar.from_value(new_value, dtype, is_cscalar=thunk._is_cscalar, name="") + return new_op, new_thunk + + class TypedBuiltinSelectOp(TypedOpBase): __slots__ = () opclass = "SelectOp" diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 20bdff5a9..0c08fd047 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -134,15 +134,17 @@ def _repr_html_(self, collapse=False, expr=None): def __eq__(self, other): """Check equality. - Equality comparison uses :meth:`isequal`. Use that directly for finer control of - what is considered equal. + Compares with ``binary.eq``, as ``==`` on a Vector or Matrix does, so the + elements of array UDTs broadcast (an ``INT8[3]`` Scalar of ``[1, 1, 1]`` + equals ``1``), but returns a bool. Two empty Scalars are equal, and + ``s == None`` checks whether ``s`` is empty. :meth:`isequal` is stricter. """ - return self.isequal(other) + return _scalar_eq(self, other) __hash__ = None def __ne__(self, other): - return not self.isequal(other) + return not _scalar_eq(self, other) def __bool__(self): """Truthiness check. @@ -199,6 +201,11 @@ def isequal(self, other, *, check_dtype=False): ------- bool + Notes + ----- + Elements of array UDTs must have the same shape, apart from leading axes + of length 1, as in ``np.array_equal``; ``==`` broadcasts instead. + See Also -------- :meth:`isclose` : For equality check of floating point dtypes @@ -207,8 +214,15 @@ def isequal(self, other, *, check_dtype=False): if type(other) is not Scalar: if other is None: return self._is_empty + dtype = None + if self.dtype._is_udt: + # A literal is compared as given, as ``==`` types it, not converted + # into the UDT (0.5 would become 0 in an int field), and an array + # literal must have the element's shape, as in np.array_equal. + if _literal_shape_differs(self.dtype, other): + return False + dtype = _literal_type(self.dtype, other, exact=True) or self.dtype try: - dtype = self.dtype if self.dtype._is_udt else None other = Scalar.from_value(other, dtype, is_cscalar=None, name="s_isequal") except TypeError: other = self._expect_type( @@ -222,6 +236,8 @@ def isequal(self, other, *, check_dtype=False): check_dtype = False if check_dtype and self.dtype != other.dtype: return False + if _element_shapes_differ(self.dtype, other.dtype): + return False if self._is_empty: return other._is_empty if other._is_empty: @@ -391,7 +407,7 @@ def value(self, val): val = _Pointer(_as_scalar(val, self.dtype, is_cscalar=True)) dtype_name = "UDT" else: - val = _as_scalar(val, is_cscalar=True) + val = _as_scalar(val, _literal_store_dtype(self.dtype, val), is_cscalar=True) dtype_name = val.dtype.name call(f"GrB_Scalar_setElement_{dtype_name}", [self, val]) @@ -1289,7 +1305,7 @@ def _weak_builtin_literal_dtype(dtype, value, op=None): rv = lookup_dtype(np_type) except ValueError: # A complex literal where the backend has no complex types return None - if kind != "b" and (op is None or not exact) and rv.np_type in _SINGLE_PRECISION: + if kind != "b" and rv.np_type in _SINGLE_PRECISION: import struct parts = (value.real, value.imag) if kind == "c" else (float(value),) @@ -1301,7 +1317,8 @@ def _weak_builtin_literal_dtype(dtype, value, op=None): # A fill value widens rather than lose its value, as an int out of # range does: FP32's to_dense(fill_value=1e300) is float64. return lookup_dtype(np.complex128 if kind == "c" else np.float64) - # As numpy 2 and a UDT warn: fp32_vec + 1e300 (or + 2**200) is inf. + # As numpy 2 and a UDT warn: fp32_vec + 1e300 (or + 2**200) is inf, and + # so is the 1e300 of fp32_vec == 1e300, which numpy compares with inf too. utils._warn_from_caller( f"overflow encountered in cast: {value!r} is infinity as {rv}", RuntimeWarning ) @@ -1562,6 +1579,78 @@ def _ewise_add_cast_error(op, left_dtype, right_dtype): ) +def _scalar_eq(scalar, other): + """Return ``scalar == other`` as a bool, by ``binary.eq`` (see ``Scalar.__eq__``).""" + if other is None: + return scalar._is_empty + if type(other) is not Scalar and output_type(other) is Scalar: + other = other.new(name="s_eq_other") # an expression, such as ``-s`` or ``v[0]`` + if type(other) is Scalar and (scalar._is_empty or other._is_empty): + return scalar._is_empty and other._is_empty + if scalar._is_empty: + return False + if not scalar.dtype._is_udt: + # For built-in dtypes, compare the two values in the type ``eq`` is typed + # for, which is what it computes, without a GraphBLAS call: ``s == 0`` + # takes about 2 us instead of 11. + if type(other) is Scalar: + other_dtype, value = other.dtype, other.value + else: + other_dtype = _weak_builtin_literal_dtype(scalar.dtype, other, binary.eq) + value = other + if other_dtype is not None and not other_dtype._is_udt: + if other_dtype is not scalar.dtype: + other_dtype = get_typed_op(binary.eq, scalar.dtype, other_dtype).type + convert = other_dtype.np_type.type + if other_dtype.np_type in _SINGLE_PRECISION: + # 1e300 beside FP32 is infinity, and _weak_builtin_literal_dtype + # has warned at the caller's line, so numpy need not warn again. + with np.errstate(over="ignore"): + return (convert(scalar.value) == convert(value)).item() + return (convert(scalar.value) == convert(value)).item() + return bool(scalar.ewise_mult(other, binary.eq).new(name="s_eq").value) + + +def _element_shapes_differ(dtype1, dtype2): + """Whether elements of ``dtype1`` and ``dtype2`` differ in shape, for ``isequal``. + + ``==`` on array UDTs broadcasts, as numpy's does, but ``isequal`` asks + whether two values are the same array, as ``np.array_equal`` does: the + shapes must match, apart from leading axes of length 1 (the rule a store + uses), and an element of a built-in dtype has shape ``()``. Records pair + field by field instead (``udt_utils._check_udt_pair``). + """ + if dtype1 is dtype2 or not (dtype1._is_udt or dtype2._is_udt): + return False + from .operator.udt_utils import _strip_leading_ones + + shapes = [] + for np_type in (dtype1.np_type, dtype2.np_type): + if np_type.names is not None: + return False + shapes.append(_strip_leading_ones(np_type.subdtype[1]) if np_type.subdtype else ()) + return shapes[0] != shapes[1] + + +def _literal_shape_differs(dtype, value): + """Whether literal ``value`` differs in shape from an element of array UDT ``dtype``. + + As :func:`_element_shapes_differ`, for ``Scalar.isequal`` with a literal: a + number has shape ``()``, so it equals an ``FP64[1]`` element but not an + ``FP64[3]`` one, as in ``np.array_equal``. ``False`` for a record UDT. + """ + from .operator.udt_utils import _strip_leading_ones + + np_type = dtype.np_type + if np_type.subdtype is None or isinstance(value, dict): + return False + try: + shape = np.shape(value) + except ValueError: # ragged; Scalar.from_value reports it + return False + return _strip_leading_ones(shape) != _strip_leading_ones(np_type.subdtype[1]) + + def _plus_zero(dtype): """Return a Scalar of ``dtype`` that ``plus`` adds to any value without changing it. @@ -1595,6 +1684,19 @@ def _converted_to(operand, dtype, **new_kwargs): return rv +def _literal_store_dtype(dtype, value): + """Return the dtype to make literal ``value`` a Scalar of, to store it as ``dtype``. + + A UDT types the literal itself. A Python int is INT64 by itself, so one + beyond INT64 (the top half of UINT64) is made ``dtype``: + ``uint64_vec[0] = 2**64 - 1`` works, and a type that cannot hold it raises. + ``None`` lets ``Scalar.from_value`` choose. + """ + if dtype._is_udt or type(value) is int and not -(2**63) <= value < 2**63: + return dtype + return None + + def _cast_for_store(value, dtype): """Return Scalar, Vector or Matrix ``value``, or its expression, to assign as ``dtype``. diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index 1c8a220f2..bc8bcf9bc 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -18,6 +18,7 @@ get_typed_op, op_from_string, ) +from .operator.select import _value_select_in_own_type from .scalar import ( _COMPLETE, _MATERIALIZE, @@ -28,9 +29,11 @@ _cast_for_store, _check_literal_fits, _check_scalar_fits, + _element_shapes_differ, _ewise_add_cast_error, _ewise_add_needs_cast, _literal_dtype, + _literal_store_dtype, _scalar_index, _weak_builtin_literal_dtype, ) @@ -358,6 +361,11 @@ def isequal(self, other, *, check_dtype=False, **opts): ------- bool + Notes + ----- + Elements of array UDTs must have the same shape, apart from leading axes + of length 1, as in ``np.array_equal``; ``==`` broadcasts instead. + See Also -------- :meth:`isclose` : For equality check of floating point dtypes @@ -366,6 +374,8 @@ def isequal(self, other, *, check_dtype=False, **opts): other = self._expect_type(other, Vector, within="isequal", argname="other") if check_dtype and self.dtype != other.dtype: return False + if _element_shapes_differ(self.dtype, other.dtype): + return False if self._size != other._size: return False if self._nvals != other._nvals: @@ -1628,6 +1638,7 @@ def select(self, op, thunk=None): if thunk is None: thunk = False # most basic form of 0 when unifying dtypes + given = thunk if type(thunk) is not Scalar: # A thunk beside a UDT becomes an element of it, so it must fit, as in apply. dtype = _literal_dtype(self.dtype, thunk, op) @@ -1643,6 +1654,7 @@ def select(self, op, thunk=None): op=op, ) op = get_typed_op(op, self.dtype, thunk.dtype, is_right_scalar=True, kind="select") + op, thunk = _value_select_in_own_type(op, self.dtype, thunk, given) self._expect_op(op, ("SelectOp", "IndexUnaryOp"), within=method_name, argname="op") if thunk._is_cscalar: if thunk.dtype._is_udt: @@ -1919,7 +1931,7 @@ def _prep_for_extract(self, resolved_indexes, mask=None, is_submask=False): def _assign_element(self, resolved_indexes, value): idx = resolved_indexes.indices[0] if type(value) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_store_dtype(self.dtype, value) try: value = Scalar.from_value(value, dtype, is_cscalar=None, name="") except TypeError: @@ -1986,7 +1998,7 @@ def _prep_for_assign(self, resolved_indexes, value, mask, is_submask, replace, o expr_repr = "[[{2._expr_name} elements]] = {0.name}" else: if type(value) is not Scalar: - dtype = self.dtype if self.dtype._is_udt else None + dtype = _literal_store_dtype(self.dtype, value) try: value = Scalar.from_value(value, dtype, is_cscalar=None, name="") except (TypeError, ValueError): diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 36cffce57..c7f5c18a4 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -5417,6 +5417,66 @@ def ge_first(x, i, j, t): # pragma: no cover (numba) x.select(ge, thunk) +@pytest.mark.skipif("not supports_udfs") +def test_udt_isequal_is_array_equal_and_eq_broadcasts(): + """Isequal compares array elements as np.array_equal; == broadcasts, on Scalars too. + + eq on array UDTs broadcasts, as numpy's == does, and isequal used it, so an + FP64[3] of [1, 1, 1] was isequal to an FP64[1] of [1]. A Scalar's == was + isequal itself, and a literal was converted into the UDT first, so an + INT8[3] of [0, 0, 0] equaled 0.5 (truncated to 0) by either spelling. + """ + f3 = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_IeqF3") + f1 = dtypes.register_anonymous(np.dtype((np.float64, (1,))), "_IeqF1") + f13 = dtypes.register_anonymous(np.dtype((np.float64, (1, 3))), "_IeqF13") + i3 = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_IeqI3") + rec = dtypes.register_anonymous(np.dtype([("ieq_a", np.int8), ("ieq_b", np.float32)]), "_IeqR") + + def vec(dtype, value): + v = Vector(dtype, size=2) + v[1] = value + return v + + ones = vec(f3, [1, 1, 1]) + assert not ones.isequal(vec(f1, [1])) + assert (ones == vec(f1, [1])).new().reduce(monoid.land).new().value # == broadcasts + assert ones.isequal(vec(f13, [[1, 1, 1]])) # leading axes of length 1 do not count + assert ones.isequal(vec(i3, [1, 1, 1])) + assert not ones.isequal(Vector.from_coo([1], [1.0], size=2)) + assert vec(f1, [5]).isequal(Vector.from_coo([1], [5.0], size=2)) + A = Matrix.from_coo([0], [1], [[1, 1, 1]], dtype=f3, nrows=1, ncols=2) + assert not A.isequal(Matrix.from_coo([0], [1], [[1]], dtype=f1, nrows=1, ncols=2)) + assert A.isequal(Matrix.from_coo([0], [1], [[1, 1, 1]], dtype=i3, nrows=1, ncols=2)) + + u = gb.Scalar.from_value([1, 1, 1], i3) + zeros = gb.Scalar.from_value([0, 0, 0], i3) + assert u == 1 + assert u != 0 + assert u == gb.Scalar.from_value(1) + assert gb.Scalar.from_value(1) == u + assert u == gb.Scalar.from_value([1.0], f1) + assert not u.isequal(1) + assert not u.isequal([1]) + assert not u.isequal(gb.Scalar.from_value([1.0], f1)) + assert u.isequal((1, 1, 1)) + assert u.isequal(np.ones((1, 3), dtype=np.int8)) + assert gb.Scalar.from_value([5], f1).isequal(5) + # A literal is compared as given, not converted into the UDT first. + assert zeros != 0.5 + assert not zeros.isequal(0.5) + assert not zeros.isequal((0, 0, 0.5)) + r = gb.Scalar.from_value((0, 0.1), rec) + assert r == (0, 0.1) # 0.1 beside a float32 field is float32, as == types it + assert r.isequal((0, 0.1)) + assert r != (0.5, 0.1) + assert not r.isequal((0.5, 0.1)) + # Empty Scalars are equal to each other, and to None. + assert gb.Scalar(i3) == gb.Scalar(f3) + assert gb.Scalar(i3) == None # noqa: E711 + assert gb.Scalar(i3) != u + assert u != gb.Scalar(i3) + + @pytest.mark.skipif("not supports_udfs") def test_udt_monoid_literal_is_typed_as_its_binaryop(): """A literal beside a Monoid is typed as beside its BinaryOp, everywhere. diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index b8c00ffec..6961300dd 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -367,6 +367,53 @@ def test_scalar_to_numpy(s): assert a.shape == b.shape +def test_uint64_top_half_from_value(): + """An int from 2**63 to 2**64 - 1 goes into a UINT64 GrB_Scalar, Vector or Matrix. + + A Python int was typed INT64 by itself, so it raised OverflowError. + """ + for value in [2**63, 2**64 - 1]: + for is_cscalar in [False, True]: + s = Scalar.from_value(value, dtype=dtypes.UINT64, is_cscalar=is_cscalar) + assert s.value == value + assert s == value + v = Vector(dtypes.UINT64, size=3) + v[0] = value + v[1:] = value + assert v.to_coo()[1].tolist() == [value] * 3 + v << value + assert v.to_coo()[1].tolist() == [value] * 3 + A = Matrix(dtypes.UINT64, nrows=2, ncols=2) + A[0, 1] = value + A[1, :] = value + assert A.to_coo()[2].tolist() == [value] * 3 + f = Vector(dtypes.FP64, size=1) + f[0] = value + assert f[0].new().value == float(value) + with pytest.raises(OverflowError): + Scalar.from_value(2**63, dtype=dtypes.INT8, is_cscalar=False) + i64 = Vector(dtypes.INT64, size=1) + with pytest.raises(OverflowError): + i64[0] = 2**63 # not wrapped + u64 = Vector(dtypes.UINT64, size=1) + with pytest.raises(OverflowError): + u64[0] = 2**64 + + +def test_eq_types_a_literal_as_binary_eq_does(): + """``==`` on a Scalar compares as ``binary.eq`` does, as on a Vector, and returns a bool.""" + assert Scalar.from_value(0.1, dtype=dtypes.FP32) == 0.1 # in FP32, as numpy 2 compares it + assert Scalar.from_value(1, dtype=dtypes.INT8) != 300 # exact, as a comparison + assert Scalar.from_value(1, dtype=dtypes.INT8) == Scalar.from_value(1.0) + not_a_number = float("nan") + assert Scalar.from_value(not_a_number) != not_a_number + assert not (Scalar.from_value(not_a_number) == not_a_number) # noqa: SIM201 + empty = Scalar(dtypes.INT64) + assert empty == Scalar(dtypes.FP64) + assert empty == None # noqa: E711 + assert empty != 0 + + @autocompute def test_neg(): for dtype in sorted( diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 61fb22f27..e7aef8cba 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -1559,6 +1559,16 @@ def test_python_number_is_weak_beside_builtin_dtypes(): assert (u8 > -1).new().to_coo()[1].tolist() == [True, True] assert i8.apply(indexunary.valuelt, 300).new().to_coo()[1].tolist() == [True, True] assert (f32 == 0.1).new().to_coo()[1].tolist() == [True, False] # compared as float32 + # A float too large for FP32 is infinity in a comparison too, with a warning, as in + # numpy 2 (an int out of range is exact; a float has no type that is both FP32 and + # holds it). The Scalar's == warns once, at the caller's line. + finf = Vector.from_coo([0, 1], [np.inf, 1], dtype=dtypes.FP32) + with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: + assert (finf == 1e300).new().to_coo()[1].tolist() == [True, False] + assert {w.filename for w in record} == {__file__} + with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: + assert Scalar.from_value(np.inf, dtype=dtypes.FP32) == 1e300 + assert [w.filename for w in record] == [__file__] # ewise_union defaults and fill values follow the rule too. w8 = Vector.from_coo([1], [1], dtype=dtypes.INT8, size=2) assert i8.ewise_union(w8, binary.plus, 0, 0).dtype == dtypes.INT8 @@ -1581,6 +1591,49 @@ def test_python_number_is_weak_beside_builtin_dtypes(): assert dtypes.unify(dtypes.INT8, dtypes.INT64, is_right_scalar=True) == dtypes.INT64 +def test_select_value_thunk_of_another_type(): + """Select with a value op compares exactly when the thunk's type differs from the values'. + + SuiteSparse:GraphBLAS's select reads such a thunk's bytes as the values' + type: on an INT8 vector, ``valueeq`` with 1.0 compared with 0 and kept + nothing, ``valuelt`` with 300 compared with 44, and ``valueeq`` with 356 + kept 100. python-graphblas now compares in the values' own type with a + thunk that keeps the same elements. + """ + i8 = Vector.from_coo([0, 1, 2, 3], [-5, 1, 2, 100], dtype=dtypes.INT8) + u8 = Vector.from_coo([0, 1], [0, 255], dtype=dtypes.UINT8) + f32 = Vector.from_coo([0, 1, 2], [np.float32(0.1), 1.0, np.nan], dtype=dtypes.FP32) + bools = Vector.from_coo([0, 1], [False, True]) + for x, op, thunk, kept in [ + (i8, select.valueeq, 1.0, [1]), + (i8, select.valueeq, 356, []), + (i8, select.valuene, 356, [0, 1, 2, 3]), + (i8, select.valuelt, 300, [0, 1, 2, 3]), + (i8, select.valuege, 2.5, [3]), + (i8, select.valuele, 2.5, [0, 1, 2]), + (i8, select.valuegt, -300, [0, 1, 2, 3]), + (i8, select.valuegt, np.nan, []), + (i8, select.valuene, np.nan, [0, 1, 2, 3]), + (i8, indexunary.valueeq, np.float64(2.0), [2]), + (i8, select.valueeq, Scalar.from_value(1.0), [1]), + (i8, select.valuelt, Scalar.from_value(300, dtype=dtypes.INT64), [0, 1, 2, 3]), + (u8, select.valuelt, np.int64(-1), []), + (u8, select.valuegt, np.int64(-1), [0, 1]), + (bools, select.valueeq, 256, []), + (bools, select.valueeq, 1, [1]), + # A strong FP64 thunk compares in FP64: float32(0.1) is not 0.1. + (f32, select.valueeq, np.float64(0.1), []), + (f32, select.valuegt, np.float64(0.1), [0, 1]), + (f32, select.valuene, np.float64(0.1), [0, 1, 2]), + # A Python float beside FP32 is FP32 (NEP 50), as numpy compares it. + (f32, select.valueeq, 0.1, [0]), + ]: + assert x.select(op, thunk).new().to_coo()[0].tolist() == kept, (x.dtype, op, thunk) + A = Matrix.from_coo([0, 0, 1], [0, 1, 1], [1, 100, -5], dtype=dtypes.INT8) + assert A.select(select.valueeq, 1.0).new().to_coo()[2].tolist() == [1] + assert A.select(select.valuelt, 300).new().nvals == 3 + + def test_inner(v): R = Matrix(v.dtype, nrows=1, ncols=v.size) # row vector C = Matrix(v.dtype, nrows=v.size, ncols=1) # column vector From 66908142eb1629d5e115780f3e90671b3c9ce003 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 8 Oct 2026 13:35:55 -0500 Subject: [PATCH 08/54] Compare a Python int beyond a 64-bit element's range exactly Commit 2 typed an int beyond an integer element's range by a type that holds it, so a comparison no longer raised, but beside INT64 and UINT64 that comparison runs in FP64, which rounds both: int64_vec < 2**63 was False for 2**63 - 1, uint64_vec == 2**64 was True for 2**64 - 1, and Scalar(2**63 - 1) == 2**63 was True, as for UDT int64 and uint64 leaves (main raised OverflowError). Such an int now compares as infinity of its sign, which every integer element compares with as with the int, so infix comparisons, apply, ewise_union defaults and Scalar == are exact, as in numpy 2, and beside a UDT a literal with one equals no element. An int too large for a float compares too (int64_vec < 10**400 and select(valueeq, 10**400) raised OverflowError). A float too large for a float32 UDT leaf (fp32_udt + 1e300, or a store) now warns at the caller's line, as a built-in FP32 does, instead of from inside python-graphblas, where the default filter showed it once per process. Drops a test helper nothing used. --- docs/user_guide/types.rst | 3 +- docs/user_guide/udt.rst | 6 +- graphblas/core/dtypes.py | 2 +- graphblas/core/matrix.py | 12 +-- graphblas/core/operator/select.py | 5 +- graphblas/core/operator/udt_utils.py | 30 ++++++- graphblas/core/scalar.py | 112 ++++++++++++++++++++------- graphblas/core/vector.py | 12 +-- graphblas/tests/test_op.py | 54 ++++++++++--- graphblas/tests/test_vector.py | 50 ++++++++++++ 10 files changed, 225 insertions(+), 61 deletions(-) diff --git a/docs/user_guide/types.rst b/docs/user_guide/types.rst index e6fd91756..36d93e8a4 100644 --- a/docs/user_guide/types.rst +++ b/docs/user_guide/types.rst @@ -46,7 +46,8 @@ so it does not depend on the installed version of numpy. User-defined types follow the same rule, field by field (see :doc:`udt`). A comparison has no result type to keep, so an integer outside the operand's -range compares exactly, as in numpy: ``v < 300`` is True for every element. +range compares exactly, as in numpy: ``v < 300`` is True for every element, as +is ``w < 2**63`` for an INT64 vector ``w``. Nor does ``v / 300`` raise, since its result is FP64 whatever the integer. An FP32 vector ``== 0.1`` compares in FP32, and ``== 1e300`` compares with infinity, with a ``RuntimeWarning`` as in numpy. The defaults of ``ewise_union`` diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index e7aab686b..25a41ad1d 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -176,9 +176,9 @@ dtype, so it never narrows. Below, ``v`` is a Vector of an ``INT8[3]`` UDT. - ``v + 1`` is ``INT8[3]``, ``v * 2.5`` and ``v + (0.5, 1, 2)`` are ``FP64[3]``, ``v + 300`` raises * - ``eq`` and ``ne`` - - Type a literal as above, as built-in comparisons do, except that an int - out of a field's range takes a type that holds it, so the comparison is - exact. Return ``BOOL``. Array elements broadcast, as numpy's ``==`` + - Type a literal as above, as built-in comparisons do, except that a + literal with an int out of a field's range equals no element, rather + than raise. Return ``BOOL``. Array elements broadcast, as numpy's ``==`` does, and so does ``==`` on a Scalar, which returns a bool. ``isequal`` compares as ``np.array_equal``: elements must have the same shape, apart from leading axes of length 1. diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 777fbc2a6..130055e30 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -568,7 +568,7 @@ def unify(type1, type2, *, is_left_scalar=False, is_right_scalar=False): unify(FP32, INT32) -> FP64 ``is_left_scalar`` and ``is_right_scalar`` no longer change the result. A - Python number is typed before it gets here (``scalar._literal_dtype``), and + Python number is typed before it gets here (``scalar._literal_operand``), and a typed scalar promotes as any operand does, as numpy 2 promotes it. They used to stand the scalar in as a 0-d array of value 0, which numpy 1 typed by value, so ``int8_vec + 300`` wrapped there. diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 92998e9a0..a89b52610 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -33,7 +33,7 @@ _element_shapes_differ, _ewise_add_cast_error, _ewise_add_needs_cast, - _literal_dtype, + _literal_operand, _literal_store_dtype, _scalar_index, _weak_builtin_literal_dtype, @@ -2145,7 +2145,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): if dtype is not None: _check_literal_fits(dtype, left_default) else: - dtype = _literal_dtype(left_dtype, left_default, op) + left_default, dtype = _literal_operand(left_dtype, left_default, op) try: left = Scalar.from_value( left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -2168,7 +2168,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): if dtype is not None: _check_literal_fits(dtype, right_default) else: - dtype = _literal_dtype(right_dtype, right_default, op) + right_default, dtype = _literal_operand(right_dtype, right_default, op) try: right = Scalar.from_value( right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -2495,7 +2495,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = None elif right is None: if type(left) is not Scalar: - dtype = _literal_dtype(self.dtype, left, op) + left, dtype = _literal_operand(self.dtype, left, op) try: left = Scalar.from_value(left, dtype, is_cscalar=None, name="") except TypeError: @@ -2531,7 +2531,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = "{1.name}.apply({op}, left={0._expr_name})" elif left is None: if type(right) is not Scalar: - dtype = _literal_dtype(self.dtype, right, op) + right, dtype = _literal_operand(self.dtype, right, op) try: right = Scalar.from_value(right, dtype, is_cscalar=None, name="") except TypeError: @@ -2654,7 +2654,7 @@ def select(self, op, thunk=None): given = thunk if type(thunk) is not Scalar: # A thunk beside a UDT becomes an element of it, so it must fit, as in apply. - dtype = _literal_dtype(self.dtype, thunk, op) + thunk, dtype = _literal_operand(self.dtype, thunk, op) try: thunk = Scalar.from_value(thunk, dtype, is_cscalar=None, name="") except TypeError: diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 4d12bcbd6..3c99df360 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -89,14 +89,15 @@ def _value_select_in_own_type(op, dtype, thunk, given): if value != value: new_name, new_value = keep_all if name == "valuene" else keep_none elif name in {"valueeq", "valuene"}: - exact = value == math.floor(value) if math.isfinite(value) else False + # An int is whole, and math.isfinite cannot take one too large for a float. + exact = type(value) is int or (math.isfinite(value) and value == math.floor(value)) if exact and lo <= value <= hi: new_name, new_value = name, int(value) else: new_name, new_value = keep_all if name == "valuene" else keep_none else: # Integers: x < t is x < ceil(t), x <= t is x <= floor(t), and so on. - if math.isinf(value): + if type(value) is float and math.isinf(value): bound = value elif name in {"valuelt", "valuege"}: bound = math.ceil(value) diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index abcbab340..7b5dce731 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -615,9 +615,10 @@ def _weak_literal_leaf(values, kind, elem, *, exact=False): Each is typed by :func:`_weak_literal_element`, so an int out of an integer ``elem``'s range raises ``OverflowError``. With ``exact``, for a comparison, - which keeps no result type, such an int takes int64, uint64 or float64, - whichever holds the values, so it compares exactly, as numpy compares it: - ``int8_vec < 300`` is all True. Built-in dtypes and UDT leaves both use this. + which keeps no result type, a fill value or truediv, such an int takes + int64, uint64 or float64, whichever holds the values (a comparison then + compares it as infinity: ``scalar._literal_operand``). Built-in dtypes and + UDT leaves both use this. """ if exact and kind == "i" and elem.kind in "biu": # The integer type an int would take: ``elem``, or INT64 beside a bool. @@ -733,6 +734,29 @@ def record_leaves(np_type, part): return dtype if literal_type == np_type else lookup_dtype(literal_type) +def _int_beyond_leaf(np_type, value): + """Return a Python int in literal ``value`` beyond its integer leaf's range, or ``None``. + + ``value`` is laid out like an element of ``np_type``, as + :func:`_weak_literal_udt` reads it (a record's dict already a tuple). No + element equals a literal with such an int, whatever its other leaves. + """ + if np_type.names is not None: + parts = value if isinstance(value, (tuple, list)) else [value] * len(np_type.names) + for name, part in zip(np_type.names, parts, strict=True): + found = _int_beyond_leaf(np_type.fields[name][0], part) + if found is not None: + return found + return None + base, shape = np_type.subdtype or (np_type, ()) + if base.kind in "iu": + lo, hi = _int_bounds(base) + for x in _literal_leaf_values(value, shape) or (): + if type(x) is int and not lo <= x <= hi: + return x + return None + + def _array_literal_udt(dtype, value): """Return the UDT a numpy array ``value`` is beside UDT ``dtype``, or ``None`` to convert it. diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 0c08fd047..b28595be1 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -1,4 +1,5 @@ import itertools +import math from functools import lru_cache import numpy as np @@ -397,7 +398,23 @@ def value(self, val): val = _dict_to_record(np_type, val) else: arr = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) - arr[:] = _view_if_same_layout(val, arr.dtype) + val = _view_if_same_layout(val, arr.dtype) + if _has_narrow_float_leaf(np_type) and np.geterr()["over"] == "warn": + # A float too large for a float32 leaf is infinity, and numpy + # warns; warn at the caller's line instead, as a built-in + # FP32 does, so the default filter shows each place it happens. + try: + with np.errstate(over="raise"): + arr[:] = val + except FloatingPointError: + utils._warn_from_caller( + f"overflow encountered in cast: {val!r} as {self.dtype}", + RuntimeWarning, + ) + with np.errstate(over="ignore"): + arr[:] = val + else: + arr[:] = val self.gb_obj[0 : np_type.itemsize] = _udt_bytes(np_type, arr) else: self.gb_obj[0] = val @@ -689,7 +706,7 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): other = op(other).new() if type(other) is not Scalar: - dtype = _literal_dtype(self.dtype, other, op) + other, dtype = _literal_operand(self.dtype, other, op) try: other = Scalar.from_value(other, dtype, is_cscalar=False, name="") except TypeError: @@ -773,7 +790,7 @@ def _ewise_mult(self, other, op=binary.times, is_infix=False): other = op(other).new() if type(other) is not Scalar: - dtype = _literal_dtype(self.dtype, other, op) + other, dtype = _literal_operand(self.dtype, other, op) try: other = Scalar.from_value(other, dtype, is_cscalar=False, name="") except TypeError: @@ -850,10 +867,9 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): right_dtype = self.dtype dtype = right_dtype if right_dtype._is_udt else None if type(other) is not Scalar: + other, literal_dtype = _literal_operand(right_dtype, other, op) try: - other = Scalar.from_value( - other, _literal_dtype(right_dtype, other, op), is_cscalar=False, name="" - ) + other = Scalar.from_value(other, literal_dtype, is_cscalar=False, name="") except TypeError: other = self._expect_type( other, @@ -879,7 +895,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): if dtype is not None: _check_literal_fits(dtype, left_default) else: - dtype = _literal_dtype(left_dtype, left_default, op) + left_default, dtype = _literal_operand(left_dtype, left_default, op) try: left = Scalar.from_value( left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -902,7 +918,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): if dtype is not None: _check_literal_fits(dtype, right_default) else: - dtype = _literal_dtype(right_dtype, right_default, op) + right_default, dtype = _literal_operand(right_dtype, right_default, op) try: right = Scalar.from_value( right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -1191,12 +1207,12 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=False, name=None, **opts): # End auto-generated code: Scalar -def _literal_dtype(dtype, value, op): - """Return the dtype to build ``value`` with, as the other operand of ``op`` beside ``dtype``. +def _literal_operand(dtype, value, op): + """Return literal ``value`` and the dtype to make it, as ``op``'s operand beside ``dtype``. - ``None`` means the value's own dtype. This is where python-graphblas - decides how a literal operand is typed; beside a built-in dtype, see - :func:`_weak_builtin_literal_dtype`. + A dtype of ``None`` means the value's own dtype. This is where + python-graphblas decides how a literal operand is typed; beside a built-in + dtype, see :func:`_weak_builtin_literal_dtype`. Beside a UDT, a literal becomes an element of that UDT, as it must for a user's op, when its type allows it (:func:`_check_literal_fits`). With a @@ -1215,21 +1231,41 @@ def _literal_dtype(dtype, value, op): weak in the same way, leaf by leaf: ``int8_udt + (1, 2, 3)`` stays ``int8_udt`` and ``+ (0.5, 1.5, 2.5)`` has float64 elements. - ``eq`` and ``ne`` type a literal the same way, as numpy compares, but a - comparison keeps no result type, so an int out of a leaf's range takes a - type that holds it: ``int8_udt == 300`` is False, not an error, and + ``eq`` and ``ne`` type a literal the same way, as numpy compares, so ``fp32_udt == 0.1`` compares in float32, as ``fp32_vec == 0.1`` does. + + A comparison keeps no result type, so a Python int beyond an integer + element's range is compared exactly, as numpy 2 compares it, as infinity + of its sign: every element compares with that infinity as with the int, + while a type that held both would round them (in FP64, the INT64 + ``2**63 - 1`` and the int ``2**63`` are equal). So ``int64_vec < 2**63`` is + all True and ``int8_vec == 300`` all False. Beside a UDT, where only ``eq`` + and ``ne`` compare, a literal with such an int in an integer leaf equals no + element: ``int8_udt == 300`` is False, not an error. """ if getattr(op, "is_positional", False): - return None + return value, None if not dtype._is_udt: - return _weak_builtin_literal_dtype(dtype, value, op) + literal_dtype = _weak_builtin_literal_dtype(dtype, value, op) + if ( + literal_dtype is not dtype + and literal_dtype is not None + and type(value) is int + and dtype.np_type.kind in "iu" + ): + # Only an int beyond dtype's range takes another type here, typed so + # by a comparison or by an op that gives another type (truediv). + if isinstance(op, str): + op = binary_from_string(op) + if op.name in _COMPARISONS: + return (math.inf if value > 0 else -math.inf), FP64 + return value, literal_dtype if getattr(op, "_udt_types", True) is None: # An op that does not take UDTs at all (``lt``) reports that itself. - return dtype + return value, dtype from .operator.binary import BinaryOp from .operator.monoid import Monoid - from .operator.udt_utils import BUILTIN_UDT_BINARY_OPS + from .operator.udt_utils import BUILTIN_UDT_BINARY_OPS, _int_beyond_leaf if isinstance(op, Monoid) and not op._anonymous: # Element-wise, a Monoid is its BinaryOp (``get_typed_op`` types it so). @@ -1242,9 +1278,17 @@ def _literal_dtype(dtype, value, op): exact = is_comparison or op.name == "truediv" literal_type = _literal_type(dtype, value, exact=exact) if literal_type is not None: - return literal_type + if is_comparison and not isinstance(value, (np.generic, np.ndarray)): + np_type = dtype.np_type + parts = _dict_to_record(np_type, value) if isinstance(value, dict) else value + beyond = _int_beyond_leaf(np_type, parts) + if beyond is not None: + # As infinity in every leaf, it still equals no element. + value = math.inf if beyond > 0 else -math.inf + literal_type = _literal_type(dtype, value) + return value, literal_type _check_literal_fits(dtype, value) - return dtype + return value, dtype # Built-in ops that compare their operands, and so keep no result type. @@ -1270,9 +1314,9 @@ def _weak_builtin_literal_dtype(dtype, value, op=None): not compare it with the values, such as a row index. A comparison keeps no result type, so an int beyond ``dtype``'s range - takes a type that holds it and compares exactly, as numpy compares it: - ``int8_vec < 300`` is all True. ``op`` of ``None``, for a fill value, does - the same, so the value is not lost. + takes a type that holds it rather than raise (:func:`_literal_operand` then + compares it as infinity). ``op`` of ``None``, for a fill value, does the + same, so the value is not lost. """ kind = _LITERAL_KINDS.get(type(value)) if kind is None: @@ -1596,15 +1640,14 @@ def _scalar_eq(scalar, other): if type(other) is Scalar: other_dtype, value = other.dtype, other.value else: - other_dtype = _weak_builtin_literal_dtype(scalar.dtype, other, binary.eq) - value = other + value, other_dtype = _literal_operand(scalar.dtype, other, binary.eq) if other_dtype is not None and not other_dtype._is_udt: if other_dtype is not scalar.dtype: other_dtype = get_typed_op(binary.eq, scalar.dtype, other_dtype).type convert = other_dtype.np_type.type if other_dtype.np_type in _SINGLE_PRECISION: - # 1e300 beside FP32 is infinity, and _weak_builtin_literal_dtype - # has warned at the caller's line, so numpy need not warn again. + # 1e300 beside FP32 is infinity, and _literal_operand has + # warned at the caller's line, so numpy need not warn again. with np.errstate(over="ignore"): return (convert(scalar.value) == convert(value)).item() return (convert(scalar.value) == convert(value)).item() @@ -1743,6 +1786,17 @@ def _dict_to_record(np_type, d): return tuple(rv) +@lru_cache +def _has_narrow_float_leaf(np_type): + """Whether ``np_type`` has a float leaf narrower than float64, which a float may overflow.""" + if np_type.subdtype is not None: + return _has_narrow_float_leaf(np_type.subdtype[0]) + if np_type.names is not None: + return any(_has_narrow_float_leaf(np_type.fields[name][0]) for name in np_type.names) + kind, size = np_type.kind, np_type.itemsize + return (kind == "f" and size < 8) or (kind == "c" and size < 16) + + @lru_cache def _padding_ranges(np_type): """The ``(start, stop)`` byte ranges of ``np_type`` that are alignment padding. diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index bc8bcf9bc..c8aaeb1e6 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -32,7 +32,7 @@ _element_shapes_differ, _ewise_add_cast_error, _ewise_add_needs_cast, - _literal_dtype, + _literal_operand, _literal_store_dtype, _scalar_index, _weak_builtin_literal_dtype, @@ -1240,7 +1240,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): if dtype is not None: _check_literal_fits(dtype, left_default) else: - dtype = _literal_dtype(left_dtype, left_default, op) + left_default, dtype = _literal_operand(left_dtype, left_default, op) try: left = Scalar.from_value( left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -1263,7 +1263,7 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): if dtype is not None: _check_literal_fits(dtype, right_default) else: - dtype = _literal_dtype(right_dtype, right_default, op) + right_default, dtype = _literal_operand(right_dtype, right_default, op) try: right = Scalar.from_value( right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar @@ -1486,7 +1486,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = None elif right is None: if type(left) is not Scalar: - dtype = _literal_dtype(self.dtype, left, op) + left, dtype = _literal_operand(self.dtype, left, op) try: left = Scalar.from_value(left, dtype, is_cscalar=None, name="") except TypeError: @@ -1522,7 +1522,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = "{1.name}.apply({op}, left={0._expr_name})" elif left is None: if type(right) is not Scalar: - dtype = _literal_dtype(self.dtype, right, op) + right, dtype = _literal_operand(self.dtype, right, op) try: right = Scalar.from_value(right, dtype, is_cscalar=None, name="") except TypeError: @@ -1641,7 +1641,7 @@ def select(self, op, thunk=None): given = thunk if type(thunk) is not Scalar: # A thunk beside a UDT becomes an element of it, so it must fit, as in apply. - dtype = _literal_dtype(self.dtype, thunk, op) + thunk, dtype = _literal_operand(self.dtype, thunk, op) try: thunk = Scalar.from_value(thunk, dtype, is_cscalar=None, name="") except TypeError: diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index c7f5c18a4..d1680f34d 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -3526,8 +3526,10 @@ def test_udt_eq_ne_type_a_literal_as_builtin_comparisons_do(): The number is weak, leaf by leaf, so ``fp32_udt == 0.1`` compares in float32 and is True, as ``fp32_vec == 0.1`` and numpy are; it was strong, and compared float32 0.1 with float64 0.1. A comparison keeps no result - type, so an int out of a leaf's range takes a type that holds it, even past - int64, where it raised. numpy scalars stay strong. + type, so a literal with an int out of a leaf's range equals no element, + even past int64, where it raised. Beside an int64 leaf, it is not rounded + to equal one (in float64, the int64 ``2**63 - 1`` and ``2**63`` are equal). + numpy scalars stay strong. """ f32 = dtypes.register_anonymous(np.dtype((np.float32, (7,))), "_EqWeakF32") i8 = dtypes.register_anonymous(np.dtype((np.int8, (7,))), "_EqWeakI8") @@ -3546,6 +3548,46 @@ def test_udt_eq_ne_type_a_literal_as_builtin_comparisons_do(): assert (w != number).new()[0].new().value assert (w == tuple(range(7))).new()[0].new().value assert not (w == (0.5, *range(1, 7))).new()[0].new().value + i64 = dtypes.register_anonymous(np.dtype((np.int64, (2,))), "_EqWeakI64") + x = Vector(i64, size=2) + x[0] = [2**63 - 1, 2**63 - 1] + x[1] = [-(2**63), 0] + for literal in [2**63, (2**63, 2**63), [2**63 - 1, 2**63], -(2**63) - 1, 10**400]: + assert (x == literal).new().to_coo()[1].tolist() == [False, False] + assert (x != literal).new().to_coo()[1].tolist() == [True, True] + assert (x == (2**63 - 1, 2**63 - 1)).new().to_coo()[1].tolist() == [True, False] + rec = dtypes.register_anonymous( + np.dtype([("a", np.int64), ("b", np.uint64)], align=True), "_EqWeakRec64" + ) + r = Vector(rec, size=1) + r[0] = (2**63 - 1, 2**64 - 1) + for literal in [(2**63, 2**64 - 1), {"a": 2**63 - 1, "b": 2**64}, (2**63 - 1, -1)]: + assert not (r == literal).new()[0].new().value + assert (r != literal).new()[0].new().value + assert (r == {"a": 2**63 - 1, "b": 2**64 - 1}).new()[0].new().value + + +def test_udt_float32_overflow_warns_at_the_callers_line(): + """A float too large for a float32 leaf is infinity, with a warning at the caller's line. + + numpy warned from inside python-graphblas, so the default filter showed it + once per process; a built-in FP32 warns at the caller's line, as numpy does. + """ + f32 = dtypes.register_anonymous(np.dtype((np.float32, (2,))), "_OverflowF32") + v = Vector(f32, size=1) + with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: + v[0] = [1e300, 1] + assert [w.filename for w in record] == [__file__] + assert v[0].new().value.tolist() == [np.inf, 1] + with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: + (v == 1e300).new() + assert [w.filename for w in record] == [__file__] + if supports_udfs: + with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: + (v + 1e300).new() + assert [w.filename for w in record] == [__file__] + with np.errstate(over="ignore"): + v[0] = [1e300, 1] # no warning, as numpy @pytest.mark.skipif("not supports_udfs") @@ -4930,14 +4972,6 @@ def test_udt_fused_store_matches_compute_then_cast(): assert C.isequal(A.T.apply(binary.times, right=2).new().dup(dtype=i16s)) -def _udt_leaf_arrays(vector, size): - """The values of a record UDT vector's entries 0 to ``size - 1``, one array per leaf.""" - values = [vector[k].new().value for k in range(size)] - return { - name: np.array([value[name] for value in values]) for name in vector.dtype.np_type.names - } - - def _assert_same_floats(got, expected, msg): """Equal, NaN equal to NaN, and with the same sign on every zero (of each complex part).""" got, expected = np.asarray(got), np.asarray(expected) diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index e7aef8cba..936b6b036 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -1591,6 +1591,56 @@ def test_python_number_is_weak_beside_builtin_dtypes(): assert dtypes.unify(dtypes.INT8, dtypes.INT64, is_right_scalar=True) == dtypes.INT64 +def test_python_int_beyond_64_bit_element_compares_exactly(): + """A Python int beyond a 64-bit integer element's range compares exactly, as in numpy 2. + + It compares as infinity of its sign. A type that held both would round + them: in FP64, the INT64 ``2**63 - 1`` and the int ``2**63`` are equal, so + ``int64_vec < 2**63`` was False there (main raised OverflowError). + """ + compare = { + "eq": lambda x, y: x == y, + "ne": lambda x, y: x != y, + "lt": lambda x, y: x < y, + "le": lambda x, y: x <= y, + "gt": lambda x, y: x > y, + "ge": lambda x, y: x >= y, + } + i64_values = [-(2**63), -(2**63) + 1, 0, 2**63 - 2, 2**63 - 1] + u64_values = [0, 2**63, 2**64 - 2, 2**64 - 1] + i64 = Vector.from_coo(range(5), i64_values, dtype=dtypes.INT64) + u64 = Vector.from_coo(range(4), u64_values, dtype=dtypes.UINT64) + for vec, values, numbers in [ + (i64, i64_values, [2**63, -(2**63) - 1, 2**64, -(2**64), 10**400]), + (u64, u64_values, [2**64, 2**64 + 1, -1, -(10**400)]), + ]: + empty = Vector(vec.dtype, size=vec.size) + for name, py_op in compare.items(): + for number in numbers: + expected = [py_op(x, number) for x in values] + for expr in [ + py_op(vec, number), + vec.apply(getattr(binary, name), right=number), + vec.apply(getattr(indexunary, f"value{name}"), number), + vec.ewise_union(empty, getattr(binary, name), 0, number), + ]: + assert expr.new().to_coo()[1].tolist() == expected + kept = vec.select(getattr(select, f"value{name}"), number).new() + assert kept.to_coo()[0].tolist() == [i for i, x in enumerate(expected) if x] + flipped = py_op(number, vec).new() + assert flipped.to_coo()[1].tolist() == [py_op(number, x) for x in values] + A = Matrix.from_coo([0, 1], [0, 1], [2**63 - 1, -(2**63)]) + assert (A < 2**63).new().to_coo()[2].tolist() == [True, True] + assert (A == -(2**63) - 1).new().to_coo()[2].tolist() == [False, False] + s = Scalar.from_value(2**63 - 1) + assert (s == 2**63) is False + assert s != 2**63 + assert (s < 2**63).new().value + assert (Scalar.from_value(2**64 - 1, dtype=dtypes.UINT64) == 2**64) is False + # A float element compares with the int as a float, as numpy 2 does. + assert (Vector.from_coo([0], [2.0**63]) == 2**63 + 1).new().to_coo()[1].tolist() == [True] + + def test_select_value_thunk_of_another_type(): """Select with a value op compares exactly when the thunk's type differs from the values'. From 3ffd904521fb49d69e0f5a6b76278e63a3d4f13b Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 8 Oct 2026 21:15:45 -0500 Subject: [PATCH 09/54] Give each UDT layout its own JIT C name; Scalar == without Numba; CI fixes SuiteSparse's JIT knows a type by its C name alone, and a UDT's C name was the name it was registered under, so two layouts registered under one name shared JIT kernels: making an iso Vector non-iso ran a kernel compiled for the other layout, which corrupted the heap or gave wrong values (on main too; the new broadcasting test registers four layouts as _BcastX, and crashed 9 Linux CI jobs). A name already given to another layout now gets a fresh _gbudt_NNN, and gb.ss.register_new refuses a name another type has. Scalar == on a UDT runs binary.eq, which needs Numba; without Numba it raised KeyError, where it compared in Python before. It now compares the values in numpy, leaf by leaf in the type binary.eq compares in, and raises binary.eq's KeyError for UDTs that do not pair. Tests: guard the cases that need Numba, skip the 32 KB UDT on Windows (MSVC builds of SuiteSparse refuse a UDT over 1024 bytes), give the JIT name tests array lengths no other test registers, and expect numpy's own int dtype (int32 on Windows with numpy 1). --- .pre-commit-config.yaml | 1 - graphblas/core/dtypes.py | 12 ++++++- graphblas/core/operator/udt_utils.py | 6 ++++ graphblas/core/scalar.py | 43 +++++++++++++++++++++++ graphblas/core/ss/dtypes.py | 7 +++- graphblas/tests/test_dtype.py | 38 ++++++++++++++++++++ graphblas/tests/test_op.py | 51 ++++++++++++++++----------- graphblas/tests/test_scalar.py | 52 ++++++++++++++++++++++++++++ graphblas/tests/test_vector.py | 5 ++- 9 files changed, 191 insertions(+), 24 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3b36390ad..def06f9d5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -29,7 +29,6 @@ repos: - id: check-json # no JSON files yet; enabled for the future - id: check-toml - id: check-yaml - - id: check-executables-have-shebangs - id: check-vcs-permalinks - id: debug-statements - id: destroyed-symlinks diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 130055e30..4800efaf9 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -143,7 +143,12 @@ def _set_udt_jit_c_definition(datatype): Supports both record UDTs (``{'x': float, 'y': float}``) and array UDTs (``np.dtype((np.float64, (3,)))``). """ - from .operator.udt_utils import _has_jit_set + from .operator.udt_utils import ( + _has_jit_set, + _jit_c_type_defs, + _synthetic_udt_counter, + _udt_c_typedef, + ) if not _has_jit_set or not datatype._is_udt: return @@ -151,6 +156,11 @@ def _set_udt_jit_c_definition(datatype): if info is None: return c_name, typedef = info + # A name already given to another layout (two UDTs registered under one + # name) gets a fresh one: SuiteSparse would share their JIT kernels. + while _jit_c_type_defs.setdefault(c_name, typedef) != typedef: + info = _udt_c_typedef(f"_gbudt_{next(_synthetic_udt_counter)}", datatype.np_type) + c_name, typedef = info lib.GrB_Type_set_String(datatype._carg, ffi.new("char[]", c_name.encode()), lib.GxB_JIT_C_NAME) lib.GrB_Type_set_String( datatype._carg, ffi.new("char[]", typedef.encode()), lib.GxB_JIT_C_DEFINITION diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index 7b5dce731..a0ca8856e 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -294,6 +294,12 @@ def _is_valid_c_identifier(name): # process state. _synthetic_udt_counter = itertools.count(1) +# The C typedef given to SuiteSparse under each ``GxB_JIT_C_NAME`` in this +# process. SuiteSparse's JIT knows a type by that name alone (its hash is the +# name's), so two layouts under one name would share compiled kernels, and a +# kernel compiled for one would read and write the other's element size. +_jit_c_type_defs = {} + def _pick_c_type_name(python_name): """Return a valid C identifier for use as ``GxB_JIT_C_NAME`` on a UDT. diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index b28595be1..4420ee75f 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -139,6 +139,8 @@ def __eq__(self, other): elements of array UDTs broadcast (an ``INT8[3]`` Scalar of ``[1, 1, 1]`` equals ``1``), but returns a bool. Two empty Scalars are equal, and ``s == None`` checks whether ``s`` is empty. :meth:`isequal` is stricter. + Without Numba, which ``binary.eq`` needs for UDTs, a UDT's values are + compared in numpy, which broadcasts them the same way. """ return _scalar_eq(self, other) @@ -1651,9 +1653,50 @@ def _scalar_eq(scalar, other): with np.errstate(over="ignore"): return (convert(scalar.value) == convert(value)).item() return (convert(scalar.value) == convert(value)).item() + if not _has_numba and (scalar.dtype._is_udt or (type(other) is Scalar and other.dtype._is_udt)): + # binary.eq on a UDT needs Numba, so compare the values in numpy, as == + # on a Scalar did before it used binary.eq, with the literal typed as + # binary.eq types it, and leaf by leaf as binary.eq compares. + from .operator.udt_utils import _check_udt_pair, _get_udt_info + + if type(other) is not Scalar: + value, dtype = _literal_operand(scalar.dtype, other, binary.eq) + other = Scalar.from_value(value, dtype, is_cscalar=None, name="s_eq_other") + # Two UDTs pair as for binary.eq, or raise its KeyError (records with other + # field names, array shapes that do not broadcast, a record and an array). + info = _get_udt_info(scalar.dtype), _get_udt_info(other.dtype) + _check_udt_pair("eq", scalar.dtype, other.dtype, *info) + return _leaves_equal(scalar.value, scalar.dtype.np_type, other.value, other.dtype.np_type) return bool(scalar.ewise_mult(other, binary.eq).new(name="s_eq").value) +def _leaves_equal(x, x_type, y, y_type): + """Whether values ``x`` and ``y`` are equal in every leaf, as ``binary.eq`` compares them. + + For ``==`` on a UDT Scalar without Numba. Records pair their fields in + order (the caller checks, by ``_check_udt_pair``, that two records pair as + for ``binary.eq``), and a value that is not a record stands in every + field. Each pair of leaves compares in the type their dtypes promote to, + whatever the values (numpy 1 types a scalar by its value), and broadcasts. + """ + if x_type.names is not None or y_type.names is not None: + x_names = x_type.names or [None] * len(y_type.names) + y_names = y_type.names or [None] * len(x_type.names) + return all( + _leaves_equal( + x if x_name is None else x[x_name], + x_type if x_name is None else x_type.fields[x_name][0], + y if y_name is None else y[y_name], + y_type if y_name is None else y_type.fields[y_name][0], + ) + for x_name, y_name in zip(x_names, y_names, strict=True) + ) + x = np.asarray(x, dtype=x_type.base) + y = np.asarray(y, dtype=y_type.base) + common = np.result_type(x.dtype, y.dtype) + return bool(np.all(x.astype(common) == y.astype(common))) + + def _element_shapes_differ(dtype1, dtype2): """Whether elements of ``dtype1`` and ``dtype2`` differ in shape, for ``isequal``. diff --git a/graphblas/core/ss/dtypes.py b/graphblas/core/ss/dtypes.py index 12f74fbce..4dafd9c82 100644 --- a/graphblas/core/ss/dtypes.py +++ b/graphblas/core/ss/dtypes.py @@ -29,7 +29,11 @@ def register_new(name, jit_c_definition, *, np_type=None): ) if not name.isidentifier(): raise ValueError(f"`name` argument must be a valid Python identifier; got: {name!r}") - if name in core.dtypes._registry or hasattr(dtypes.ss, name): + from ..operator.udt_utils import _jit_c_type_defs + + # SuiteSparse's JIT knows a type by its C name alone, so it must not be + # another type's C name either. + if name in core.dtypes._registry or hasattr(dtypes.ss, name) or name in _jit_c_type_defs: raise ValueError(f"{name!r} name for dtype is unavailable") # GxB_MAX_NAME_LEN counts the terminating NUL, so SuiteSparse keeps one fewer # character; a name that long would no longer match its typedef and would @@ -57,6 +61,7 @@ def register_new(name, jit_c_definition, *, np_type=None): ffi_new("char[]", jit_c_definition.encode()), ) check_status_carg(status, "Type", gb_obj[0]) + _jit_c_type_defs[name] = jit_c_definition # Let SuiteSparse:GraphBLAS determine the size (we gave 0 as size above) size_ptr = ffi_new("size_t*") diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index a9550ff51..c1a8cd7ea 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -539,6 +539,44 @@ def test_udt_jit_c_info_pinned_at_first_register(): assert "_PinUDT_first" in udt.jit_c_definition +@pytest.mark.skipif("not _has_jit_set") +def test_udt_layouts_under_one_name_get_their_own_jit_c_names(): + """Two UDT layouts registered under one name get distinct JIT C names. + + SuiteSparse's JIT knows a type by its C name alone, so the second layout + ran kernels compiled for the first: making an iso Vector non-iso expanded + 28-byte elements with a kernel for 6-byte ones (wrong values), or 6-byte + elements with one for 28-byte ones (heap corruption, in either order of + the tests that did this). + """ + small = dtypes.register_anonymous(np.dtype((np.int16, (1, 3, 1))), "_SameJitName") + big = dtypes.register_anonymous(np.dtype((np.int16, (2, 1, 7))), "_SameJitName") + assert small.jit_c_name == "_SameJitName" + assert big.jit_c_name != small.jit_c_name + assert big.jit_c_name in big.jit_c_definition + jit_c_control = gb.ss.config["jit_c_control"] if gb.backend == "suitesparse" else None + if jit_c_control is not None: + gb.ss.config["jit_c_control"] = "on" + try: + for udt in [small, big]: + shape = udt.np_type.shape + one = np.full(shape, 1, np.int16) + two = np.full(shape, 2, np.int16) + v = gb.Vector(udt, size=4) + v[0] = one + v[1] = one # iso + v[2] = two # no longer iso: SuiteSparse expands the iso value + values = [v[i].new().value.tolist() for i in range(3)] + assert values == [one.tolist(), one.tolist(), two.tolist()] + finally: + if jit_c_control is not None: + gb.ss.config["jit_c_control"] = jit_c_control + if gb.backend == "suitesparse": + # Nor may a typedef registered in C take a name a UDT already has. + with pytest.raises(ValueError, match="unavailable"): + dtypes.ss.register_new("_SameJitName", "typedef struct { int8_t q ; } _SameJitName ;") + + def test_dtype_to_from_string(): types = [dtypes.BOOL, dtypes.FP64] for c in string.ascii_letters: diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index d1680f34d..89dd9316b 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1,4 +1,5 @@ import itertools +import sys import numpy as np import pytest @@ -42,7 +43,7 @@ UINT32, UINT64, ) -from graphblas.exceptions import DimensionMismatch, DomainMismatch, UdfParseError +from graphblas.exceptions import DimensionMismatch, DomainMismatch, InvalidValue, UdfParseError from .conftest import shouldhave @@ -2222,13 +2223,14 @@ def test_udt_jit_typedef(): assert "double jy" in defn assert "JitTypeTest" in defn - # Array UDT - arr_dtype = np.dtype((np.float64, (7,))) - arr_udt = dtypes.register_anonymous(arr_dtype, "Vec7") + # Array UDT, of a length no other test uses: a literal can register float64 + # arrays of common lengths under another name first, in any test order. + arr_dtype = np.dtype((np.float64, (31,))) + arr_udt = dtypes.register_anonymous(arr_dtype, "Vec31") lib.GrB_Type_get_String(arr_udt._carg, buf, lib.GxB_JIT_C_DEFINITION) defn = ffi.string(buf).decode() - assert "double v [7]" in defn - assert "Vec7" in defn + assert "double v [31]" in defn + assert "Vec31" in defn # 2D array UDT mat_dtype = np.dtype((np.int32, (5, 5))) @@ -2270,20 +2272,20 @@ def test_udt_jit_op_definitions(): assert "ainv_JitOpTest" in defn assert "jp" in defn - # Array UDT JIT definitions - arr_dtype = np.dtype((np.float64, (5,))) - arr_udt = dtypes.register_anonymous(arr_dtype, "Vec5Jit") + # Array UDT JIT definitions, of a length no other test uses (see test_udt_jit_typedef) + arr_dtype = np.dtype((np.float64, (37,))) + arr_udt = dtypes.register_anonymous(arr_dtype, "Vec37Jit") typed = binary.plus[arr_udt] lib.GrB_BinaryOp_get_String(typed.gb_obj, buf, lib.GxB_JIT_C_DEFINITION) defn = ffi.string(buf).decode() - assert "plus_Vec5Jit" in defn + assert "plus_Vec37Jit" in defn assert "z->v[i] = (x->v[i]) + (y->v[i])" in defn - assert "i < 5" in defn + assert "i < 37" in defn typed = unary.ainv[arr_udt] lib.GrB_UnaryOp_get_String(typed.gb_obj, buf, lib.GxB_JIT_C_DEFINITION) defn = ffi.string(buf).decode() - assert "ainv_Vec5Jit" in defn + assert "ainv_Vec37Jit" in defn @pytest.mark.skipif("not supports_udfs") @@ -3579,10 +3581,10 @@ def test_udt_float32_overflow_warns_at_the_callers_line(): v[0] = [1e300, 1] assert [w.filename for w in record] == [__file__] assert v[0].new().value.tolist() == [np.inf, 1] - with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: - (v == 1e300).new() - assert [w.filename for w in record] == [__file__] - if supports_udfs: + if supports_udfs: # UDT eq and plus need Numba where there is no C JIT + with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: + (v == 1e300).new() + assert [w.filename for w in record] == [__file__] with pytest.warns(RuntimeWarning, match="overflow encountered in cast") as record: (v + 1e300).new() assert [w.filename for w in record] == [__file__] @@ -4078,7 +4080,14 @@ def test_udt_array_ops_on_large_elements(): Unrolled, Numba took about two minutes to type ``minus`` on this element. """ - T = dtypes.register_anonymous(np.dtype((np.float64, (64, 64))), "_LoopBig") + try: + T = dtypes.register_anonymous(np.dtype((np.float64, (64, 64))), "_LoopBig") + except InvalidValue: + if sys.platform != "win32": + raise + # SuiteSparse:GraphBLAS built by MSVC (no variable-length arrays) refuses a + # UDT larger than 1024 bytes (GB_VLA_MAXSIZE); this one is 32768. + pytest.skip("SuiteSparse:GraphBLAS limits a UDT to 1024 bytes on Windows") S = dtypes.register_anonymous(np.dtype((np.float64, (64, 1))), "_LoopBigCol") x = np.arange(64 * 64, dtype=np.float64).reshape(64, 64) col = np.arange(64, dtype=np.float64).reshape(64, 1) @@ -4498,8 +4507,10 @@ def test_udt_literal_converts_only_when_exact(): assert r.apply(second, right=(1, 0.1)).new()[0].new().value.tolist() == (1, np.float32(0.1)) assert r.apply(second, right={"ex_i": 2, "ex_f": 0.5}).new()[0].new().value.tolist() == (2, 0.5) assert r.apply(second, right=np.int8(2)).new()[0].new().value.tolist() == (2, 2.0) - # A sequence holding numpy values is strong, as the array numpy makes of it. - assert binary.plus(v, (np.int8(1), 2, 3)).new().dtype.np_type == np.dtype((np.int64, (3,))) + # A sequence holding numpy values is strong, as the array numpy makes of it + # (int64, or int32 on Windows with numpy 1). + made = np.asarray((np.int8(1), 2, 3)).dtype + assert binary.plus(v, (np.int8(1), 2, 3)).new().dtype.np_type == np.dtype((made, (3,))) # The error names the literal's type; a layout's name is whatever UDT # registered it first, so only the built-in scalar types are pinned here. for expr, typed_as in [ @@ -4596,7 +4607,7 @@ def test_udt_sequence_and_array_literals_type_like_numbers(): # Strong: an array is its own dtype, in its shape, or broadcast to the UDT's. for literal, expected in [ (np.array([0.5, 1, 2]), f64s), - (np.array([1, 2, 3]), np.dtype((np.int64, (3,)))), + (np.array([1, 2, 3]), np.dtype((np.array([1]).dtype, (3,)))), # int32 on numpy 1 Windows (np.array([1], np.int8), int8s.np_type), (np.array([0.5]), f64s), (np.ones((1, 3)), np.dtype((np.float64, (1, 3)))), diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index 6961300dd..6c14e70df 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -707,6 +707,58 @@ def test_record_from_dict(): assert s == (1, (2, 3)) +@pytest.mark.skipif("not gb.core._has_numba") +def test_udt_eq_without_numba_matches_binary_eq(monkeypatch): + """Without Numba, ``==`` on a UDT Scalar compares in numpy, as ``binary.eq`` would. + + ``binary.eq`` on a UDT needs Numba; ``==`` on a Scalar used to compare in + Python, so it works without Numba still. Here both ways must agree. + """ + rec = dtypes.register_anonymous(np.dtype([("x", np.int64), ("y", np.float64)], align=True)) + arr = dtypes.register_anonymous(np.dtype((np.int8, (3,)))) + f32 = dtypes.register_anonymous(np.dtype((np.float32, (2,)))) + r = Scalar.from_value((1, 2.5), rec) + a = Scalar.from_value([1, 1, 1], arr) + f = Scalar.from_value([0.1, 0.1], f32) + cases = [ + (r, (1, 2.5)), + (r, (1, 2)), + (r, {"x": 1, "y": 2.5}), + (a, 1), + (a, [1, 1, 1]), + (a, [1]), + (a, [1, 2, 1]), + (a, 0.5), + (a, 300), + (a, 2**63), + (a, Scalar.from_value([1, 1, 1], arr)), + (f, 0.1), + (f, np.float64(0.1)), + ] + # UDTs that binary.eq does not pair raise its KeyError both ways. + other_names = dtypes.register_anonymous( + np.dtype([("p", np.int64), ("q", np.float64)], align=True) + ) + arr2 = dtypes.register_anonymous(np.dtype((np.int8, (2,)))) + unpaired = [ + (r, Scalar.from_value((1, 2.5), other_names)), + (a, Scalar.from_value([1, 1], arr2)), + (r, a), + ] + + def compare_all(): + for s, other in unpaired: + with pytest.raises(KeyError, match="does not work with"): + s == other # noqa: B015 + return [(s == other, s != other) for s, other in cases] + + with_numba = compare_all() + monkeypatch.setattr(gb.core.scalar, "_has_numba", False) + without_numba = compare_all() + assert without_numba == with_numba + assert with_numba[:3] == [(True, False), (False, True), (True, False)] + + def test_get(s): assert s.get() == 5 assert s.get("mittens") == 5 diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 936b6b036..34d1ae353 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -1538,7 +1538,10 @@ def test_python_number_is_weak_beside_builtin_dtypes(): assert (Scalar.from_value(0.5, dtype=dtypes.FP32) * 0.5).new().dtype == dtypes.FP32 A = Matrix.from_coo([0], [0], [100], dtype=dtypes.INT8) assert (A * 2).dtype == dtypes.INT8 - for expr in [lambda: i8 + 300, lambda: u8 + -1, lambda: i8 // 300, lambda: s8 * 300]: + exprs = [lambda: i8 + 300, lambda: u8 + -1, lambda: s8 * 300] + if supports_udfs: # integer floordiv is a Numba op + exprs.append(lambda: i8 // 300) + for expr in exprs: with pytest.raises(OverflowError, match="out of bounds"): expr() # truediv gives FP64 whatever the int, so an int out of range does not raise: From b6fdcec2a9a5e1fb99d4c24d6d26750b0aebf7f9 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 9 Oct 2026 00:24:02 -0500 Subject: [PATCH 10/54] Spell out complex * and / for UDTs; Numba 0.68 follows Python 3.14 From Python 3.14, CPython recovers infinities that the textbook complex formulas leave as nan (C99 Annex G), and Numba 0.68 follows it there: 1j * (inf+infj) is -inf+infj and (inf+nanj) / (1+1j) is inf-infj, where numpy and the C JIT kernel give nan+nanj. The UDT cfunc used Numba's operators, so it disagreed with the C JIT kernel on CI's macOS job with Python 3.14 and Numba 0.68. The cfunc now spells out the formulas the C kernel uses (the textbook product, and Smith's division as CPython divided before 3.14), so the two agree on any Python and Numba. The user guide no longer calls a complex quotient numpy's: numpy multiplies by the divisor's reciprocal, so the last bit can differ, and numpy's overflows where the quotient is finite. --- docs/user_guide/udt.rst | 6 +- graphblas/core/operator/udt_utils.py | 101 ++++++++++++++++++++------- graphblas/tests/test_op.py | 33 +++++++++ graphblas/tests/test_ssjit.py | 13 ++-- 4 files changed, 121 insertions(+), 32 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 25a41ad1d..7b731fb44 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -205,8 +205,10 @@ built-in op computes it, and the Numba cfunc and a C JIT kernel give the same bits: ``min`` and ``max`` ignore NaN (C's ``fmin`` and ``fmax``) and resolve a tie between ``-0.0`` and ``0.0`` to ``-0.0`` for ``min`` and ``0.0`` for ``max``, ``floordiv`` is numpy's, with 0 for an integer divided by zero, and -``truediv`` is numpy's (Smith's method for complex numbers), with numpy's -infinities for a zero divisor. A few edge cases can differ from built-in +``truediv`` is numpy's, with numpy's infinities for a zero divisor. A complex +quotient takes Smith's method, as numpy's does, but divides where numpy +multiplies by a reciprocal, so it can differ from numpy's in the last bit (and +stays finite where numpy's overflows). A few edge cases can differ from built-in vectors, where the UDT gives numpy's answer: ``INT64_MIN // -1`` is ``INT64_MIN`` (an ``INT64`` vector gives 0), a complex product or quotient with an infinite or NaN part, a complex quotient with a ``-0.0`` divisor, and the diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index a0ca8856e..30385b354 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -71,10 +71,13 @@ def _compile_codegen(src, *, func_name, source_label, extra_ns=None): # Generated code never calls Python's ``min`` / ``max`` (still reachable as # builtins): they order a NaN by argument position, where SuiteSparse's # ``fmin`` ignores it, so :func:`_minmax_expr` spells the comparison out, - # with ``signbit`` for the signed-zero tie. + # with ``signbit`` for the signed-zero tie. Complex ``*`` and ``/`` call + # :func:`_complex_times` and :func:`_complex_truediv` (see _expr_binary). namespace = {"abs": abs, "np": np, "signbit": np.signbit} if _has_numba: namespace["numba"] = numba + namespace["_complex_times"] = _complex_times + namespace["_complex_truediv"] = _complex_truediv if extra_ns: namespace.update(extra_ns) exec(code, namespace) @@ -979,7 +982,7 @@ def _expr_binary(py_op, x_expr, y_expr, x_dtype, y_dtype): ``x_expr`` and ``y_expr`` are scalars, never arrays. ``x_dtype`` and ``y_dtype`` are the numpy dtypes the op computes in for each operand: the built-in op's input types, which the operands are converted to - first. Three ops consult them: + first. These ops consult them: - ``min`` / ``max`` follow SuiteSparse's own ``GrB_MIN_FP64``, which is C99 ``fmin`` and ignores a NaN operand rather than ordering or @@ -991,11 +994,12 @@ def _expr_binary(py_op, x_expr, y_expr, x_dtype, y_dtype): unrepresentable quotient, while numpy wraps to ``INT64_MIN``. Since ``a // -1`` is exactly ``-a``, routing that divisor through negation (which wraps) reaches numpy's answer without the trap. - - ``/`` where either side is complex. Numba's complex division raises - ``ZeroDivisionError`` unconditionally, outside the error model's - control, so a zero divisor left the element unwritten while the C - kernel returned numpy's infinities. Spelling out the zero case here - makes both paths match numpy. + - ``*`` and ``/`` where either side is complex go through + :func:`_complex_times` and :func:`_complex_truediv`, the formulas + the C kernel spells out. Numba's own operators are CPython's, so + its ``/`` raises ``ZeroDivisionError`` on a zero divisor, and from + Python 3.14 (Numba 0.68) both recover infinities that numpy and the + C kernel leave as ``nan``. ``/`` on two integers needs nothing: the built-in typing makes its result float64, as ``truediv`` on two INT64 vectors is FP64, so @@ -1005,13 +1009,60 @@ def _expr_binary(py_op, x_expr, y_expr, x_dtype, y_dtype): return _minmax_expr(py_op, x_expr, y_expr, x_dtype, y_dtype) if py_op == "//" and _is_signed_int(x_dtype) and _is_signed_int(y_dtype): return f"(-{x_expr} if {y_expr} == -1 else {x_expr} // {y_expr})" + if py_op == "*" and "c" in (x_dtype.kind, y_dtype.kind): + return f"_complex_times({x_expr}, {y_expr})" if py_op == "/" and "c" in (x_dtype.kind, y_dtype.kind): - return ( - f"({x_expr} / {y_expr} if {y_expr} != 0 " - f"else complex({x_expr}.real / 0.0, {x_expr}.imag / 0.0))" - ) + return f"_complex_truediv({x_expr}, {y_expr})" return f"{x_expr} {py_op} {y_expr}" + @numba.njit(error_model="numpy") + def _complex_times(x, y): + """Multiply two complex numbers by the textbook formula, as numpy does. + + Numba's own ``*`` is CPython's, which from Python 3.14 (Numba 0.68) + adds C99 Annex G's recovery of infinities: ``1j * (inf+infj)`` is + ``-inf+infj`` there and ``nan+nanj`` in numpy. This is + :func:`_c_complex_times_stmt`, product for product, so the cfunc and + the C kernel agree whatever the Python and Numba versions. The parts + keep the operands' precision; the complex128 result holds them + exactly, for the caller to store in the field's type. + """ + ar = x.real + ai = x.imag + br = y.real + bi = y.imag + return complex(ar * br - ai * bi, ar * bi + ai * br) + + @numba.njit(error_model="numpy") + def _complex_truediv(x, y): + """Divide two complex numbers by Smith's method, as the C kernel does. + + This is CPython's ``_Py_c_quot`` before Python 3.14 (scale by the + larger part of the divisor), statement for statement as + :func:`_c_complex_truediv_stmt` spells it. Numba's own ``/`` raises + ``ZeroDivisionError`` on a zero divisor, outside the error model's + control, and from Python 3.14 (Numba 0.68) recovers infinities: + ``(inf+nanj) / (1+1j)`` is ``inf-infj`` there and ``nan+nanj`` in + numpy. A zero divisor divides the parts by ``0.0`` as reals, for + numpy's infinities. The result is complex128, as in + :func:`_complex_times`. + """ + ar = x.real + ai = x.imag + br = y.real + bi = y.imag + if br == 0 and bi == 0: + return complex(ar / 0.0, ai / 0.0) + if abs(br) >= abs(bi): + r = bi / br + d = br + bi * r + return complex((ar + ai * r) / d, (ai - ar * r) / d) + if bi == 0: # the real part is NaN + return complex(np.nan, np.nan) + r = br / bi + d = br * r + bi + return complex((ar * r + ai) / d, (ai * r - ar) / d) + def _minmax_expr(py_op, x_expr, y_expr, x_dtype, y_dtype): """Return a Python expression for ``min`` / ``max`` matching C ``fmin``. @@ -1686,16 +1737,14 @@ def _c_complex_truediv_stmt(target, lhs, rhs, field_dtype): with ``-fcx-limited-range`` under GCC, which replaces C99 Annex G division with the naive formula: it squares the divisor's parts, so ``1e20 / 1e20`` in complex64 gave ``nan`` and ``(1+1j) / (1e160+1e160j)`` gave ``0`` on - Linux. This block is Numba's complex division, which is CPython's - ``_Py_c_quot`` (Smith's method: scale by the larger part of the divisor), - computed in the field's own precision with real arithmetic only, so no - flag changes it. Each product is its own statement: clang contracts a - multiply and an add within one expression into a fused multiply-add, - which rounds once where Numba rounds twice, and GCC under ``-std=c11`` - does not contract at all. A zero divisor divides the parts by ``0.0`` as - reals, for numpy's infinities, as ``_expr_binary`` spells it for the - cfunc; ``rhs == 0`` on a ``_Complex`` compares both parts, mirroring the - cfunc's ``y != 0``. + Linux. This block is the cfunc's :func:`_complex_truediv`, statement for + statement: CPython's ``_Py_c_quot`` before Python 3.14 (Smith's method: + scale by the larger part of the divisor), computed in the field's own + precision with real arithmetic only, so no flag changes it. Each product + is its own statement: clang contracts a multiply and an add within one + expression into a fused multiply-add, which rounds once where Numba + rounds twice, and GCC under ``-std=c11`` does not contract at all. A zero + divisor divides the parts by ``0.0`` as reals, for numpy's infinities. The constructor is SuiteSparse's ``GB_CMPLX32``/``GB_CMPLX64`` rather than C11's ``CMPLXF``/``CMPLX``, which macOS does not define. @@ -1725,13 +1774,15 @@ def _c_complex_truediv_stmt(target, lhs, rhs, field_dtype): def _c_complex_times_stmt(target, lhs, rhs, field_dtype): """Return a C block multiplying two complex fields into ``target``, as the Numba cfunc does. - Numba and numpy multiply complex numbers by the textbook formula. C99 - Annex G adds a recovery step that turns some of its ``nan`` results into + numpy multiplies complex numbers by the textbook formula. C99 Annex G + adds a recovery step that turns some of its ``nan`` results into infinities, and whether C's ``*`` on ``_Complex`` takes it depends on the compiler: GCC under SuiteSparse's ``-fcx-limited-range`` does not, clang does. So ``1j * (inf+infj)`` was ``nan+nanj`` on Linux and the cfunc but - ``-inf+infj`` from a C kernel on macOS. This block is the textbook formula - in real arithmetic, one product per statement so clang cannot contract it + ``-inf+infj`` from a C kernel on macOS. CPython takes it too from 3.14, + and Numba follows CPython, so the cfunc does not use Numba's ``*`` + either: this block is :func:`_complex_times`, the textbook formula in + real arithmetic, one product per statement so clang cannot contract it into a fused multiply-add (see :func:`_c_complex_truediv_stmt`). """ ctype, cmplx, creal, cimag, _, _ = _c_complex_names(field_dtype) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 89dd9316b..4ed64a54c 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1764,6 +1764,39 @@ def test_udt_complex_truediv_by_zero(udt_op_path): np.testing.assert_array_equal(got, expected, err_msg=udt_op_path) +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("not dtypes._supports_complex") +def test_udt_complex_times_truediv_with_infinities_as_numpy(udt_op_path): + """Complex ``*`` and ``/`` on a UDT field answer as numpy does at infinities. + + From Python 3.14, CPython recovers infinities that the textbook formulas + leave as ``nan`` (C99 Annex G), so ``1j * (inf+infj)`` is ``-inf+infj`` + there and ``nan+nanj`` in numpy. Numba 0.68 follows CPython, so a cfunc + built on Numba's operators answered as Python and the C JIT kernel as + numpy. Both now spell numpy's formulas out, on any Python and Numba. + """ + inf, nan = np.inf, np.nan + # The last pair is finite, with an exact product and quotient: numpy + # divides by multiplying with ``1 / denominator``, which can round the last + # bit differently from CPython's division, which the kernels use. + xs = [1j, -2.5j, complex(nan, inf), complex(inf, nan), 1 + 1j, complex(inf, 0.0), -1 + 7j] + ys = [complex(inf, inf), complex(inf, -inf), 2 + 1j, 1 + 1j, complex(inf, inf), 1j, 1 + 1j] + rec = dtypes.register_anonymous( + np.dtype([("cxi_s", np.complex64), ("cxi_d", np.complex128)], align=True), "_CxInfRec" + ) + arr = dtypes.register_anonymous(np.dtype((np.complex64, (3,))), "_CxInfArr3") + for udt, leaves in [(rec, [0, 1]), (arr, [0, 2])]: + v, w = _udt_vectors(udt, xs, ys) + for gb_op, reference in [(binary.times, np.multiply), (binary.truediv, np.true_divide)]: + result = gb_op(v & w).new() + for leaf in leaves: + got = np.array([result[i].new().value[leaf] for i in range(len(xs))]) + with np.errstate(invalid="ignore"): + expected = reference(np.array(xs, got.dtype), np.array(ys, got.dtype)) + msg = f"{gb_op.name} {got.dtype} {udt_op_path}" + np.testing.assert_array_equal(got, expected, err_msg=msg) + + @pytest.mark.skipif("not supports_udfs") # 136-byte UDT, which SS < 9 rejects; see test_udt_large_array. @pytest.mark.skipif( diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 228d69b91..097fc6866 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -1187,8 +1187,10 @@ def test_complex_truediv_udt_jit_matches_cfunc_at_extremes(): SuiteSparse compiles C JIT kernels with ``-fcx-limited-range`` under GCC, which turns C's ``_Complex`` division into the naive formula: complex64 ``1e20 / 1e20`` came out ``nan`` and ``(1+1j) / (1e160+1e160j)`` came out - ``0`` on Linux. The kernel divides with real arithmetic instead, by the - algorithm Numba (and CPython) use, so no flag changes it. + ``0`` on Linux. The kernel divides with real arithmetic instead, by + CPython's algorithm before 3.14, so no flag changes it; the cfunc spells + out the same steps, since from Python 3.14 Numba's own ``/`` recovers + infinities as CPython's does (``(inf+nanj) / (1+1j)`` is ``inf-infj``). """ if not _has_jit_set: pytest.skip("jit_c_source introspection requires SuiteSparse:GraphBLAS >= 9") @@ -1236,12 +1238,13 @@ def test_complex_truediv_udt_jit_matches_cfunc_at_extremes(): def test_complex_times_udt_jit_matches_cfunc_with_infinities(): """Complex multiplication in the C JIT kernel is the cfunc's, bit for bit. - Numba and numpy multiply by the textbook formula; clang's ``_Complex`` - multiply adds C99 Annex G's recovery, so ``1j * (inf+infj)`` came out + numpy multiplies by the textbook formula; clang's ``_Complex`` multiply + adds C99 Annex G's recovery, so ``1j * (inf+infj)`` came out ``-inf+infj`` from a C kernel on macOS and ``nan+nanj`` from the cfunc (GCC's ``-fcx-limited-range`` drops the recovery, so Linux agreed). The kernel multiplies with real arithmetic instead, for record fields and - array elements alike. + array elements alike. CPython adds the recovery from 3.14 and Numba 0.68 + follows it, so the cfunc spells the formula out too. """ if not _has_jit_set: pytest.skip("jit_c_source introspection requires SuiteSparse:GraphBLAS >= 9") From 428d00fbc6341d6980f85f914e25a780c8049783 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 14 Oct 2026 12:40:07 -0500 Subject: [PATCH 11/54] Move UDT and UDF logic out of the general modules Code moves, plus import changes and a few small helpers, so scalar, vector, matrix and base read mostly as the built-in path. - core/udt.py (new): rules for UDT elements: layout, op typing, store casts, the ewise_add rule, equality without Numba, and element values as bytes. The type rules moved here from udt_utils, which now only generates code and is renamed udt_codegen. - core/literals.py (new): NEP 50 typing of literal operands, for built-in and UDT types (moved from scalar.py and udt_utils). - operator/udf.py (new): the Numba-only UDF-on-UDT block from operator/base.py, plus _bool_to_int8. - scalar._union_default replaces the ewise_union default block, which was copied six times across Scalar, Vector and Matrix. - New helpers pulled out of scalar.py's inline code: _scalar_ewise_add_converted, _udt_eq_without_numba, _element_bytes and _element_value. core/udt.py imports operator, base and scalar names inside functions, after the built-in early returns, so built-in calls pay nothing extra. --- graphblas/core/base.py | 92 +- graphblas/core/dtypes.py | 6 +- graphblas/core/infixmethods.py | 18 +- graphblas/core/literals.py | 615 +++++++++++++ graphblas/core/matrix.py | 63 +- graphblas/core/operator/base.py | 653 +------------- graphblas/core/operator/binary.py | 22 +- graphblas/core/operator/indexbinary.py | 2 +- graphblas/core/operator/indexunary.py | 2 +- graphblas/core/operator/select.py | 2 +- graphblas/core/operator/udf.py | 665 ++++++++++++++ .../operator/{udt_utils.py => udt_codegen.py} | 605 +------------ graphblas/core/operator/unary.py | 10 +- graphblas/core/scalar.py | 801 ++--------------- graphblas/core/ss/dtypes.py | 2 +- graphblas/core/ss/jit_config.py | 2 +- graphblas/core/udt.py | 809 ++++++++++++++++++ graphblas/core/vector.py | 59 +- graphblas/tests/test_dtype.py | 2 +- graphblas/tests/test_op.py | 28 +- graphblas/tests/test_ssjit.py | 2 +- pyproject.toml | 2 +- 22 files changed, 2247 insertions(+), 2215 deletions(-) create mode 100644 graphblas/core/literals.py create mode 100644 graphblas/core/operator/udf.py rename graphblas/core/operator/{udt_utils.py => udt_codegen.py} (72%) create mode 100644 graphblas/core/udt.py diff --git a/graphblas/core/base.py b/graphblas/core/base.py index 796771595..d4e76e159 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -7,8 +7,10 @@ from . import NULL from .descriptor import lookup as descriptor_lookup from .expr import AmbiguousAssignOrExtract, Updater +from .literals import _literal_store_dtype from .mask import Mask from .operator import UNKNOWN_OPCLASS, binary_from_string, find_opclass, get_typed_op +from .udt import _fused_store_expr, _scalar_accum_op, _store_cast_op from .utils import _Pointer, libget, output_type _recorder = ContextVar("recorder") @@ -162,94 +164,6 @@ def _expect_op(self, op, values, *, within, **kwargs): raise TypeError(message) from None -def _store_cast_op(result_dtype, output_dtype): - """Return the UnaryOp that converts a result for an object of ``output_dtype``, or ``None``. - - GraphBLAS casts between built-in types on store (C's casts, with a float - saturating into an integer type), but it cannot cast to or from a UDT; - left to GraphBLAS, that is a bare GrB_DOMAIN_MISMATCH or, for a single - element, bytes read as the wrong type. So when a UDT is involved, the - result is stored by applying the op returned here, which casts each - element or record leaf the way GraphBLAS casts built-in types, between - UDTs whose layouts correspond (``udt_utils._udt_cast_error``). Any other - pair raises ``DomainMismatch``. ``None`` means GraphBLAS stores the result - as it is. This is the one place that decides what may be stored where. - """ - if result_dtype is output_dtype or not (result_dtype._is_udt or output_dtype._is_udt): - return None - from . import _has_numba - - reason = "UDTs do not cast" - if _has_numba and result_dtype._is_udt and output_dtype._is_udt: - from .operator.udt_utils import _udt_cast_error - - if (reason := _udt_cast_error(result_dtype, output_dtype)) is None: - from .operator.unary import _udt_cast_op - - return _udt_cast_op(result_dtype, output_dtype) - elif _has_numba: - reason = "a UDT casts only to another UDT" - from ..exceptions import DomainMismatch - - raise DomainMismatch( - f"cannot store a result of type {result_dtype} in an object of type " - f"{output_dtype}: {reason}. Use .new() for an object of the result's type, or " - "compute the result in the type you want to store." - ) - - -def _fused_store_expr(expr, output_dtype): - """Return ``expr`` computing straight into ``output_dtype``, or raise ``DomainMismatch``. - - A result of another UDT type is stored by the op that computes it, which - casts each element as it writes it (``udt_utils.fused_store_op``), so no - temporary of the result's own type is made: data GraphBLAS works on is - often too large to hold twice. That works for an element-wise - computation by a built-in op that lifts to UDTs (``apply``, - ``ewise_mult``, ``ewise_union``). ``ewise_add`` copies an entry present in - only one input without the op, and a semiring or reduction would add in - the narrower type, so those and any other computation raise. - """ - import copy - - from .operator.udt_utils import fused_store_op - - op, cfunc_name = expr.op, expr.cfunc_name - if op is not None and op.opclass == "Monoid": - # Element-wise, a Monoid is its BinaryOp. - op, cfunc_name = op.binaryop, cfunc_name.replace("_Monoid", "_BinaryOp") - fused = None - if any(k in cfunc_name for k in ("_apply", "eWiseMult", "eWiseUnion")): - fused = fused_store_op(op, output_dtype) - if fused is None: - from ..exceptions import DomainMismatch - - method = "extraction" if expr.method_name == "__getitem__" else expr.method_name - what = method if op is None else f"{method} with {op.name}" - raise DomainMismatch( - f"cannot store a result of type {expr.dtype} in an object of type " - f"{output_dtype}: GraphBLAS cannot cast a UDT, and only a built-in op that lifts " - "to UDTs, in apply, ewise_mult or ewise_union, can write another type as it " - f"computes ({what} cannot), so this would need a temporary copy of the whole " - "result. Make it explicitly by storing the computed result, as in " - "`obj << expr.new()`, which casts it in one more pass." - ) - rv = copy.copy(expr) - rv.op, rv.cfunc_name, rv.dtype = fused, cfunc_name, output_dtype - return rv - - -def _scalar_accum_op(scalar, accum): - """Return the op to accumulate into ``scalar`` with, as ``scalar.ewise_add`` runs it. - - ``__call__`` types ``accum`` on the object's dtype. For a UDT Scalar, that - would convert a literal into the UDT (``s += 0.5``), where a Vector - computes ``s + 0.5`` in its own type and casts as it stores; the untyped - op does the same here. - """ - return accum.parent if scalar.dtype._is_udt else accum - - AmbiguousAssignOrExtract._expect_op = _expect_op AmbiguousAssignOrExtract._expect_type = _expect_type @@ -509,7 +423,7 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * # w << (v & v) expr = expr._to_expr() else: - from .scalar import Scalar, _literal_store_dtype + from .scalar import Scalar if type(expr) is Scalar: scalar = expr diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 4800efaf9..71cce87e0 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -117,7 +117,7 @@ def _get_udt_c_type_info(datatype): word or isn't an identifier. Supports record UDTs (including nested records) and array UDTs (e.g., ``FP64[3]``). """ - from .operator.udt_utils import _udt_c_typedef + from .operator.udt_codegen import _udt_c_typedef if datatype.np_type is None: return None @@ -143,7 +143,7 @@ def _set_udt_jit_c_definition(datatype): Supports both record UDTs (``{'x': float, 'y': float}``) and array UDTs (``np.dtype((np.float64, (3,)))``). """ - from .operator.udt_utils import ( + from .operator.udt_codegen import ( _has_jit_set, _jit_c_type_defs, _synthetic_udt_counter, @@ -578,7 +578,7 @@ def unify(type1, type2, *, is_left_scalar=False, is_right_scalar=False): unify(FP32, INT32) -> FP64 ``is_left_scalar`` and ``is_right_scalar`` no longer change the result. A - Python number is typed before it gets here (``scalar._literal_operand``), and + Python number is typed before it gets here (``literals._literal_operand``), and a typed scalar promotes as any operand does, as numpy 2 promotes it. They used to stand the scalar in as a 0-d array of value 0, which numpy 1 typed by value, so ``int8_vec + 300`` wrapped there. diff --git a/graphblas/core/infixmethods.py b/graphblas/core/infixmethods.py index b0fbf0d3d..53f9d76f0 100644 --- a/graphblas/core/infixmethods.py +++ b/graphblas/core/infixmethods.py @@ -2,26 +2,12 @@ from ..dtypes import BOOL from .infix import MatrixInfixExpr, ScalarInfixExpr, VectorInfixExpr from .matrix import Matrix, MatrixExpression, MatrixIndexExpr, TransposedMatrix -from .operator import get_typed_op -from .scalar import Scalar, ScalarExpression, ScalarIndexExpr, _ewise_add_needs_cast, _plus_zero +from .scalar import Scalar, ScalarExpression, ScalarIndexExpr +from .udt import _ewise_add_casts_udt, _plus_zero from .utils import output_type from .vector import Vector, VectorExpression, VectorIndexExpr -def _ewise_add_casts_udt(left, right, op): - """Whether ``left.ewise_add(right, op)`` would need GraphBLAS to cast a UDT.""" - try: - left_dtype, right_dtype = left.dtype, right.dtype - except TypeError: - # An infix operand that has no dtype here (``x & y`` of non-BOOL); - # ewise_add reports it as before. - return False - if not (left_dtype._is_udt or right_dtype._is_udt): - return False - typed_op = get_typed_op(op, left_dtype, right_dtype, kind="binary") - return _ewise_add_needs_cast(typed_op, left_dtype, right_dtype) - - def call_op(self, other, op, *, outer=False, union=False): type1 = output_type(self) type2 = output_type(other) diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py new file mode 100644 index 000000000..9392cbca6 --- /dev/null +++ b/graphblas/core/literals.py @@ -0,0 +1,615 @@ +"""How python-graphblas types a literal operand, such as the ``1`` in ``v + 1``. + +A Python number is weak, as numpy 2 types it (NEP 50): it takes the dtype of +the operand beside it when that dtype's kind holds it, so ``int8_vec + 1`` is +INT8 and ``int8_vec * 0.5`` is FP64. A numpy value keeps its own dtype. Beside +a UDT, the same rules apply leaf by leaf. :func:`_literal_operand` decides for +each operand; a literal that must become an element of a UDT is checked by +:func:`_check_literal_fits`. + +The module also types a literal stored into an object +(:func:`_literal_store_dtype`) and a fill value +(:func:`_weak_builtin_literal_dtype` with no op). A literal that converts into +a UDT without fitting it is deprecated in :func:`_refuse_unless_exact` (MAINT). +""" + +import math +import struct +import warnings +from functools import lru_cache + +import numpy as np + +from ..dtypes import FP64, lookup_dtype +from . import utils +from .dtypes import _view_if_same_layout +from .operator import IndexUnaryOp, SelectOp, TypedOpBase, binary_from_string +from .operator.binary import BinaryOp +from .operator.monoid import Monoid +from .udt import ( + BUILTIN_UDT_BINARY_OPS, + _casts_safely_into, + _dict_to_record, + _leaf_arrays, + _promoted_record, + _strip_leading_ones, +) + + +def _literal_operand(dtype, value, op): + """Return literal ``value`` and the dtype to make it, as ``op``'s operand beside ``dtype``. + + A dtype of ``None`` means the value's own dtype. This is where + python-graphblas decides how a literal operand is typed; beside a built-in + dtype, see :func:`_weak_builtin_literal_dtype`. + + Beside a UDT, a literal becomes an element of that UDT, as it must for a + user's op, when its type allows it (:func:`_check_literal_fits`). With a + built-in op that lifts to UDTs, a literal is an operand in its own right, + typed as numpy 2 types it (NEP 50), by :func:`_literal_type`: + + - a numpy scalar or array keeps its dtype (strong); an array's type is the + array layout of its dtype (:func:`_array_literal_udt`); + - a Python ``bool``, ``int``, ``float`` or ``complex`` is weak: in each + element or leaf it takes that element's dtype when it fits the kind, so + ``int8_udt + 1`` stays ``int8_udt`` and ``int8_udt += 1`` can be stored, + while ``int8_udt * 2.5`` has float64 elements (see + :func:`_weak_literal_udt`), where converting it to the UDT would + multiply by 2; + - a tuple, list or dict of Python numbers, laid out like an element, is + weak in the same way, leaf by leaf: ``int8_udt + (1, 2, 3)`` stays + ``int8_udt`` and ``+ (0.5, 1.5, 2.5)`` has float64 elements. + + ``eq`` and ``ne`` type a literal the same way, as numpy compares, so + ``fp32_udt == 0.1`` compares in float32, as ``fp32_vec == 0.1`` does. + + A comparison keeps no result type, so a Python int beyond an integer + element's range is compared exactly, as numpy 2 compares it, as infinity + of its sign: every element compares with that infinity as with the int, + while a type that held both would round them (in FP64, the INT64 + ``2**63 - 1`` and the int ``2**63`` are equal). So ``int64_vec < 2**63`` is + all True and ``int8_vec == 300`` all False. Beside a UDT, where only ``eq`` + and ``ne`` compare, a literal with such an int in an integer leaf equals no + element: ``int8_udt == 300`` is False, not an error. + """ + if getattr(op, "is_positional", False): + return value, None + if not dtype._is_udt: + literal_dtype = _weak_builtin_literal_dtype(dtype, value, op) + if ( + literal_dtype is not dtype + and literal_dtype is not None + and type(value) is int + and dtype.np_type.kind in "iu" + ): + # Only an int beyond dtype's range takes another type here, typed so + # by a comparison or by an op that gives another type (truediv). + if isinstance(op, str): + op = binary_from_string(op) + if op.name in _COMPARISONS: + return (math.inf if value > 0 else -math.inf), FP64 + return value, literal_dtype + if getattr(op, "_udt_types", True) is None: + # An op that does not take UDTs at all (``lt``) reports that itself. + return value, dtype + if isinstance(op, Monoid) and not op._anonymous: + # Element-wise, a Monoid is its BinaryOp (``get_typed_op`` types it so). + op = op.binaryop + if type(op) is BinaryOp and not op._anonymous: + is_comparison = op.name in ("eq", "ne") + if is_comparison or op.name in BUILTIN_UDT_BINARY_OPS: + # truediv gives float leaves whatever the int, so, as a comparison, it + # types an int out of an integer leaf's range exactly instead of raising. + exact = is_comparison or op.name == "truediv" + literal_type = _literal_type(dtype, value, exact=exact) + if literal_type is not None: + if is_comparison and not isinstance(value, (np.generic, np.ndarray)): + np_type = dtype.np_type + parts = _dict_to_record(np_type, value) if isinstance(value, dict) else value + beyond = _int_beyond_leaf(np_type, parts) + if beyond is not None: + # As infinity in every leaf, it still equals no element. + value = math.inf if beyond > 0 else -math.inf + literal_type = _literal_type(dtype, value) + return value, literal_type + _check_literal_fits(dtype, value) + return value, dtype + + +# 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)]) + + +# The types a Python float or complex can overflow. +_SINGLE_PRECISION = {np.dtype(np.float32), np.dtype(np.complex64)} + + +def _weak_builtin_literal_dtype(dtype, value, op=None): + """Return the dtype of Python number ``value`` beside built-in ``dtype``, or ``None``. + + A Python ``bool``, ``int``, ``float`` or ``complex`` is weak, as numpy 2 + types it (NEP 50), by the table UDT leaves use + (:func:`_weak_literal_leaf`): ``int8_vec + 1`` is INT8, + ``fp32_vec * 0.5`` is FP32, ``int8_vec * 0.5`` is FP64, and + ``int8_vec + 300`` raises ``OverflowError``. python-graphblas applies the + table itself, so numpy 1 gives the same answers. ``None`` keeps the value's + own dtype (strong): for anything but a Python number, for a typed op, which + fixes the operand types itself, and for the thunk of an index op that does + not compare it with the values, such as a row index. + + A comparison keeps no result type, so an int beyond ``dtype``'s range + takes a type that holds it rather than raise (:func:`_literal_operand` then + compares it as infinity). ``op`` of ``None``, for a fill value, does the + same, so the value is not lost. + """ + kind = _LITERAL_KINDS.get(type(value)) + if kind is None: + return None + if isinstance(op, str): + try: + op = binary_from_string(op) + except Exception: + return None + exact = op is None + if op is not None: + if isinstance(op, TypedOpBase): + return None + exact = getattr(op, "name", None) in _COMPARISONS and not op._anonymous + if isinstance(op, (IndexUnaryOp, SelectOp)) and not exact: + return None + try: + np_type = _weak_literal_leaf([value], kind, dtype.np_type, exact=exact) + except OverflowError: + if exact or not _gives_another_type(op, dtype): + raise + # An op whose result is not the operand's type has no narrow result to + # keep, so an int out of range is typed exactly, as for a comparison: + # int8_vec / 300 is FP64, as numpy divides integers in float64. + np_type = _weak_literal_leaf([value], kind, dtype.np_type, exact=True) + if np_type == dtype.np_type: + rv = dtype + else: + try: + rv = lookup_dtype(np_type) + except ValueError: # A complex literal where the backend has no complex types + return None + if kind != "b" and rv.np_type in _SINGLE_PRECISION: + parts = (value.real, value.imag) if kind == "c" else (float(value),) + try: + for part in parts: + struct.pack(" len(expected) and shape[0] == 1: - shape = shape[1:] - try: - return np.broadcast_shapes(shape, expected) == expected - except ValueError: - return False - - def _udf_probe_value(dtype): - """Build a stand-in operand of ``dtype`` for :func:`_run_udf_probe`. - - Only the shape of the UDF's result is read, so any values it accepts - will do. - """ - if dtype._is_udt: - np_type = dtype.np_type - if np_type.subdtype is None: - return np.ones(1, dtype=np_type)[0] - base_np_type, shape = np_type.subdtype - return np.ones(shape, 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. - - 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, - so ``udt in some_op`` can run it too. It costs no extra compile: Numba - types the stand-ins as it types the wrapper's operands (a C-contiguous - 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: - raise UdfParseError(_summarize_numba_typing_error(exc)) from exc - except Exception: - return None - - def _raise_unless_fits(shape, expected, return_type, path, *, probed): - """Raise ``UdfParseError`` unless a return of ``shape`` fills an ``expected`` slot. - - ``path`` is the record field, or ``None`` for an array UDT's element. - """ - if _fits_by_broadcast(shape, expected): - return - field = "" if path is None else f" for field {path} of {return_type}" - # A probed shape holds for the values the probe used. Saying so helps - # the author of a UDF whose output shape depends on its inputs' values. - how = " when run on sample values" if probed else "" - slot = f"{return_type} elements are" if path is None else "that field holds" - raise UdfParseError( - f"UDT UDF returned an array of shape {tuple(shape)}{field}{how}, but {slot} " - f"{tuple(expected)}. Return an array of that shape, or one that broadcasts to it." - ) - - def _raise_unless_sequence_fits(kind, length, expected, return_type, path, *, probed): - """Raise ``UdfParseError`` unless a ``kind`` of ``length`` fills record field ``path``. - - Numba copies a tuple or list into an array field element by element, - which works only for a one-dimensional field and only at its exact - length. Unlike an array, a sequence does not broadcast, so a 1-tuple - does not fill a 3-element field. A wrong length raises inside the cfunc - just as a wrong-shape array does; a field of higher rank fails the - wrapper compile. - """ - if expected == (length,): - return - how = " when run on sample values" if probed else "" - if len(expected) == 1: - fix = f"Return {expected[0]} values, or an array that broadcasts to {expected}." - else: - fix = f"A {kind} only fills a one-dimensional field; return an array." - raise UdfParseError( - f"UDT UDF returned a {kind} of length {length} for field {path} of {return_type}" - f"{how}, but that field holds {expected}. {fix}" - ) - - # Numba types a record UDF might return for a scalar field, none of which the - # wrapper can write there, each with the kind the error names. ``Array`` - # covers ``NestedArray``. - _NON_SCALAR_KINDS = ( - ("an array", numba.core.types.Array), - ("a tuple", numba.core.types.BaseTuple), - ("a list", numba.core.types.List), - ) - - def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes): - """Reject a UDF whose return cannot fill the output element. - - The wrapper slice-assigns an array return, and each array-typed record - leaf, into the element, so a shape that does not fit raises inside the - cfunc, where the caller never sees it (see - :func:`_compose_wrapper_body`). So does a record leaf the UDF can - return as ``None`` (``v if cond else None``, typed ``Optional``), array - or scalar, on the path that returns it; a leaf that is always ``None`` - does not compile at all, and Numba's traceback says why at length. Nor - does an array, tuple or list returned for a scalar leaf: Numba has no - record setitem for them. Checking here makes each an error when the op - is typed. - - An ``Optional`` leaf is rejected from its type, whichever branch the - sample values would take. Numba also keeps the extents of an operand, - or operand field, returned as-is (a ``NestedArray``), so those are - checked from the type alone. An array the UDF builds (``x + y``, - ``x[:2]``) is a plain ``Array``, which has a rank but no extents, so - only running the UDF, on stand-ins built from ``arg_dtypes``, can say - what shape it has. A tuple or list for an array leaf must match its - 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. - """ - if zkind == "array_elements": - # (record field path, Numba type, slot shape, index into the tuple) - slots = [(None, numba_ret_type, return_type.np_type.shape, None)] - elif zkind == "record_fields": - from .udt_utils import _iter_record_leaves - - # The tuple is flat over the leaves, one element per leaf. - leaves = list(_iter_record_leaves(return_type.np_type)) - if len(numba_ret_type.types) != len(leaves): - return # a wrong field count, not a wrong shape - slots = [ - (path, numba_ret_type.types[i], leaf.shape, i) - for i, (path, _c, leaf) in enumerate(leaves) - ] - else: - return - unknown = [] - for path, nb_type, expected, i in slots: - if isinstance(nb_type, (numba.core.types.Optional, numba.core.types.NoneType)): - field = "" if path is None else f" for field {path}" - raise UdfParseError( - f"UDT UDF can return None{field}, but {return_type} cannot hold None. " - f"Return a value on every path." - ) - if not expected: - for kind, cls in _NON_SCALAR_KINDS: - if isinstance(nb_type, cls): - raise UdfParseError( - f"UDT UDF returned {kind} for field {path} of {return_type}, but " - f"that field holds a scalar. Return a scalar for it." - ) - continue # a scalar leaf: nothing to fit - if isinstance(nb_type, numba.core.types.NestedArray): - _raise_unless_fits(nb_type.shape, expected, return_type, path, probed=False) - elif isinstance(nb_type, numba.core.types.Array) or nb_type is None: - # No extents to read (or, with no return type given, no type - # at all), so learn the shape by running the UDF. - unknown.append((path, expected, i)) - elif isinstance(nb_type, numba.core.types.BaseTuple): - _raise_unless_sequence_fits( - "tuple", len(nb_type.types), expected, return_type, path, probed=False - ) - 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: - for path, expected, i in unknown: - value = ran[0] if i is None else ran[0][i] - if isinstance(value, list): - _raise_unless_sequence_fits( - "list", len(value), expected, return_type, path, probed=True - ) - else: - _raise_unless_fits(np.shape(value), expected, return_type, path, probed=True) - - def _get_udt_wrapper( - numba_func, return_type, dtype, dtype2=None, *, include_indexes=False, numba_ret_type=None - ): - """Build a Numba cfunc wrapper for unary, binary, indexunary, or select UDFs on UDTs. - - ``include_indexes=True`` inserts ``(row, col)`` between ``x`` and - ``y`` in both the wrapper signature and the call to ``numba_func``, - matching the IndexUnaryOp and SelectOp shape. The IndexBinaryOp path - (four indices plus a theta operand) uses - :func:`_get_udt_wrapper_indexbinary`. - """ - nt = numba.types - zsetup, zptr_type, zkind, zinfo = _output_handler(return_type, numba_ret_type) - xsetup, xderef, xptr_type = _input_operand(dtype, "x") - wrapper_args = [zptr_type, xptr_type] - arg_dtypes = [dtype] - if include_indexes: - wrapper_args.extend([UINT64.numba_type, UINT64.numba_type]) - arg_dtypes.extend([UINT64, UINT64]) - ysetup, yderef_expr, yarg = "", "", "" - if dtype2 is not None: - ysetup, yderef, yptr_type = _input_operand(dtype2, "y") - wrapper_args.append(yptr_type) - arg_dtypes.append(dtype2) - 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) - - rcidx = ", row, col" if include_indexes else "" - signature_line = f"def wrapper(z_ptr, x_ptr{rcidx}{yarg}):" - body_setup = f"{zsetup}{xsetup}{ysetup}" - call_expr = f"numba_func({xderef}{rcidx}{yderef_expr})" - - text = _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr) - from .udt_utils import _compile_codegen - - kind = "indexunary" if include_indexes else ("binary" if dtype2 is not None else "unary") - wrapper = _compile_codegen( - text, - func_name="wrapper", - source_label=f"", - extra_ns={"numba_func": numba_func}, - ) - return wrapper, wrapper_sig - - def _get_udt_wrapper_indexbinary( - numba_func, return_type, dtype, dtype2, *, numba_ret_type=None - ): - """Build a Numba cfunc wrapper for IndexBinaryOp UDFs on UDTs. - - Signature: ``f(x, ix, jx, y, iy, jy, theta) -> z``. ``dtype2`` is the - shared type of ``y`` and ``theta``. - """ - nt = numba.types - zsetup, zptr_type, zkind, zinfo = _output_handler(return_type, numba_ret_type) - xsetup, xderef, xptr_type = _input_operand(dtype, "x") - ysetup, yderef, yptr_type = _input_operand(dtype2, "y") - tsetup, tderef, tptr_type = _input_operand(dtype2, "t") - wrapper_sig = nt.void( - zptr_type, - xptr_type, - UINT64.numba_type, - UINT64.numba_type, - yptr_type, - UINT64.numba_type, - UINT64.numba_type, - tptr_type, - ) - arg_dtypes = [dtype, UINT64, UINT64, dtype2, UINT64, UINT64, dtype2] - _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes) - - signature_line = "def wrapper(z_ptr, x_ptr, ix, jx, y_ptr, iy, jy, t_ptr):" - body_setup = f"{zsetup}{xsetup}{ysetup}{tsetup}" - call_expr = f"numba_func({xderef}, ix, jx, {yderef}, iy, jy, {tderef})" - - text = _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr) - from .udt_utils import _compile_codegen - - wrapper = _compile_codegen( - text, - func_name="wrapper", - source_label=f"", - extra_ns={"numba_func": numba_func}, - ) - return wrapper, wrapper_sig - - # Maps ``opclass`` to the SuiteSparse C type name SS allocates for it. # ``IndexBinaryOp`` is in the GxB_ namespace (SS-specific, added in 9.4); # the rest are GrB_. ``SelectOp`` is implemented on top of # ``GrB_IndexUnaryOp`` (BOOL-returning), so its handle frees through the # same ``GrB_IndexUnaryOp_free``. Used by ``TypedOpBase.__del__`` to -# synthesize the pointer cell for ``_free``. +# synthesize the pointer cell for ``_free``, and by +# ``udf._finalize_udt_op`` to call ``_new``. _GB_OBJ_C_TYPENAME = { "UnaryOp": "GrB_UnaryOp", "BinaryOp": "GrB_BinaryOp", diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 25b48b8e3..966a49047 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -26,6 +26,10 @@ from .. import _has_numba, _supports_udfs, ffi, lib from ..dtypes import _sample_values from ..expr import InfixExprBase + +# Imported unconditionally (plain dict, no numba): ``_compile_udt`` consults it +# even when numba is absent. +from ..udt import BUILTIN_UDT_BINARY_OPS as _BUILTIN_UDT_BINARY_OPS from .base import ( _SS_OPERATORS, OpBase, @@ -35,14 +39,10 @@ _hasop, ) -# Imported unconditionally (plain dict, no numba): ``_compile_udt`` consults it -# even when numba is absent. -from .udt_utils import BUILTIN_UDT_BINARY_OPS as _BUILTIN_UDT_BINARY_OPS - if _has_numba: import numba - from .base import ( + from .udf import ( _bool_to_int8, _compile_udf_for_udt, _finalize_udt_op, @@ -64,12 +64,8 @@ def _make_udt_comparison(dtype, dtype2, *, is_eq): UDT. Mismatched-shape UDT pairs raise ``KeyError``. Returns ``(wrapper_func, wrapper_sig)``. """ - from .udt_utils import ( - _array_operand_source, - _check_udt_pair, - _get_udt_info, - _iter_record_leaves, - ) + from ..udt import _check_udt_pair, _get_udt_info, _iter_record_leaves + from .udt_codegen import _array_operand_source info_x = _get_udt_info(dtype) info_y = _get_udt_info(dtype2) @@ -387,7 +383,7 @@ def _pair_dtype(op, dtype, dtype2): if _has_numba: - from .udt_utils import ( + from .udt_codegen import ( _compile_codegen, _has_jit_set, _maybe_warn_jit_skipped, @@ -693,7 +689,7 @@ def _compile_udt(self, dtype, dtype2): else: # eq/ne return BOOL, not the UDT; the kernel is a chain of # per-leaf scalar comparisons. - from .udt_utils import set_jit_c_comparison_on_op + from .udt_codegen import set_jit_c_comparison_on_op op._jit_c_info = set_jit_c_comparison_on_op( op.gb_obj, diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index de2ca5f71..c07968881 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -13,7 +13,7 @@ if _has_numba: import numba - from .base import ( + from .udf import ( _bool_to_int8, _compile_udf_for_udt, _get_udt_wrapper_indexbinary, diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index 5820e66c5..25eecc0b0 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -12,7 +12,7 @@ if _has_numba: import numba - from .base import ( + from .udf import ( _bool_to_int8, _compile_udf_for_udt, _finalize_udt_op, diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 3c99df360..802219b53 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -7,7 +7,7 @@ from .indexunary import IndexUnaryOp, TypedBuiltinIndexUnaryOp if _has_numba: - from .base import _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper + from .udf import _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper _VALUE_OPS = {"valueeq", "valuene", "valuelt", "valuele", "valuegt", "valuege"} diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py new file mode 100644 index 000000000..41e1b7c9d --- /dev/null +++ b/graphblas/core/operator/udf.py @@ -0,0 +1,665 @@ +"""Compiling a user's function (UDF) on UDTs into a GraphBLAS operator. + +GraphBLAS calls an operator as a C function on pointers to raw elements. A UDF +written for a record or array UDT takes Numba records and arrays instead, and +may return a tuple, an array or a number to fill the result's element. The +functions here generate the Numba cfunc between the two: it reads each operand +through its pointer, calls the UDF, and writes what it returns into the output +element. A return that cannot fill that element is an error when the op is +typed, not inside the cfunc. ``_finalize_udt_op`` then registers the cfunc as +a GraphBLAS operator. + +Built-in operators that lift to UDTs are generated by ``udt_codegen``, which uses +these too. ``_bool_to_int8``, the numba#5395 workaround, is here as well: every +Numba UDF on BOOL uses it, not only those on UDTs. + +Importing this module needs Numba. +""" + +import re + +import numba +import numpy as np +from numba import NumbaError + +from ...dtypes import BOOL, INT8, UINT64, lookup_dtype +from ...exceptions import UdfParseError, check_status_carg +from .. import ffi, lib +from ..dtypes import _sample_values +from ..udt import _iter_record_leaves +from .base import _GB_OBJ_C_TYPENAME +from .udt_codegen import _compile_codegen + + +def _bool_to_int8(dtype): + """Return INT8 if ``dtype`` is BOOL, else the dtype unchanged. + + Numba can't compile cfuncs that read or write ``CPointer(boolean)`` + (errors like ``cannot store i1 to i8*`` / ``cond is not i1: i8``); see + numba/numba#5395. Routing BOOL through INT8 sidesteps that; GraphBLAS + coerces back at the cfunc boundary. + + MAINT 2026-05-24: still hits on Numba 0.65. Re-test periodically and + drop the INT8 routing when upstream is fixed. + """ + return INT8 if dtype == BOOL else dtype + + +def _finalize_udt_op(parent_op, dtype, dtype2, ret_type, wrapper, wrapper_sig, typed_user_cls): + """Compile the cfunc, allocate the ``GrB`` op, wrap it, and cache it. + + Shared tail for ``_compile_udt`` in UnaryOp / BinaryOp / IndexUnaryOp / + SelectOp. Looks up the SuiteSparse handle type and ``_new`` symbol + from ``typed_user_cls.opclass``. ``dtype2`` is ``None`` for unary + ops; the rest pass both. Returns the cached ``TypedUser*Op``. + """ + try: + wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) + except NumbaError as exc: + dtypes_str = str(dtype) if dtype2 is None else f"{dtype}, {dtype2}" + raise UdfParseError( + f"{parent_op._modname}.{parent_op.name} does not work with ({dtypes_str}): " + f"{_summarize_numba_typing_error(exc)}" + ) from exc + c_typename = _GB_OBJ_C_TYPENAME[typed_user_cls.opclass] + error_label = c_typename.removeprefix("GrB_").removeprefix("GxB_") + gb_obj = ffi.new(f"{c_typename}*") + new_func = getattr(lib, f"{c_typename}_new") + if dtype2 is None: + check_status_carg( + new_func(gb_obj, wrapper.cffi, ret_type._carg, dtype._carg), + error_label, + gb_obj[0], + ) + op = typed_user_cls(parent_op, parent_op.name, dtype, ret_type, gb_obj[0]) + key = dtype + else: + check_status_carg( + new_func(gb_obj, wrapper.cffi, ret_type._carg, dtype._carg, dtype2._carg), + error_label, + gb_obj[0], + ) + op = typed_user_cls(parent_op, parent_op.name, dtype, ret_type, gb_obj[0], dtype2=dtype2) + key = (dtype, dtype2) + parent_op._udt_types[key] = ret_type + parent_op._udt_ops[key] = op + return op + + +def _compile_udf_for_udt(numba_func, sig, *, op_kind, op_name, dtypes): + """Compile ``sig`` and re-raise Numba compilation errors as ``UdfParseError``. + + Catches the full ``NumbaError`` hierarchy (TypingError, LoweringError, + UnsupportedError, ...) so any compilation failure produces a + single-line UDT diagnostic. + """ + try: + numba_func.compile(sig) + except NumbaError as exc: + dtypes_str = ", ".join(str(d) for d in dtypes) + snippet = _summarize_numba_typing_error(exc) + raise UdfParseError( + f"{op_kind}.{op_name} does not work with ({dtypes_str}): {snippet}" + ) from exc + + +# Numba prefixes its actionable diagnostic line with one of these. +_NUMBA_DIAG_PREFIXES = ( + "No implementation of function", + "No conversion from", + "Field ", + "Cannot infer", + "Operator Overload", + "Invalid use of", + "use of undeclared", + "Untyped global", +) + +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") + + +def _summarize_numba_typing_error(exc): + """Pull the most actionable lines out of a Numba TypingError. + + Numba's "No implementation of function ... found for signature:" + diagnostic puts the signature on the line *below* the prefix; the + signature is the actionable part. When the matched line ends with a + ``:``, append the next non-empty line so the user sees both. + """ + lines = [_ANSI_RE.sub("", line).strip() for line in str(exc).splitlines()] + for i, line in enumerate(lines): + if line.startswith(_NUMBA_DIAG_PREFIXES): + if line.endswith(":"): + for follow in lines[i + 1 :]: + if follow: + return f"{line} {follow}" + return line + for line in lines: + if line: + return line + return "Numba could not compile the function for these input types" + + +def _resolve_udt_return_type(numba_ret_type, *dtypes): + """Resolve a Numba return type to a DataType, matching Tuple returns to an input UDT. + + When a UDF returns a tuple, Numba infers ``Tuple(...)`` rather than a + Record type. Match by field count, preferring a candidate whose field + types align with the Tuple's element types. An array the UDF builds + (``x + y``) is typed without its extents, so it takes the type of the + array-UDT operand with the same element type and rank. + """ + try: + return lookup_dtype(numba_ret_type) + except (ValueError, TypeError): + pass + if isinstance(numba_ret_type, numba.core.types.BaseTuple): + n = len(numba_ret_type.types) + + def _leaves(d): + """Return ``(python_path, c_path, leaf_dtype)`` for each leaf of a record UDT.""" + return list(_iter_record_leaves(d.np_type)) + + # Match by total leaf count (flat for shallow records, total + # leaves across nesting for nested records). + same_arity = [ + d for d in dtypes if d._is_udt and d.np_type.names is not None and len(_leaves(d)) == n + ] + # Prefer a UDT whose leaf types match the tuple elements element-wise. + for d in same_arity: + leaf_dtypes = [leaf for _py, _c, leaf in _leaves(d)] + try: + expected = [lookup_dtype(t).numba_type for t in leaf_dtypes] + except (ValueError, TypeError): + continue + if all(et == tt for et, tt in zip(expected, numba_ret_type.types, strict=True)): + return d + # No perfect match; fall back to the first arity-compatible UDT. + if same_arity: + return same_arity[0] + # Tuple return whose arity matches no input UDT: most likely the + # user is returning the wrong number of fields. List the candidate + # arities so the fix is obvious. + record_arities = sorted( + {len(_leaves(d)) for d in dtypes if d._is_udt and d.np_type.names is not None} + ) + if record_arities: + expected = " or ".join(str(a) for a in record_arities) + raise UdfParseError( + f"UDT UDF returned a tuple of length {n}; expected {expected} " + f"to match one of the input record UDTs." + ) + # All UDT inputs are array UDTs. Tuples don't map to those: the + # function should return a numpy array of the right shape (or a + # scalar) instead. + array_inputs = [d for d in dtypes if d._is_udt and d.np_type.subdtype is not None] + if array_inputs: + shape = array_inputs[0].np_type.subdtype[1] + raise UdfParseError( + f"UDT UDF returned a tuple of length {n}, but inputs are array UDTs of " + f"shape {shape}. Return a numpy array (e.g., ``np.array(...)``) or a " + f"scalar; tuple returns are only matched to record UDTs." + ) + elif isinstance(numba_ret_type, numba.core.types.Array): + # Numba types an array the UDF builds as a plain ``Array``: element + # type and rank, but no extents, which only an operand can supply. + # (An operand returned as-is keeps its ``NestedArray`` type, extents + # and all, and ``lookup_dtype`` resolved it above.) + array_inputs = list( + dict.fromkeys( # dedup: ``x`` and ``y`` are often the same UDT + d for d in dtypes if d._is_udt and d.np_type.subdtype is not None + ) + ) + if isinstance(numba_ret_type, numba.core.types.NestedArray): + # Extents survived, so ``lookup_dtype`` above named a UDT for + # them unless SuiteSparse refused to create it: builds without + # variable-length arrays cap a UDT at 128 bytes before SS 9.0 + # and 1024 from then on. Matching those extents to a same-rank + # operand would hand SuiteSparse an element of the wrong size. + shaped_like = [] + else: + shaped_like = [ + d + for d in array_inputs + if d.numba_type.dtype == numba_ret_type.dtype + and d.numba_type.ndim == numba_ret_type.ndim + ] + if len(shaped_like) == 1: + return shaped_like[0] + if shaped_like: + # Operands that differ only in length are indistinguishable + # here. Guessing would hand SuiteSparse an element of the + # wrong size, so say so. + raise UdfParseError( + f"UDT UDF returned {numba_ret_type!r}, which matches more than one " + f"input array UDT ({', '.join(str(d) for d in shaped_like)}). " + f"Return one of the operands, or make the operands the same type." + ) + # An array UDT went in and an array came out, but not one that + # fits: name the mismatch rather than fall through to the generic + # "unsupported type", whose advice the user already followed. + if array_inputs: + described = "; ".join( + f"{d} elements are {d.numba_type.dtype} with shape {d.numba_type.shape}" + for d in array_inputs + ) + raise UdfParseError( + f"UDT UDF returned {numba_ret_type!r}, which matches no input array " + f"UDT: {described}. Return an array matching one of those, or one of " + f"the operands." + ) + raise UdfParseError( + f"UDT UDF returned an unsupported type {numba_ret_type!r}. " + f"Return a scalar, a tuple matching a record UDT's fields, or a numpy array " + f"matching an array UDT's shape." + ) + + +def _array_udt_view(dtype): + """Return ``(base_element_numba_type, shape)`` for an array UDT. + + Array UDTs are addressed as a ``carray`` over their base elements, in + the UDT's declared shape, rather than as Numba's ``NestedArray``. Numba + models a ``NestedArray`` *value* as an array descriptor (data pointer, + shape, strides, ...), so loading or storing one through a ``CPointer`` + moves the descriptor rather than the element payload, corrupting + whatever follows it (and overrunning the element outright once the + descriptor is the wider of the two). + """ + nested = dtype.numba_type + return nested.dtype, nested.shape + + +def _input_operand(dtype, var): + """Return ``(setup_line, deref_expr, ptr_arg_type)`` for one input operand. + + - ``setup_line`` is the optional ``var = numba.carray(var_ptr, n)`` line + to add to the wrapper body, where ``n`` is 1 for a record UDT and the + element's shape for an array UDT (empty for non-UDT cases). + - ``deref_expr`` is the value to pass to ``numba_func`` for this operand. + - ``ptr_arg_type`` is the Numba ``CPointer(...)`` type for the wrapper signature. + """ + nt = numba.types + if dtype._is_udt: + if dtype.np_type.subdtype is None: + return ( + f" {var} = numba.carray({var}_ptr, 1)\n", + f"{var}[0]", + nt.CPointer(dtype.numba_type), + ) + base_numba, shape = _array_udt_view(dtype) + return ( + f" {var} = numba.carray({var}_ptr, {shape})\n", + var, + nt.CPointer(base_numba), + ) + if dtype == BOOL: + # Numba can't compile bool ptrs (numba/numba#5395); expose them + # as int8 and cast on deref. + # MAINT 2026-05-24: still hits on Numba 0.65; re-test periodically. + return "", f"bool({var}_ptr[0])", nt.CPointer(INT8.numba_type) + return "", f"{var}_ptr[0]", nt.CPointer(dtype.numba_type) + + +def _output_handler(return_type, numba_ret_type): + """Return ``(setup_line, ret_ptr_type, write_kind, write_info)``. + + ``write_kind`` is ``"record_fields"`` when the UDF returns a Tuple to be + unpacked into a record output, ``"array_elements"`` for an array UDT + output, otherwise ``"direct"``. ``write_info`` is one + ``(access_path, is_array)`` pair per leaf for ``"record_fields"`` and a + ``(BL, BR, zname)`` 3-tuple for ``"direct"``; ``"array_elements"`` needs + none, since ``setup_line`` already binds ``z`` to the right shape. + """ + nt = numba.types + ztype = INT8 if return_type == BOOL else return_type + ret_ptr_type = nt.CPointer(ztype.numba_type) + + if ( + numba_ret_type is not None + and isinstance(numba_ret_type, numba.core.types.BaseTuple) + and return_type._is_udt + and return_type.np_type.names is not None + ): + # ``write_info`` pairs each leaf field's Python access path + # (``"['a']"`` for flat, ``"['outer']['inner_a']"`` for nested) + # with whether that leaf is array-typed. The wrapper iterates + # these to write the flat tuple ``_result`` back leaf-by-leaf. + leaves = tuple( + (py, d.subdtype is not None) for py, _c, d in _iter_record_leaves(return_type.np_type) + ) + return ( + " z = numba.carray(z_ptr, 1)\n", + ret_ptr_type, + "record_fields", + leaves, + ) + if return_type._is_udt: + if return_type.np_type.subdtype is None: + return ( + " z = numba.carray(z_ptr, 1)\n", + ret_ptr_type, + "direct", + ("", "", "z[0]"), + ) + base_numba, shape = _array_udt_view(return_type) + return ( + f" z = numba.carray(z_ptr, {shape})\n", + nt.CPointer(base_numba), + "array_elements", + None, + ) + if return_type == BOOL: + return "", ret_ptr_type, "direct", ("bool(", ")", "z_ptr[0]") + return "", ret_ptr_type, "direct", ("", "", "z_ptr[0]") + + +def _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr): + """Assemble the Python source for a UDT cfunc wrapper. + + For record returns, the wrapper writes leaf-by-leaf; Numba does not + 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. + """ + if zkind == "array_elements": + return ( + f"{signature_line}\n" + f"{body_setup}" + f" _result = {call_expr}\n" + f" z[:] = _result\n" + ) + if zkind == "record_fields": + # Array-typed leaves slice-assign. Numba's record-field setitem + # copies the destination's extent regardless of the source's, so a + # short source is read past its end and a long one is silently + # truncated; ``[:]`` makes both a shape error instead. + field_assigns = "".join( + f" z[0]{path}{'[:]' if is_array else ''} = _result[{i}]\n" + for i, (path, is_array) in enumerate(zinfo) + ) + return f"{signature_line}\n{body_setup} _result = {call_expr}\n{field_assigns}" + BL, BR, zname = zinfo + return f"{signature_line}\n{body_setup} {zname} = {BL}{call_expr}{BR}\n" + + +def _fits_by_broadcast(shape, expected): + """Whether a return of ``shape`` fills an ``expected``-shaped destination. + + The wrapper slice-assigns (``z[:] = ...``), and numpy broadcasts on + assignment, so an exact match is not the requirement: a ``(1,)`` return + legitimately fills a ``(6,)`` element, and a ``(1, 3)`` return fills + every row of a ``(2, 3)`` one. Rejecting those would refuse code that + works today. + + numpy's assignment rule is broadcasting plus a leading-``1`` strip when + the source has the higher rank, so ``(1, 6)`` fits ``(6,)`` while + ``(6, 1)`` does not. ``test_udt_broadcast_matches_numba_slice_assign`` + pins this against Numba's own slice-assign for both. + """ + shape = tuple(shape) + expected = tuple(expected) + while len(shape) > len(expected) and shape[0] == 1: + shape = shape[1:] + try: + return np.broadcast_shapes(shape, expected) == expected + except ValueError: + return False + + +def _udf_probe_value(dtype): + """Build a stand-in operand of ``dtype`` for :func:`_run_udf_probe`. + + Only the shape of the UDF's result is read, so any values it accepts + will do. + """ + if dtype._is_udt: + np_type = dtype.np_type + if np_type.subdtype is None: + return np.ones(1, dtype=np_type)[0] + base_np_type, shape = np_type.subdtype + return np.ones(shape, 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. + + 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, + so ``udt in some_op`` can run it too. It costs no extra compile: Numba + types the stand-ins as it types the wrapper's operands (a C-contiguous + 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: + raise UdfParseError(_summarize_numba_typing_error(exc)) from exc + except Exception: + return None + + +def _raise_unless_fits(shape, expected, return_type, path, *, probed): + """Raise ``UdfParseError`` unless a return of ``shape`` fills an ``expected`` slot. + + ``path`` is the record field, or ``None`` for an array UDT's element. + """ + if _fits_by_broadcast(shape, expected): + return + field = "" if path is None else f" for field {path} of {return_type}" + # A probed shape holds for the values the probe used. Saying so helps + # the author of a UDF whose output shape depends on its inputs' values. + how = " when run on sample values" if probed else "" + slot = f"{return_type} elements are" if path is None else "that field holds" + raise UdfParseError( + f"UDT UDF returned an array of shape {tuple(shape)}{field}{how}, but {slot} " + f"{tuple(expected)}. Return an array of that shape, or one that broadcasts to it." + ) + + +def _raise_unless_sequence_fits(kind, length, expected, return_type, path, *, probed): + """Raise ``UdfParseError`` unless a ``kind`` of ``length`` fills record field ``path``. + + Numba copies a tuple or list into an array field element by element, + which works only for a one-dimensional field and only at its exact + length. Unlike an array, a sequence does not broadcast, so a 1-tuple + does not fill a 3-element field. A wrong length raises inside the cfunc + just as a wrong-shape array does; a field of higher rank fails the + wrapper compile. + """ + if expected == (length,): + return + how = " when run on sample values" if probed else "" + if len(expected) == 1: + fix = f"Return {expected[0]} values, or an array that broadcasts to {expected}." + else: + fix = f"A {kind} only fills a one-dimensional field; return an array." + raise UdfParseError( + f"UDT UDF returned a {kind} of length {length} for field {path} of {return_type}" + f"{how}, but that field holds {expected}. {fix}" + ) + + +# Numba types a record UDF might return for a scalar field, none of which the +# wrapper can write there, each with the kind the error names. ``Array`` +# covers ``NestedArray``. +_NON_SCALAR_KINDS = ( + ("an array", numba.core.types.Array), + ("a tuple", numba.core.types.BaseTuple), + ("a list", numba.core.types.List), +) + + +def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes): + """Reject a UDF whose return cannot fill the output element. + + The wrapper slice-assigns an array return, and each array-typed record + leaf, into the element, so a shape that does not fit raises inside the + cfunc, where the caller never sees it (see + :func:`_compose_wrapper_body`). So does a record leaf the UDF can + return as ``None`` (``v if cond else None``, typed ``Optional``), array + or scalar, on the path that returns it; a leaf that is always ``None`` + does not compile at all, and Numba's traceback says why at length. Nor + does an array, tuple or list returned for a scalar leaf: Numba has no + record setitem for them. Checking here makes each an error when the op + is typed. + + An ``Optional`` leaf is rejected from its type, whichever branch the + sample values would take. Numba also keeps the extents of an operand, + or operand field, returned as-is (a ``NestedArray``), so those are + checked from the type alone. An array the UDF builds (``x + y``, + ``x[:2]``) is a plain ``Array``, which has a rank but no extents, so + only running the UDF, on stand-ins built from ``arg_dtypes``, can say + what shape it has. A tuple or list for an array leaf must match its + 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. + """ + if zkind == "array_elements": + # (record field path, Numba type, slot shape, index into the tuple) + slots = [(None, numba_ret_type, return_type.np_type.shape, None)] + elif zkind == "record_fields": + # The tuple is flat over the leaves, one element per leaf. + leaves = list(_iter_record_leaves(return_type.np_type)) + if len(numba_ret_type.types) != len(leaves): + return # a wrong field count, not a wrong shape + slots = [ + (path, numba_ret_type.types[i], leaf.shape, i) + for i, (path, _c, leaf) in enumerate(leaves) + ] + else: + return + unknown = [] + for path, nb_type, expected, i in slots: + if isinstance(nb_type, (numba.core.types.Optional, numba.core.types.NoneType)): + field = "" if path is None else f" for field {path}" + raise UdfParseError( + f"UDT UDF can return None{field}, but {return_type} cannot hold None. " + f"Return a value on every path." + ) + if not expected: + for kind, cls in _NON_SCALAR_KINDS: + if isinstance(nb_type, cls): + raise UdfParseError( + f"UDT UDF returned {kind} for field {path} of {return_type}, but " + f"that field holds a scalar. Return a scalar for it." + ) + continue # a scalar leaf: nothing to fit + if isinstance(nb_type, numba.core.types.NestedArray): + _raise_unless_fits(nb_type.shape, expected, return_type, path, probed=False) + elif isinstance(nb_type, numba.core.types.Array) or nb_type is None: + # No extents to read (or, with no return type given, no type + # at all), so learn the shape by running the UDF. + unknown.append((path, expected, i)) + elif isinstance(nb_type, numba.core.types.BaseTuple): + _raise_unless_sequence_fits( + "tuple", len(nb_type.types), expected, return_type, path, probed=False + ) + 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: + for path, expected, i in unknown: + value = ran[0] if i is None else ran[0][i] + if isinstance(value, list): + _raise_unless_sequence_fits( + "list", len(value), expected, return_type, path, probed=True + ) + else: + _raise_unless_fits(np.shape(value), expected, return_type, path, probed=True) + + +def _get_udt_wrapper( + numba_func, return_type, dtype, dtype2=None, *, include_indexes=False, numba_ret_type=None +): + """Build a Numba cfunc wrapper for unary, binary, indexunary, or select UDFs on UDTs. + + ``include_indexes=True`` inserts ``(row, col)`` between ``x`` and + ``y`` in both the wrapper signature and the call to ``numba_func``, + matching the IndexUnaryOp and SelectOp shape. The IndexBinaryOp path + (four indices plus a theta operand) uses + :func:`_get_udt_wrapper_indexbinary`. + """ + nt = numba.types + zsetup, zptr_type, zkind, zinfo = _output_handler(return_type, numba_ret_type) + xsetup, xderef, xptr_type = _input_operand(dtype, "x") + wrapper_args = [zptr_type, xptr_type] + arg_dtypes = [dtype] + if include_indexes: + wrapper_args.extend([UINT64.numba_type, UINT64.numba_type]) + arg_dtypes.extend([UINT64, UINT64]) + ysetup, yderef_expr, yarg = "", "", "" + if dtype2 is not None: + ysetup, yderef, yptr_type = _input_operand(dtype2, "y") + wrapper_args.append(yptr_type) + arg_dtypes.append(dtype2) + 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) + + rcidx = ", row, col" if include_indexes else "" + signature_line = f"def wrapper(z_ptr, x_ptr{rcidx}{yarg}):" + body_setup = f"{zsetup}{xsetup}{ysetup}" + call_expr = f"numba_func({xderef}{rcidx}{yderef_expr})" + + text = _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr) + kind = "indexunary" if include_indexes else ("binary" if dtype2 is not None else "unary") + wrapper = _compile_codegen( + text, + func_name="wrapper", + source_label=f"", + extra_ns={"numba_func": numba_func}, + ) + return wrapper, wrapper_sig + + +def _get_udt_wrapper_indexbinary(numba_func, return_type, dtype, dtype2, *, numba_ret_type=None): + """Build a Numba cfunc wrapper for IndexBinaryOp UDFs on UDTs. + + Signature: ``f(x, ix, jx, y, iy, jy, theta) -> z``. ``dtype2`` is the + shared type of ``y`` and ``theta``. + """ + nt = numba.types + zsetup, zptr_type, zkind, zinfo = _output_handler(return_type, numba_ret_type) + xsetup, xderef, xptr_type = _input_operand(dtype, "x") + ysetup, yderef, yptr_type = _input_operand(dtype2, "y") + tsetup, tderef, tptr_type = _input_operand(dtype2, "t") + wrapper_sig = nt.void( + zptr_type, + xptr_type, + UINT64.numba_type, + UINT64.numba_type, + yptr_type, + UINT64.numba_type, + UINT64.numba_type, + tptr_type, + ) + arg_dtypes = [dtype, UINT64, UINT64, dtype2, UINT64, UINT64, dtype2] + _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes) + + signature_line = "def wrapper(z_ptr, x_ptr, ix, jx, y_ptr, iy, jy, t_ptr):" + body_setup = f"{zsetup}{xsetup}{ysetup}{tsetup}" + call_expr = f"numba_func({xderef}, ix, jx, {yderef}, iy, jy, {tderef})" + + text = _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr) + wrapper = _compile_codegen( + text, + func_name="wrapper", + source_label=f"", + extra_ns={"numba_func": numba_func}, + ) + return wrapper, wrapper_sig diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_codegen.py similarity index 72% rename from graphblas/core/operator/udt_utils.py rename to graphblas/core/operator/udt_codegen.py index 30385b354..09bed96fa 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_codegen.py @@ -1,10 +1,19 @@ -"""Shared utilities for auto-generating UDT operator implementations. - -Provides Numba wrapper generators and JIT C code generators for element-wise -operations on record UDTs (struct types) and array UDTs (e.g., FP64[3]). - -Compilation is lazy: functions here are only called from ``_compile_udt`` -methods when an operator is first used with a particular UDT. +"""Code generation for the operators that compute on UDTs. + +``core/udt.py`` decides which operand pairs an op takes and what type its +result has; this module writes the code that computes it, for record UDTs +(C structs) and array UDTs (such as ``FP64[3]``): + +- ``_udt_c_typedef``, a UDT's C struct, set when the UDT is registered; +- under ``if _has_numba:``, the Numba cfuncs: ``compile_udt_binary_wrapper``, + ``compile_udt_unary_wrapper``, ``compile_udt_cast_wrapper`` and + ``fused_store_op``; +- the JIT C definitions SuiteSparse compiles in place of a cfunc + (``set_jit_c_on_op``, ``set_jit_c_comparison_on_op``). + +``_compile_codegen`` compiles the generated Python, for this module, +``operator.udf`` and the comparisons in ``operator.binary``. An op is compiled +when it is first typed for a UDT. """ import ast @@ -17,6 +26,14 @@ from ... import backend from .. import _has_numba, ffi, lib +from ..udt import ( + BUILTIN_UDT_BINARY_OPS, + BUILTIN_UDT_UNARY_OPS, + _check_udt_pair, + _get_udt_info, + _iter_record_leaves, + _udt_op_types, +) if _has_numba: import numba @@ -84,21 +101,6 @@ def _compile_codegen(src, *, func_name, source_label, extra_ns=None): return namespace[func_name] -BUILTIN_UDT_BINARY_OPS = { - "plus": "+", - "minus": "-", - "times": "*", - "truediv": "/", - "floordiv": "//", - "min": "min", - "max": "max", -} - -BUILTIN_UDT_UNARY_OPS = { - "ainv": "-", - "abs": "abs", -} - # Ops that use function-call syntax rather than infix (e.g., min(a, b), not "a min b") _FUNC_BINARY_OPS = {"min", "max"} _FUNC_UNARY_OPS = {"abs"} @@ -121,14 +123,11 @@ def _compile_codegen(src, *, func_name, source_label, extra_ns=None): # JIT include chain pulls in GraphBLAS.h, so these names are in scope. # Plus, minus, times, truediv, ainv compile natively on _Complex; abs # uses cabs / cabsf (see _c_expr_unary). min, max, floordiv don't make - # sense on complex; ``_element_types`` rejects them before any codegen. + # sense on complex; ``udt._element_types`` rejects them before any codegen. np.dtype(np.complex64): "GxB_FC32_t", np.dtype(np.complex128): "GxB_FC64_t", } -# Ops with no meaning on complex fields (no ordering for min/max, no integer -# mod for floordiv); ``_element_types`` rejects them on both paths. -_OPS_NOT_FOR_COMPLEX = frozenset({"min", "max", "floordiv"}) # ``_MIN`` spellings from , keyed by itemsize. Signed integer # division traps on ``MIN / -1``; the generated C tests for it explicitly. @@ -320,534 +319,6 @@ def _pick_c_type_name(python_name): return f"_gbudt_{next(_synthetic_udt_counter)}" -def _get_udt_info(dtype): - """Return ('record', field_names) or ('array', (base_dtype, flat_size)) or None.""" - np_type = dtype.np_type - if np_type.names is not None: - return "record", np_type.names - if np_type.subdtype is not None: - base, shape = np_type.subdtype - return "array", (base, reduce(mul, shape)) - return None - - -def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): - """Raise ``KeyError`` if two UDT operands disagree on shape. - - ``info_x`` and ``info_y`` are ``_get_udt_info`` results. When either is - ``None`` (one side is a scalar broadcast), no check fires. Without this - gate the generated wrapper code raises a cryptic Numba ``TypingError`` - on the first field access, or, for array fields that do not broadcast, - fails inside the cfunc without telling the caller. Element dtypes may - differ; :func:`_udt_op_types` promotes them. - """ - if info_x is None or info_y is None: - return - kind_x, detail_x = info_x - kind_y, detail_y = info_y - if kind_x != kind_y: - raise KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"cannot mix record and array UDTs in a single element-wise op." - ) - if kind_x == "record": - if detail_x != detail_y: - raise KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"record UDTs must share field names, in the same order; got " - f"{list(detail_x)} vs {list(detail_y)}." - ) - # Matching top-level names is not enough. The codegen reads both - # operands at the left one's leaf paths, so a field that is a - # sub-record on one side and a scalar on the other, or sub-records - # with different names, fails Numba's typing: a compile error - # reported for what is really the same shape disagreement the check - # above reports as a KeyError. The paths must also come in the same - # order, as the top-level names must, so that leaf ``i`` of one - # operand is leaf ``i`` of the other. - leaves_x = list(_iter_record_leaves(dtype.np_type)) - leaves_y = list(_iter_record_leaves(dtype2.np_type)) - if [py for py, _c, _d in leaves_x] != [py for py, _c, _d in leaves_y]: - raise KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"record UDTs must nest the same way, with the same field names in the " - f"same order at every level; got {[c for _py, c, _d in leaves_x]} vs " - f"{[c for _py, c, _d in leaves_y]}." - ) - # Each field must also have the same shape in both (field dtypes may - # differ), as numpy requires before it promotes or compares two record - # dtypes. Shapes that do not broadcast would raise inside the cfunc, - # unseen, and leave whatever was in the output buffer. Shapes that do - # broadcast would work one way round only, since the arithmetic ops - # write into a field shaped like the left operand's: ``(3,)`` plus - # ``(1,)`` works, ``(1,)`` plus ``(3,)`` does not. - for (_py, c, leaf_x), (_py_y, _c_y, leaf_y) in zip(leaves_x, leaves_y, strict=True): - if leaf_x.shape != leaf_y.shape: - raise KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"record UDTs must have the same shape for each field; field {c} is " - f"{leaf_x.shape} vs {leaf_y.shape}." - ) - return - # Array UDTs broadcast as numpy arrays do, so ``(3, 1)`` with ``(1, 4)`` - # gives ``(3, 4)``; the generated code reads each operand through a flat - # index map (``_array_operand_source``). Shapes such as ``(2, 3)`` with - # ``(3, 2)`` do not broadcast, as numpy refuses them. - (base_x, shape_x), (base_y, shape_y) = dtype.np_type.subdtype, dtype2.np_type.subdtype - try: - np.broadcast_shapes(shape_x, shape_y) - except ValueError: - raise KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): array UDTs must " - f"have shapes that broadcast together; got {base_x}{list(shape_x)} vs " - f"{base_y}{list(shape_y)}." - ) from None - - -def _corresponding_leaves(np_x, np_z): - """Return ``(reason, pairs)`` for putting an element of ``np_x`` into one of ``np_z``. - - ``reason`` says why the two layouts do not correspond, or is ``None``, and - then ``pairs`` holds the base dtypes of each pair of leaves. Records must - line up as the operands of a lifted op do (:func:`_check_udt_pair`): the - same field names, order and nesting and the same shape for each field. - Arrays must have the same shape apart from leading axes of length 1; - operands broadcast, but a stored element keeps every value in its place. - numpy casts records by position whatever their names, and repeats or - drops elements of an array field of another length, which would put - values in places nobody chose. - """ - if (np_x.names is None) != (np_z.names is None): - return "a record UDT and an array UDT do not correspond", None - if np_x.names is None: - if np_x.subdtype is None or np_z.subdtype is None: - return "only record and array UDTs cast", None - (base_x, shape_x), (base_z, shape_z) = np_x.subdtype, np_z.subdtype - if _strip_leading_ones(shape_x) != _strip_leading_ones(shape_z): - return "array UDTs must have the same shape, apart from leading axes of length 1", None - return None, [(base_x, base_z)] - leaves_x = list(_iter_record_leaves(np_x)) - leaves_z = list(_iter_record_leaves(np_z)) - if [py for py, _c, _d in leaves_x] != [py for py, _c, _d in leaves_z]: - return "records must have the same field names, in the same order, at every level", None - pairs = [(x, z) for (_, _, x), (_, _, z) in zip(leaves_x, leaves_z, strict=True)] - if any(x.shape != z.shape for x, z in pairs): - return "records must have the same shape for each field", None - return None, [(x.base, z.base) for x, z in pairs] - - -def _udt_cast_error(dtype, to_dtype): - """Return why an element of UDT ``dtype`` cannot be stored as UDT ``to_dtype``, or ``None``. - - The layouts must correspond (:func:`_corresponding_leaves`), and each pair - of elements must be built-in numbers, or the same type. - """ - reason, pairs = _corresponding_leaves(dtype.np_type, to_dtype.np_type) - for x, z in pairs or (): - if x != z and not (x.kind in "biufc" and z.kind in "biufc"): - return f"elements of {x} do not cast to {z}" - return reason - - -def _strip_leading_ones(shape): - i = 0 - while i < len(shape) and shape[i] == 1: - i += 1 - return shape[i:] - - -def _builtin_dtype(np_type): - """Return the built-in ``DataType`` for ``np_type``, or ``None`` if it is not one.""" - from ..dtypes import _registry - - dt = _registry.get(np_type) - return None if dt is None or dt._is_udt else dt - - -def _complex_op_error(op_name, dtype, dtype2): - return KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"this op is not defined on complex fields. Use ``binary.plus``, " - f"``minus``, ``times``, or ``truediv`` for complex element-wise " - f"arithmetic, or register a custom binary op." - ) - - -def _element_types(op_name, dtype, dtype2, elem_x, elem_y): - """Return ``(in_x, in_y, out)``, the numpy dtypes ``binary.{op_name}`` uses on one element. - - They come from the op's typing on the built-in dtypes, so each element of an - array UDT, or each record leaf, is promoted as a Vector of that dtype would - be, and ``truediv`` on integers gives floats. ``in_x`` and ``in_y`` are the - types GraphBLAS casts the operands to before it applies the op. The - generated code casts to them too, because Numba promotes differently: it - adds uint64 and int64 in int64, where GraphBLAS uses float64. ``dtype`` and - ``dtype2`` are the operands, named in errors. Two elements that are not - both built-in dtypes (bytes, datetime, records inside an array field) have - no promotion rule, so they must be the same dtype, which the result keeps. - """ - # ``min``/``max``/``floordiv`` have no defined semantics on complex - # operands (no ordering, no integer mod). Reject early with a clear - # message; otherwise Numba's cfunc compile blows up several frames - # down with ``NotImplementedError: No definition for lowering lt``. - if op_name in _OPS_NOT_FOR_COMPLEX and "c" in (elem_x.kind, elem_y.kind): - raise _complex_op_error(op_name, dtype, dtype2) - builtin_x = _builtin_dtype(elem_x) - builtin_y = _builtin_dtype(elem_y) - if builtin_x is None or builtin_y is None: - if elem_x == elem_y: - return elem_x, elem_y, elem_x - # Builds without complex dtypes (the vanilla backend) still lift complex - # fields, so promote them as the built-in typing does where it has them. - if "c" in (elem_x.kind, elem_y.kind) and {elem_x.kind, elem_y.kind} <= set("biufc"): - promoted = np.promote_types(elem_x, elem_y) - return promoted, promoted, promoted - raise KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): elements of " - f"{elem_x} and {elem_y} have no common type." - ) - from ... import binary - from .utils import get_typed_op - - typed = get_typed_op(getattr(binary, op_name), builtin_x, builtin_y) - return typed.type.np_type, typed.type2.np_type, typed.return_type.np_type - - -def _udt_op_types(op_name, dtype, dtype2): - """Return ``(ret_type, input_types)`` for ``binary.{op_name}`` on ``dtype`` and ``dtype2``. - - At least one operand is a record or array UDT, and ``_check_udt_pair`` has - accepted the pair; the other may be a plain scalar type. Each element, or - record leaf, follows :func:`_element_types`, so the result never narrows: - an int64 array UDT plus ``0.5`` is ``FP64[3]``. ``ret_type`` is an - operand's own type when that holds the result, so same-type arithmetic - keeps its UDT. Otherwise it is the promoted layout, which is the UDT a user - registered with that layout if there is one, else an anonymous one. - ``input_types`` holds the ``(in_x, in_y)`` pair of each record leaf, in - leaf order, or the single pair for an array UDT's elements. - """ - from ..dtypes import lookup_dtype - - np_x, np_y = dtype.np_type, dtype2.np_type - if np_x.names is not None or np_y.names is not None: - leaves_x = list(_iter_record_leaves(np_x)) if np_x.names is not None else None - leaves_y = list(_iter_record_leaves(np_y)) if np_y.names is not None else None - input_types = [] - result_leaves = [] - for i, (_py, _c, leaf) in enumerate(leaves_x or leaves_y): - elem_x = np_x if leaves_x is None else leaves_x[i][2].base - elem_y = np_y if leaves_y is None else leaves_y[i][2].base - in_x, in_y, elem = _element_types(op_name, dtype, dtype2, elem_x, elem_y) - input_types.append((in_x, in_y)) - result_leaves.append(np.dtype((elem, leaf.shape)) if leaf.shape else elem) - for operand, leaves in [(dtype, leaves_x), (dtype2, leaves_y)]: - if leaves is not None and all( - result_leaf == leaf - for result_leaf, (_py, _c, leaf) in zip(result_leaves, leaves, strict=True) - ): - return operand, input_types - records = [np_type for np_type in (np_x, np_y) if np_type.names is not None] - return lookup_dtype(_promoted_record(records, iter(result_leaves))), input_types - base_x, shape_x = np_x.subdtype or (np_x, ()) - base_y, shape_y = np_y.subdtype or (np_y, ()) - in_x, in_y, elem = _element_types(op_name, dtype, dtype2, base_x, base_y) - result = np.dtype((elem, np.broadcast_shapes(shape_x, shape_y))) - for operand in (dtype, dtype2): - if operand.np_type == result: - return operand, [(in_x, in_y)] - return lookup_dtype(result), [(in_x, in_y)] - - -def _promoted_record(records, result_leaves): - """Build a record laid out like ``records`` with ``result_leaves`` for its leaves. - - ``records`` holds one or two record dtypes with the same field names and - nesting. The new record aligns its fields, at each level, if any of them - does there, so a C-compatible operand gives a C-compatible result. - """ - first = records[0] - fields = [] - for name in first.names: - field = first.fields[name][0] - if field.names is not None: - field = _promoted_record([r.fields[name][0] for r in records], result_leaves) - else: - field = next(result_leaves) - fields.append((name, field)) - return np.dtype(fields, align=any(r.isalignedstruct for r in records)) - - -# The kinds of element a weak Python literal fits without changing kind, and -# the dtype it takes otherwise: numpy 2's rules for Python scalars (NEP 50). -_WEAK_LITERAL_FITS = {"b": "biufc", "i": "iufc", "f": "fc", "c": "c"} -_WEAK_LITERAL_DEFAULT = { - "b": np.dtype(np.bool_), - "i": np.dtype(np.int64), - "f": np.dtype(np.float64), - "c": np.dtype(np.complex128), -} - - -@lru_cache -def _int_bounds(elem): - info = np.iinfo(elem) - return info.min, info.max - - -def _weak_literal_element(value, kind, elem): - """Return the dtype a Python literal of ``kind`` takes beside an element of ``elem``. - - It is ``elem`` when that kind of element holds the literal (``1`` beside - int8, ``0.5`` beside float32), else the literal's default dtype (``0.5`` - beside int8 is float64), and complex of ``elem``'s precision for a complex - literal beside a float. An int literal outside an integer element's range - raises ``OverflowError``, as numpy 2 does, rather than wrapping. - """ - if elem.kind in _WEAK_LITERAL_FITS[kind]: - if kind == "i" and elem.kind in "iu": - lo, hi = _int_bounds(elem) - if not lo <= value <= hi: - raise OverflowError(f"Python integer {value} out of bounds for {elem}") - return elem - if kind == "c" and elem.kind == "f": - return np.result_type(elem, np.complex64) - rv = _WEAK_LITERAL_DEFAULT[kind] - if kind == "i": # beside a bool, an int is INT64, and must fit it - lo, hi = _int_bounds(rv) - if not lo <= value <= hi: - raise OverflowError(f"Python integer {value} out of bounds for {rv}") - return rv - - -def _weak_literal_leaf(values, kind, elem, *, exact=False): - """Return the one dtype Python numbers ``values`` of widest kind ``kind`` take beside ``elem``. - - Each is typed by :func:`_weak_literal_element`, so an int out of an integer - ``elem``'s range raises ``OverflowError``. With ``exact``, for a comparison, - which keeps no result type, a fill value or truediv, such an int takes - int64, uint64 or float64, whichever holds the values (a comparison then - compares it as infinity: ``scalar._literal_operand``). Built-in dtypes and - UDT leaves both use this. - """ - if exact and kind == "i" and elem.kind in "biu": - # The integer type an int would take: ``elem``, or INT64 beside a bool. - lo, hi = _int_bounds(elem if elem.kind != "b" else _WEAK_LITERAL_DEFAULT["i"]) - for x in values: - if not lo <= x <= hi: - if all(-(2**63) <= x < 2**63 for x in values): - return np.dtype(np.int64) - if all(0 <= x < 2**64 for x in values): - return np.dtype(np.uint64) - return np.dtype(np.float64) - for x in values: - elem_type = _weak_literal_element(x, kind, elem) - return elem_type - - -_LITERAL_KINDS = {bool: "b", int: "i", float: "f", complex: "c"} - - -class _LiteralDoesNotFit(Exception): - """A tuple or list literal is not laid out like an element of the UDT.""" - - -def _as_python_number(value): - """Return a ``Fraction`` or ``Decimal`` as the int or float it stands for; anything else as is. - - numpy would convert either into an integer field by truncating it (``Fraction(1, 2)`` - becomes 0), so they are typed as the Python number they are worth instead. - """ - from decimal import Decimal - from fractions import Fraction - - if isinstance(value, (Fraction, Decimal)): - try: - return int(value) if value == int(value) else float(value) - except (OverflowError, ValueError): # infinite or NaN - return float(value) - return value - - -def _literal_leaf_values(value, shape): - """Return the Python numbers in literal ``value`` read as an array of ``shape``, or ``None``. - - ``shape`` is ``()`` for a scalar leaf. A number fills any shape, as numpy - broadcasts it. A sequence that does not fill the shape, or holds anything - that is not a Python number, gives ``None``. - """ - if shape: - arr = np.empty(shape, dtype=object) - try: - arr[...] = value - except (TypeError, ValueError): - return None - values = [_as_python_number(x) for x in arr.ravel().tolist()] - else: - values = [_as_python_number(value)] - return values if all(type(x) in _LITERAL_KINDS for x in values) else None - - -def _weak_literal_udt(dtype, value, *, exact=False): - """Return the UDT a Python number, or a tuple or list of them, becomes beside UDT ``dtype``. - - The literal is read as an element of ``dtype``: a number stands in every - leaf, a sequence fills an array UDT or an array field, and a tuple or list - fills a record by position, nested for a nested record. Each leaf then takes - the dtype numpy 2 gives a Python scalar beside that leaf's element - (:func:`_weak_literal_element`), of the widest kind among the leaf's - numbers, so the pair promotes as two UDTs do: ``int8_udt + 1`` and - ``int8_udt + (1, 2, 3)`` stay ``int8_udt``, while ``int8_udt * 2.5`` and - ``int8_udt + (0.5, 1.5, 2.5)`` give float64 elements. A leaf that is not a - number (bytes, datetime) takes the literal's default dtype and fails to - pair. With ``exact``, for ``eq`` and ``ne``, an int out of a leaf's range - takes a type that holds it instead of raising (:func:`_weak_literal_leaf`), - so ``int8_udt == 300`` is False. A sequence not laid out like an element - gives ``None``. - """ - from ..dtypes import lookup_dtype - - def leaf_dtype(leaf, part): - base, shape = leaf.subdtype or (leaf, ()) - values = _literal_leaf_values(part, shape) - if values is None: - raise _LiteralDoesNotFit - kind = max((_LITERAL_KINDS[type(x)] for x in values), key="bifc".index) - if base.kind in "biufc": - elem = _weak_literal_leaf(values, kind, base, exact=exact) - else: - elem = _WEAK_LITERAL_DEFAULT[kind] - return np.dtype((elem, shape)) if shape else elem - - def record_leaves(np_type, part): - if type(_as_python_number(part)) in _LITERAL_KINDS: - parts = [part] * len(np_type.names) - elif isinstance(part, (tuple, list)) and len(part) == len(np_type.names): - parts = part - else: - raise _LiteralDoesNotFit - for name, item in zip(np_type.names, parts, strict=True): - field = np_type.fields[name][0] - if field.names is not None: - yield from record_leaves(field, item) - else: - yield leaf_dtype(field, item) - - np_type = dtype.np_type - try: - if np_type.names is not None: - literal_type = _promoted_record([np_type], record_leaves(np_type, value)) - else: - literal_type = leaf_dtype(np_type, value) - except _LiteralDoesNotFit: - return None - return dtype if literal_type == np_type else lookup_dtype(literal_type) - - -def _int_beyond_leaf(np_type, value): - """Return a Python int in literal ``value`` beyond its integer leaf's range, or ``None``. - - ``value`` is laid out like an element of ``np_type``, as - :func:`_weak_literal_udt` reads it (a record's dict already a tuple). No - element equals a literal with such an int, whatever its other leaves. - """ - if np_type.names is not None: - parts = value if isinstance(value, (tuple, list)) else [value] * len(np_type.names) - for name, part in zip(np_type.names, parts, strict=True): - found = _int_beyond_leaf(np_type.fields[name][0], part) - if found is not None: - return found - return None - base, shape = np_type.subdtype or (np_type, ()) - if base.kind in "iu": - lo, hi = _int_bounds(base) - for x in _literal_leaf_values(value, shape) or (): - if type(x) is int and not lo <= x <= hi: - return x - return None - - -def _array_literal_udt(dtype, value): - """Return the UDT a numpy array ``value`` is beside UDT ``dtype``, or ``None`` to convert it. - - An array has a dtype, so it is strong, as numpy 2 treats it: - ``int8_udt + np.array([0.5, 1, 2])`` has float64 elements and - ``+ np.array([1, 2, 3])`` int64 ones. Its type is the array layout of its - dtype, in its own shape when that pairs with the UDT's and in the UDT's - shape when it broadcasts to that (``np.array([1])`` beside ``INT8[3]``). - A 0-d structured array or ``np.void`` beside a record is its record dtype. - Anything else, such as a numeric array beside a record, is converted into - the UDT. - """ - from ..dtypes import lookup_dtype - - np_type = dtype.np_type - if value.dtype.names is not None: - return lookup_dtype(value.dtype) if value.ndim == 0 else None - # A bool array converts as it is (a BOOL array UDT cannot be registered), - # and bool casts safely into every numeric element. - if np_type.names is not None or value.dtype.kind not in "iufc": - return None - shape = np_type.subdtype[1] - if _strip_leading_ones(value.shape) == _strip_leading_ones(shape): - shape = value.shape - else: - try: - np.broadcast_to(value, shape) - except ValueError: - return None - return lookup_dtype(np.dtype((value.dtype, shape))) - - -def _leaf_dtypes(np_type): - """Yield the base dtype of each leaf of ``np_type``: a record's fields, an array's element.""" - if np_type.names is not None: - for _py, _c, leaf in _iter_record_leaves(np_type): - yield leaf.base - else: - yield np_type.base - - -def _casts_safely_into(literal_type, dtype): - """Whether each leaf of ``literal_type`` casts safely (``np.can_cast``) into UDT ``dtype``'s. - - A scalar type fills every leaf; otherwise the layouts must correspond - (:func:`_corresponding_leaves`). - """ - if literal_type == dtype: - return True - np_from, np_to = literal_type.np_type, dtype.np_type - if np_from.names is None and np_from.subdtype is None: - return all(np.can_cast(np_from, leaf, "safe") for leaf in _leaf_dtypes(np_to)) - reason, pairs = _corresponding_leaves(np_from, np_to) - return reason is None and all(np.can_cast(x, z, "safe") for x, z in pairs) - - -def _iter_record_leaves(np_type, python_prefix="", c_prefix=""): - """Yield ``(python_access, c_access, leaf_dtype)`` for each leaf in a record. - - A non-nested record yields one entry per field, with paths like - ``"['a']"`` (Python) and ``"a"`` (C struct). A nested record recurses - into structured fields: a top-level field ``"outer"`` whose dtype is - itself a struct with field ``"inner_a"`` yields - ``("['outer']['inner_a']", "outer.inner_a", float64_dtype)``. - - Only record-in-record nesting is recognised. Array-typed fields are - yielded as a single leaf, which the Numba path handles: the generated - expression operates on the whole sub-array and the wrapper slice-assigns - it back. The JIT path never sees such a record, because a subarray dtype - has no entry in ``NP_TO_C_TYPES`` and ``_udt_c_typedef`` bails out, so - the op keeps the cfunc for every call. - """ - for name in np_type.names: - field_dtype = np_type.fields[name][0] - py_path = f"{python_prefix}[{name!r}]" - c_path = f"{c_prefix}{name}" if not c_prefix else f"{c_prefix}.{name}" - if field_dtype.names is not None: - yield from _iter_record_leaves(field_dtype, py_path, c_path) - else: - yield py_path, c_path, field_dtype - - def _udt_c_typedef(python_name, np_type): """Build the JIT C typedef for a record or array UDT, or return ``None``. @@ -1124,7 +595,7 @@ def _make_record_func( ``leaf_paths`` is a sequence of Python access strings ``"['a']"`` or, for nested records, ``"['outer']['inner_a']"``. The generated function always returns a *flat* tuple of leaf values regardless of nesting - depth; the wrapper (in base.py) walks the same leaf paths when + depth; the wrapper (``udf._compose_wrapper_body``) walks the same leaf paths when writing the result back, so nested-record outputs land at the correct depth without nested tuple construction (which Numba can't ``setitem``-assign to a record field). @@ -1291,14 +762,14 @@ def compile_udt_binary_wrapper(op_name, py_op, dtype, dtype2, out=None): Handles two cases: - - Both sides are UDTs that pair (:func:`_check_udt_pair`): a + - Both sides are UDTs that pair (``udt._check_udt_pair``): a field-by-field (record) or element-by-element (array) op. - One side is a UDT and the other is a scalar type: the scalar is broadcast to all fields or elements (e.g., ``Point + int`` adds the int to every field). Each element is computed as the op computes it on the operands' - element dtypes (see :func:`_udt_op_types`). + element dtypes (see ``udt._udt_op_types``). ``out``, a UDT of the result's layout, makes the op store each element into that type instead, cast as a store casts (:func:`fused_store_op`). @@ -1308,7 +779,7 @@ def compile_udt_binary_wrapper(op_name, py_op, dtype, dtype2, out=None): common mistakes like passing two record UDTs with different field names. """ - from .base import _get_udt_wrapper + from .udf import _get_udt_wrapper info_x = _get_udt_info(dtype) info_y = _get_udt_info(dtype2) @@ -1333,7 +804,7 @@ def cast(in_dtype, elem): return None if in_dtype == elem.base else in_dtype if udt_dtype.np_type.names is not None: - from .base import _compile_udf_for_udt + from .udf import _compile_udf_for_udt # Use leaf paths so the same codegen handles nested-record UDTs # uniformly. A non-nested record's leaves are its top-level @@ -1409,7 +880,7 @@ def _make_record_store_wrapper(func, result_type, out_type, dtype, dtype2=None): Returns (wrapper_func, wrapper_sig). """ - from .base import _input_operand + from .udf import _input_operand nt = numba.types ns = {"numba_func": func} @@ -1500,7 +971,7 @@ def _cast_leaf_source(src, dst, expr, ns): def compile_udt_cast_wrapper(dtype, to_dtype): """Compile a cfunc wrapper that stores an element of UDT ``dtype`` as UDT ``to_dtype``. - :func:`_udt_cast_error` has accepted the pair. Each element of an array + ``udt._udt_cast_error`` has accepted the pair. Each element of an array UDT, and each record leaf, converts by :func:`_cast_leaf_source`. Returns (wrapper_func, wrapper_sig). @@ -1551,7 +1022,7 @@ def compile_udt_unary_wrapper(op_name, py_op, dtype, out=None): Returns (wrapper_func, wrapper_sig, ret_type). Raises KeyError if the dtype is not supported. """ - from .base import _get_udt_wrapper, _resolve_udt_return_type + from .udf import _get_udt_wrapper, _resolve_udt_return_type info = _get_udt_info(dtype) if info is None: @@ -1563,7 +1034,7 @@ def compile_udt_unary_wrapper(op_name, py_op, dtype, out=None): kind, _detail = info if kind == "record": - from .base import _compile_udf_for_udt + from .udf import _compile_udf_for_udt leaf_paths = [py for py, _c, _d in _iter_record_leaves(dtype.np_type)] func = _make_record_func(leaf_paths, 1, py_op) @@ -1598,12 +1069,12 @@ def fused_store_op(typed_op, out_type): The returned op computes what ``typed_op`` computes, in the same type, and casts each element or leaf into ``out_type`` as it writes it, as a - store casts (``base._store_cast_op``). So a result goes into an object + store casts (``udt._store_cast_op``). So a result goes into an object of another UDT type in the one GraphBLAS call that computes it, with no temporary. Only the built-in ops that lift to UDTs have such a form; - ``base._store_cast_op`` has already accepted the pair of types. + ``udt._store_cast_op`` has already accepted the pair of types. """ - from .base import _finalize_udt_op + from .udf import _finalize_udt_op parent = getattr(typed_op, "parent", None) if parent is None or parent._anonymous: @@ -1897,7 +1368,7 @@ def _make_jit_c_definition(op_name, py_op, dtype, arity): Returns ``(c_name, c_defn)``, or ``None`` if the dtype can't be expressed in C (unsupported field types, or names that collide with C reserved words). An op not defined on the field types (``min`` on a complex field) never - gets here: ``_element_types`` rejects it when the op is typed. + gets here: ``udt._element_types`` rejects it when the op is typed. """ np_type = dtype.np_type # Prefer the C name SuiteSparse already has for this type. ``GxB_JIT_C_NAME`` diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index f5bf0fc44..bffedeedd 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -35,7 +35,7 @@ if _has_numba: import numba - from .base import ( + from .udf import ( _bool_to_int8, _compile_udf_for_udt, _finalize_udt_op, @@ -120,8 +120,8 @@ def _one(x): if _has_numba: - from .udt_utils import BUILTIN_UDT_UNARY_OPS as _BUILTIN_UDT_UNARY_OPS - from .udt_utils import ( + from ..udt import BUILTIN_UDT_UNARY_OPS as _BUILTIN_UDT_UNARY_OPS + from .udt_codegen import ( _has_jit_set, _maybe_warn_jit_skipped, compile_udt_cast_wrapper, @@ -504,8 +504,8 @@ def __init__( def _udt_cast_op(dtype, to_dtype): """Return the UnaryOp storing an element of UDT ``dtype`` as UDT ``to_dtype``. - ``base._store_cast_op`` decides when to use it, after - ``udt_utils._udt_cast_error`` has accepted the pair. It is compiled once + ``udt._store_cast_op`` decides when to use it, after + ``udt._udt_cast_error`` has accepted the pair. It is compiled once per pair and has no JIT kernel. """ parent = _udt_casts.get(to_dtype) diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 4420ee75f..109067250 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -1,6 +1,4 @@ import itertools -import math -from functools import lru_cache import numpy as np @@ -9,10 +7,26 @@ from ..exceptions import EmptyObject, check_status from . import _has_numba, _supports_udfs, automethods, ffi, lib, utils from .base import BaseExpression, BaseType, call -from .dtypes import _view_if_same_layout from .expr import AmbiguousAssignOrExtract -from .operator import IndexUnaryOp, SelectOp, TypedOpBase, binary_from_string, get_typed_op -from .operator.udt_utils import _LITERAL_KINDS, _weak_literal_leaf +from .literals import ( + _SINGLE_PRECISION, + _check_literal_fits, + _check_scalar_fits, + _literal_operand, + _literal_store_dtype, + _literal_type, +) +from .operator import get_typed_op +from .udt import ( + _element_bytes, + _element_shapes_differ, + _element_value, + _ewise_add_needs_cast, + _is_lifted_op, + _literal_shape_differs, + _scalar_ewise_add_converted, + _udt_eq_without_numba, +) from .utils import _Pointer, output_type, wrapdoc if _supports_udfs: @@ -373,12 +387,7 @@ def value(self): dtype_name = "UDT" if is_udt else self.dtype.name call(f"GrB_Scalar_extractElement_{dtype_name}", [_Pointer(scalar), self]) if is_udt: - np_type = self.dtype.np_type - rv = np.array(ffi.buffer(scalar.gb_obj[0 : np_type.itemsize])) - if np_type.subdtype is None: - return rv.view(np_type)[0] - base, shape = np_type.subdtype - return rv.view(base).reshape(shape) + return _element_value(self.dtype, scalar.gb_obj) return scalar.gb_obj[0] @value.setter @@ -389,35 +398,7 @@ def value(self, val): if output_type(val) is Scalar: val = val.value # raise below if wrong type (as determined by cffi) if self.dtype._is_udt: - np_type = self.dtype.np_type - # Zeros, not empty: when numpy writes only the fields, a - # record's alignment padding is whatever the allocation - # happened to contain. ``_udt_bytes`` zeros the padding in the - # other direction, when numpy copies whole items instead. - if np_type.subdtype is None: - arr = np.zeros(1, dtype=np_type) - if isinstance(val, dict) and np_type.names is not None: - val = _dict_to_record(np_type, val) - else: - arr = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) - val = _view_if_same_layout(val, arr.dtype) - if _has_narrow_float_leaf(np_type) and np.geterr()["over"] == "warn": - # A float too large for a float32 leaf is infinity, and numpy - # warns; warn at the caller's line instead, as a built-in - # FP32 does, so the default filter shows each place it happens. - try: - with np.errstate(over="raise"): - arr[:] = val - except FloatingPointError: - utils._warn_from_caller( - f"overflow encountered in cast: {val!r} as {self.dtype}", - RuntimeWarning, - ) - with np.errstate(over="ignore"): - arr[:] = val - else: - arr[:] = val - self.gb_obj[0 : np_type.itemsize] = _udt_bytes(np_type, arr) + self.gb_obj[0 : self.dtype.np_type.itemsize] = _element_bytes(self.dtype, val) else: self.gb_obj[0] = val self._empty = False @@ -505,7 +486,7 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=None, name=None): new_scalar = Scalar(dtype, is_cscalar=is_cscalar, name=name) if not clear and not self._is_empty: if self.dtype._is_udt: - # Cast as storing it would (``base._store_cast_op``), not as numpy + # Cast as storing it would (``udt._store_cast_op``), not as numpy # does; a UDT does not cast to a built-in dtype at all. new_scalar << self elif new_scalar.is_cscalar and not new_scalar.dtype._is_udt: @@ -723,18 +704,7 @@ def _ewise_add(self, other, op=monoid.plus, is_infix=False): op = get_typed_op(op, self.dtype, other.dtype, kind="binary") self._expect_op(op, ("BinaryOp", "Monoid"), within=method_name, argname="op") if _ewise_add_needs_cast(op, self.dtype, other.dtype): - if not _is_lifted_op(op.parent): - raise _ewise_add_cast_error(op, self.dtype, other.dtype) - # GraphBLAS casts an unpaired entry to the result type, which a UDT - # cannot be. Convert both to the result type first, which changes no - # value, and apply the op in that one type: an int8 UDT plus 0.5 is - # float64, and an empty Scalar still adds as 0. A Scalar is one - # element, so the converted copies cost nothing. - if self.dtype != op.return_type: - self = _converted_to(self, op.return_type, name="") - if other.dtype != op.return_type: - other = _converted_to(other, op.return_type, name="") - op = get_typed_op(op.parent, self.dtype, other.dtype, kind="binary") + self, other, op = _scalar_ewise_add_converted(self, other, op) return ScalarExpression( method_name, f"GrB_Vector_eWiseAdd_{op.opclass}", @@ -891,52 +861,8 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): temp_op = get_typed_op(op, self.dtype, other.dtype, kind="binary") - left_dtype = temp_op.type - dtype = left_dtype if left_dtype._is_udt else None - if type(left_default) is not Scalar: - if dtype is not None: - _check_literal_fits(dtype, left_default) - else: - left_default, dtype = _literal_operand(left_dtype, left_default, op) - try: - left = Scalar.from_value( - left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar - ) - except TypeError: - left = self._expect_type( - left_default, - Scalar, - within=method_name, - keyword_name="left_default", - extra_message="Literal scalars also accepted.", - op=op, - ) - else: - _check_scalar_fits(dtype, left_default) - left = _as_scalar(left_default, dtype, is_cscalar=False) # pragma: is_grbscalar - right_dtype = temp_op.type2 - dtype = right_dtype if right_dtype._is_udt else None - if type(right_default) is not Scalar: - if dtype is not None: - _check_literal_fits(dtype, right_default) - else: - right_default, dtype = _literal_operand(right_dtype, right_default, op) - try: - right = Scalar.from_value( - right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar - ) - except TypeError: - right = self._expect_type( - right_default, - Scalar, - within=method_name, - keyword_name="right_default", - extra_message="Literal scalars also accepted.", - op=op, - ) - else: - _check_scalar_fits(dtype, right_default) - right = _as_scalar(right_default, dtype, is_cscalar=False) # pragma: is_grbscalar + left = _union_default(self, left_default, temp_op.type, op, "left_default") + right = _union_default(self, right_default, temp_op.type2, op, "right_default") op1 = get_typed_op(op, self.dtype, right.dtype, kind="binary") op2 = get_typed_op(op, left.dtype, other.dtype, kind="binary") @@ -1209,422 +1135,6 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=False, name=None, **opts): # End auto-generated code: Scalar -def _literal_operand(dtype, value, op): - """Return literal ``value`` and the dtype to make it, as ``op``'s operand beside ``dtype``. - - A dtype of ``None`` means the value's own dtype. This is where - python-graphblas decides how a literal operand is typed; beside a built-in - dtype, see :func:`_weak_builtin_literal_dtype`. - - Beside a UDT, a literal becomes an element of that UDT, as it must for a - user's op, when its type allows it (:func:`_check_literal_fits`). With a - built-in op that lifts to UDTs, a literal is an operand in its own right, - typed as numpy 2 types it (NEP 50), by :func:`_literal_type`: - - - a numpy scalar or array keeps its dtype (strong); an array's type is the - array layout of its dtype (``udt_utils._array_literal_udt``); - - a Python ``bool``, ``int``, ``float`` or ``complex`` is weak: in each - element or leaf it takes that element's dtype when it fits the kind, so - ``int8_udt + 1`` stays ``int8_udt`` and ``int8_udt += 1`` can be stored, - while ``int8_udt * 2.5`` has float64 elements (see - ``udt_utils._weak_literal_udt``), where converting it to the UDT would - multiply by 2; - - a tuple, list or dict of Python numbers, laid out like an element, is - weak in the same way, leaf by leaf: ``int8_udt + (1, 2, 3)`` stays - ``int8_udt`` and ``+ (0.5, 1.5, 2.5)`` has float64 elements. - - ``eq`` and ``ne`` type a literal the same way, as numpy compares, so - ``fp32_udt == 0.1`` compares in float32, as ``fp32_vec == 0.1`` does. - - A comparison keeps no result type, so a Python int beyond an integer - element's range is compared exactly, as numpy 2 compares it, as infinity - of its sign: every element compares with that infinity as with the int, - while a type that held both would round them (in FP64, the INT64 - ``2**63 - 1`` and the int ``2**63`` are equal). So ``int64_vec < 2**63`` is - all True and ``int8_vec == 300`` all False. Beside a UDT, where only ``eq`` - and ``ne`` compare, a literal with such an int in an integer leaf equals no - element: ``int8_udt == 300`` is False, not an error. - """ - if getattr(op, "is_positional", False): - return value, None - if not dtype._is_udt: - literal_dtype = _weak_builtin_literal_dtype(dtype, value, op) - if ( - literal_dtype is not dtype - and literal_dtype is not None - and type(value) is int - and dtype.np_type.kind in "iu" - ): - # Only an int beyond dtype's range takes another type here, typed so - # by a comparison or by an op that gives another type (truediv). - if isinstance(op, str): - op = binary_from_string(op) - if op.name in _COMPARISONS: - return (math.inf if value > 0 else -math.inf), FP64 - return value, literal_dtype - if getattr(op, "_udt_types", True) is None: - # An op that does not take UDTs at all (``lt``) reports that itself. - return value, dtype - from .operator.binary import BinaryOp - from .operator.monoid import Monoid - from .operator.udt_utils import BUILTIN_UDT_BINARY_OPS, _int_beyond_leaf - - if isinstance(op, Monoid) and not op._anonymous: - # Element-wise, a Monoid is its BinaryOp (``get_typed_op`` types it so). - op = op.binaryop - if type(op) is BinaryOp and not op._anonymous: - is_comparison = op.name in ("eq", "ne") - if is_comparison or op.name in BUILTIN_UDT_BINARY_OPS: - # truediv gives float leaves whatever the int, so, as a comparison, it - # types an int out of an integer leaf's range exactly instead of raising. - exact = is_comparison or op.name == "truediv" - literal_type = _literal_type(dtype, value, exact=exact) - if literal_type is not None: - if is_comparison and not isinstance(value, (np.generic, np.ndarray)): - np_type = dtype.np_type - parts = _dict_to_record(np_type, value) if isinstance(value, dict) else value - beyond = _int_beyond_leaf(np_type, parts) - if beyond is not None: - # As infinity in every leaf, it still equals no element. - value = math.inf if beyond > 0 else -math.inf - literal_type = _literal_type(dtype, value) - return value, literal_type - _check_literal_fits(dtype, value) - return value, dtype - - -# 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)]) - - -# The types a Python float or complex can overflow. -_SINGLE_PRECISION = {np.dtype(np.float32), np.dtype(np.complex64)} - - -def _weak_builtin_literal_dtype(dtype, value, op=None): - """Return the dtype of Python number ``value`` beside built-in ``dtype``, or ``None``. - - A Python ``bool``, ``int``, ``float`` or ``complex`` is weak, as numpy 2 - types it (NEP 50), by the table UDT leaves use - (``udt_utils._weak_literal_leaf``): ``int8_vec + 1`` is INT8, - ``fp32_vec * 0.5`` is FP32, ``int8_vec * 0.5`` is FP64, and - ``int8_vec + 300`` raises ``OverflowError``. python-graphblas applies the - table itself, so numpy 1 gives the same answers. ``None`` keeps the value's - own dtype (strong): for anything but a Python number, for a typed op, which - fixes the operand types itself, and for the thunk of an index op that does - not compare it with the values, such as a row index. - - A comparison keeps no result type, so an int beyond ``dtype``'s range - takes a type that holds it rather than raise (:func:`_literal_operand` then - compares it as infinity). ``op`` of ``None``, for a fill value, does the - same, so the value is not lost. - """ - kind = _LITERAL_KINDS.get(type(value)) - if kind is None: - return None - if isinstance(op, str): - try: - op = binary_from_string(op) - except Exception: - return None - exact = op is None - if op is not None: - if isinstance(op, TypedOpBase): - return None - exact = getattr(op, "name", None) in _COMPARISONS and not op._anonymous - if isinstance(op, (IndexUnaryOp, SelectOp)) and not exact: - return None - try: - np_type = _weak_literal_leaf([value], kind, dtype.np_type, exact=exact) - except OverflowError: - if exact or not _gives_another_type(op, dtype): - raise - # An op whose result is not the operand's type has no narrow result to - # keep, so an int out of range is typed exactly, as for a comparison: - # int8_vec / 300 is FP64, as numpy divides integers in float64. - np_type = _weak_literal_leaf([value], kind, dtype.np_type, exact=True) - if np_type == dtype.np_type: - rv = dtype - else: - try: - rv = lookup_dtype(np_type) - except ValueError: # A complex literal where the backend has no complex types - return None - if kind != "b" and rv.np_type in _SINGLE_PRECISION: - import struct - - parts = (value.real, value.imag) if kind == "c" else (float(value),) - try: - for part in parts: - struct.pack(" pos: - ranges.append((pos, start)) - pos = max(pos, stop) - if pos < np_type.itemsize: - ranges.append((pos, np_type.itemsize)) - return tuple(ranges) - - -def _collect_data_ranges(np_type, offset, out): - """Append the ``(start, stop)`` byte ranges of ``np_type`` that hold field data.""" - if np_type.subdtype is not None: - base = np_type.subdtype[0] - for start in range(offset, offset + np_type.itemsize, base.itemsize): - _collect_data_ranges(base, start, out) - elif np_type.names is not None: - for field_type, field_offset in (np_type.fields[name][:2] for name in np_type.names): - _collect_data_ranges(field_type, offset + field_offset, out) + dtype = input_dtype if input_dtype._is_udt else None + if type(default) is Scalar: + _check_scalar_fits(dtype, default) + return _as_scalar(default, dtype, is_cscalar=False) # pragma: is_grbscalar + if dtype is not None: + _check_literal_fits(dtype, default) else: - out.append((offset, offset + np_type.itemsize)) - - -def _udt_bytes(np_type, arr): - """The bytes of a UDT buffer, with any alignment padding zeroed. - - ``tobytes()`` flattens any-rank numpy buffers (in particular, multi-dim - array UDTs like ``FP64[2, 3]``) to a 1-D byte buffer that cffi can copy - into GrB_Scalar storage. - - numpy promises nothing about the padding bytes it leaves behind, and which - way it goes has changed across releases: numpy 2.5 made ``arr[:] = record`` - a whole-item copy out of an uninitialized temporary, so even a buffer from - ``np.zeros`` comes back with junk between the fields. A caller's own array - can carry junk padding in too. Since the whole record is copied into the - GrB_Scalar and read back out of ``Scalar.value``, anything that compares or - hashes those raw bytes would otherwise see an equal-valued scalar as - unequal. - """ - raw = arr.tobytes() - ranges = _padding_ranges(np_type) - if not ranges: - return raw - buf = bytearray(raw) - for start, stop in ranges: - buf[start:stop] = bytes(stop - start) - return bytes(buf) + default, dtype = _literal_operand(input_dtype, default, op) + try: + return Scalar.from_value( + default, + dtype, + is_cscalar=False, # pragma: is_grbscalar + name="", + ) + except TypeError: + return obj._expect_type( + default, + Scalar, + within="ewise_union", + keyword_name=keyword_name, + extra_message="Literal scalars also accepted.", + op=op, + ) _MATERIALIZE = Scalar.from_value(lib.GrB_MATERIALIZE, is_cscalar=True, name="GrB_MATERIALIZE") diff --git a/graphblas/core/ss/dtypes.py b/graphblas/core/ss/dtypes.py index 4dafd9c82..029721590 100644 --- a/graphblas/core/ss/dtypes.py +++ b/graphblas/core/ss/dtypes.py @@ -29,7 +29,7 @@ def register_new(name, jit_c_definition, *, np_type=None): ) if not name.isidentifier(): raise ValueError(f"`name` argument must be a valid Python identifier; got: {name!r}") - from ..operator.udt_utils import _jit_c_type_defs + from ..operator.udt_codegen import _jit_c_type_defs # SuiteSparse's JIT knows a type by its C name alone, so it must not be # another type's C name either. diff --git a/graphblas/core/ss/jit_config.py b/graphblas/core/ss/jit_config.py index 8f3aa4a9a..677b979d8 100644 --- a/graphblas/core/ss/jit_config.py +++ b/graphblas/core/ss/jit_config.py @@ -387,7 +387,7 @@ def _compiling_c_source(name): def _maybe_warn_no_jit(*, op_name="", dtype_name=""): """Emit a ``NoJITWarning`` (once per ``(op_name, dtype_name)``) when UDT auto-lift falls back. - The only caller is ``udt_utils._maybe_warn_jit_skipped``, which fires + The only caller is ``udt_codegen._maybe_warn_jit_skipped``, which fires after C codegen returned nothing. Codegen reads the dtype and neither the compiler path nor ``jit_c_control``, so the dtype is the cause here even on a host that could not have compiled a kernel anyway. Do not re-derive diff --git a/graphblas/core/udt.py b/graphblas/core/udt.py new file mode 100644 index 000000000..f24a3ab11 --- /dev/null +++ b/graphblas/core/udt.py @@ -0,0 +1,809 @@ +"""The rules for user-defined types (UDTs), apart from generating code for them. + +GraphBLAS knows a UDT only by its size: it cannot cast one to or from any other +type, and it copies an element of one as raw bytes. The rules here fill that +gap, in this order: + +- which built-in ops lift to UDTs (``BUILTIN_UDT_BINARY_OPS``), and the layout + of an element (:func:`_iter_record_leaves`); +- which operands an op on UDTs takes and the type of its result + (:func:`_check_udt_pair`, :func:`_udt_op_types`); +- what may be stored in an object of another type, and how it is converted on + the way (:func:`_store_cast_op`); +- when ``ewise_add`` would need GraphBLAS to cast a UDT; +- how ``isequal``, and ``==`` without Numba, compare elements; +- how an element's value becomes bytes and back, for ``Scalar.value``. + +``operator.udt_codegen`` generates the operators that compute on UDTs, +``operator.udf`` wraps a user's function on them, and ``literals`` types a +literal operand beside one. + +``operator.udt_codegen`` imports this module while the operator package is still +loading, so a function here that needs an operator, ``base`` or ``scalar`` +imports it inside the function, once it knows a UDT is involved. +""" + +import copy +from functools import lru_cache, reduce +from operator import mul + +import numpy as np + +from ..exceptions import DomainMismatch +from . import _has_numba, ffi, utils +from .dtypes import _registry, _view_if_same_layout, lookup_dtype +from .utils import output_type + +# The built-in ops that lift to UDTs, each with the Python operator or function +# the generated code applies. +BUILTIN_UDT_BINARY_OPS = { + "plus": "+", + "minus": "-", + "times": "*", + "truediv": "/", + "floordiv": "//", + "min": "min", + "max": "max", +} +BUILTIN_UDT_UNARY_OPS = { + "ainv": "-", + "abs": "abs", +} + + +# Ops with no meaning on complex fields: no ordering for min and max, and no +# integer mod for floordiv. +_OPS_NOT_FOR_COMPLEX = frozenset({"min", "max", "floordiv"}) + + +# The layout of an element. + + +def _get_udt_info(dtype): + """Return ('record', field_names) or ('array', (base_dtype, flat_size)) or None.""" + np_type = dtype.np_type + if np_type.names is not None: + return "record", np_type.names + if np_type.subdtype is not None: + base, shape = np_type.subdtype + return "array", (base, reduce(mul, shape)) + return None + + +def _iter_record_leaves(np_type, python_prefix="", c_prefix=""): + """Yield ``(python_access, c_access, leaf_dtype)`` for each leaf in a record. + + A non-nested record yields one entry per field, with paths like + ``"['a']"`` (Python) and ``"a"`` (C struct). A nested record recurses + into structured fields: a top-level field ``"outer"`` whose dtype is + itself a struct with field ``"inner_a"`` yields + ``("['outer']['inner_a']", "outer.inner_a", float64_dtype)``. + + Only record-in-record nesting is recognised. Array-typed fields are + yielded as a single leaf, which the Numba path handles: the generated + expression operates on the whole sub-array and the wrapper slice-assigns + it back. The JIT path never sees such a record, because a subarray dtype + has no entry in ``udt_codegen.NP_TO_C_TYPES`` and ``udt_codegen._udt_c_typedef`` + bails out, so the op keeps the cfunc for every call. + """ + for name in np_type.names: + field_dtype = np_type.fields[name][0] + py_path = f"{python_prefix}[{name!r}]" + c_path = f"{c_prefix}{name}" if not c_prefix else f"{c_prefix}.{name}" + if field_dtype.names is not None: + yield from _iter_record_leaves(field_dtype, py_path, c_path) + else: + yield py_path, c_path, field_dtype + + +def _strip_leading_ones(shape): + i = 0 + while i < len(shape) and shape[i] == 1: + i += 1 + return shape[i:] + + +# Which operands an op on UDTs takes, and the type of its result. + + +def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): + """Raise ``KeyError`` if two UDT operands disagree on shape. + + ``info_x`` and ``info_y`` are ``_get_udt_info`` results. When either is + ``None`` (one side is a scalar broadcast), no check fires. Without this + gate the generated wrapper code raises a cryptic Numba ``TypingError`` + on the first field access, or, for array fields that do not broadcast, + fails inside the cfunc without telling the caller. Element dtypes may + differ; :func:`_udt_op_types` promotes them. + """ + if info_x is None or info_y is None: + return + kind_x, detail_x = info_x + kind_y, detail_y = info_y + if kind_x != kind_y: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"cannot mix record and array UDTs in a single element-wise op." + ) + if kind_x == "record": + if detail_x != detail_y: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"record UDTs must share field names, in the same order; got " + f"{list(detail_x)} vs {list(detail_y)}." + ) + # Matching top-level names is not enough. The codegen reads both + # operands at the left one's leaf paths, so a field that is a + # sub-record on one side and a scalar on the other, or sub-records + # with different names, fails Numba's typing: a compile error + # reported for what is really the same shape disagreement the check + # above reports as a KeyError. The paths must also come in the same + # order, as the top-level names must, so that leaf ``i`` of one + # operand is leaf ``i`` of the other. + leaves_x = list(_iter_record_leaves(dtype.np_type)) + leaves_y = list(_iter_record_leaves(dtype2.np_type)) + if [py for py, _c, _d in leaves_x] != [py for py, _c, _d in leaves_y]: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"record UDTs must nest the same way, with the same field names in the " + f"same order at every level; got {[c for _py, c, _d in leaves_x]} vs " + f"{[c for _py, c, _d in leaves_y]}." + ) + # Each field must also have the same shape in both (field dtypes may + # differ), as numpy requires before it promotes or compares two record + # dtypes. Shapes that do not broadcast would raise inside the cfunc, + # unseen, and leave whatever was in the output buffer. Shapes that do + # broadcast would work one way round only, since the arithmetic ops + # write into a field shaped like the left operand's: ``(3,)`` plus + # ``(1,)`` works, ``(1,)`` plus ``(3,)`` does not. + for (_py, c, leaf_x), (_py_y, _c_y, leaf_y) in zip(leaves_x, leaves_y, strict=True): + if leaf_x.shape != leaf_y.shape: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"record UDTs must have the same shape for each field; field {c} is " + f"{leaf_x.shape} vs {leaf_y.shape}." + ) + return + # Array UDTs broadcast as numpy arrays do, so ``(3, 1)`` with ``(1, 4)`` + # gives ``(3, 4)``; the generated code reads each operand through a flat + # index map (``udt_codegen._array_operand_source``). Shapes such as ``(2, 3)`` with + # ``(3, 2)`` do not broadcast, as numpy refuses them. + (base_x, shape_x), (base_y, shape_y) = dtype.np_type.subdtype, dtype2.np_type.subdtype + try: + np.broadcast_shapes(shape_x, shape_y) + except ValueError: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): array UDTs must " + f"have shapes that broadcast together; got {base_x}{list(shape_x)} vs " + f"{base_y}{list(shape_y)}." + ) from None + + +def _builtin_dtype(np_type): + """Return the built-in ``DataType`` for ``np_type``, or ``None`` if it is not one.""" + dt = _registry.get(np_type) + return None if dt is None or dt._is_udt else dt + + +def _complex_op_error(op_name, dtype, dtype2): + return KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"this op is not defined on complex fields. Use ``binary.plus``, " + f"``minus``, ``times``, or ``truediv`` for complex element-wise " + f"arithmetic, or register a custom binary op." + ) + + +def _element_types(op_name, dtype, dtype2, elem_x, elem_y): + """Return ``(in_x, in_y, out)``, the numpy dtypes ``binary.{op_name}`` uses on one element. + + They come from the op's typing on the built-in dtypes, so each element of an + array UDT, or each record leaf, is promoted as a Vector of that dtype would + be, and ``truediv`` on integers gives floats. ``in_x`` and ``in_y`` are the + types GraphBLAS casts the operands to before it applies the op. The + generated code casts to them too, because Numba promotes differently: it + adds uint64 and int64 in int64, where GraphBLAS uses float64. ``dtype`` and + ``dtype2`` are the operands, named in errors. Two elements that are not + both built-in dtypes (bytes, datetime, records inside an array field) have + no promotion rule, so they must be the same dtype, which the result keeps. + """ + # ``min``/``max``/``floordiv`` have no defined semantics on complex + # operands (no ordering, no integer mod). Reject early with a clear + # message; otherwise Numba's cfunc compile blows up several frames + # down with ``NotImplementedError: No definition for lowering lt``. + if op_name in _OPS_NOT_FOR_COMPLEX and "c" in (elem_x.kind, elem_y.kind): + raise _complex_op_error(op_name, dtype, dtype2) + builtin_x = _builtin_dtype(elem_x) + builtin_y = _builtin_dtype(elem_y) + if builtin_x is None or builtin_y is None: + if elem_x == elem_y: + return elem_x, elem_y, elem_x + # Builds without complex dtypes (the vanilla backend) still lift complex + # fields, so promote them as the built-in typing does where it has them. + if "c" in (elem_x.kind, elem_y.kind) and {elem_x.kind, elem_y.kind} <= set("biufc"): + promoted = np.promote_types(elem_x, elem_y) + return promoted, promoted, promoted + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): elements of " + f"{elem_x} and {elem_y} have no common type." + ) + from .. import binary + from .operator import get_typed_op + + typed = get_typed_op(getattr(binary, op_name), builtin_x, builtin_y) + return typed.type.np_type, typed.type2.np_type, typed.return_type.np_type + + +def _udt_op_types(op_name, dtype, dtype2): + """Return ``(ret_type, input_types)`` for ``binary.{op_name}`` on ``dtype`` and ``dtype2``. + + At least one operand is a record or array UDT, and ``_check_udt_pair`` has + accepted the pair; the other may be a plain scalar type. Each element, or + record leaf, follows :func:`_element_types`, so the result never narrows: + an int64 array UDT plus ``0.5`` is ``FP64[3]``. ``ret_type`` is an + operand's own type when that holds the result, so same-type arithmetic + keeps its UDT. Otherwise it is the promoted layout, which is the UDT a user + registered with that layout if there is one, else an anonymous one. + ``input_types`` holds the ``(in_x, in_y)`` pair of each record leaf, in + leaf order, or the single pair for an array UDT's elements. + """ + np_x, np_y = dtype.np_type, dtype2.np_type + if np_x.names is not None or np_y.names is not None: + leaves_x = list(_iter_record_leaves(np_x)) if np_x.names is not None else None + leaves_y = list(_iter_record_leaves(np_y)) if np_y.names is not None else None + input_types = [] + result_leaves = [] + for i, (_py, _c, leaf) in enumerate(leaves_x or leaves_y): + elem_x = np_x if leaves_x is None else leaves_x[i][2].base + elem_y = np_y if leaves_y is None else leaves_y[i][2].base + in_x, in_y, elem = _element_types(op_name, dtype, dtype2, elem_x, elem_y) + input_types.append((in_x, in_y)) + result_leaves.append(np.dtype((elem, leaf.shape)) if leaf.shape else elem) + for operand, leaves in [(dtype, leaves_x), (dtype2, leaves_y)]: + if leaves is not None and all( + result_leaf == leaf + for result_leaf, (_py, _c, leaf) in zip(result_leaves, leaves, strict=True) + ): + return operand, input_types + records = [np_type for np_type in (np_x, np_y) if np_type.names is not None] + return lookup_dtype(_promoted_record(records, iter(result_leaves))), input_types + base_x, shape_x = np_x.subdtype or (np_x, ()) + base_y, shape_y = np_y.subdtype or (np_y, ()) + in_x, in_y, elem = _element_types(op_name, dtype, dtype2, base_x, base_y) + result = np.dtype((elem, np.broadcast_shapes(shape_x, shape_y))) + for operand in (dtype, dtype2): + if operand.np_type == result: + return operand, [(in_x, in_y)] + return lookup_dtype(result), [(in_x, in_y)] + + +def _promoted_record(records, result_leaves): + """Build a record laid out like ``records`` with ``result_leaves`` for its leaves. + + ``records`` holds one or two record dtypes with the same field names and + nesting. The new record aligns its fields, at each level, if any of them + does there, so a C-compatible operand gives a C-compatible result. + """ + first = records[0] + fields = [] + for name in first.names: + field = first.fields[name][0] + if field.names is not None: + field = _promoted_record([r.fields[name][0] for r in records], result_leaves) + else: + field = next(result_leaves) + fields.append((name, field)) + return np.dtype(fields, align=any(r.isalignedstruct for r in records)) + + +# Storing a value of one type in an object of another. + + +def _corresponding_leaves(np_x, np_z): + """Return ``(reason, pairs)`` for putting an element of ``np_x`` into one of ``np_z``. + + ``reason`` says why the two layouts do not correspond, or is ``None``, and + then ``pairs`` holds the base dtypes of each pair of leaves. Records must + line up as the operands of a lifted op do (:func:`_check_udt_pair`): the + same field names, order and nesting and the same shape for each field. + Arrays must have the same shape apart from leading axes of length 1; + operands broadcast, but a stored element keeps every value in its place. + numpy casts records by position whatever their names, and repeats or + drops elements of an array field of another length, which would put + values in places nobody chose. + """ + if (np_x.names is None) != (np_z.names is None): + return "a record UDT and an array UDT do not correspond", None + if np_x.names is None: + if np_x.subdtype is None or np_z.subdtype is None: + return "only record and array UDTs cast", None + (base_x, shape_x), (base_z, shape_z) = np_x.subdtype, np_z.subdtype + if _strip_leading_ones(shape_x) != _strip_leading_ones(shape_z): + return "array UDTs must have the same shape, apart from leading axes of length 1", None + return None, [(base_x, base_z)] + leaves_x = list(_iter_record_leaves(np_x)) + leaves_z = list(_iter_record_leaves(np_z)) + if [py for py, _c, _d in leaves_x] != [py for py, _c, _d in leaves_z]: + return "records must have the same field names, in the same order, at every level", None + pairs = [(x, z) for (_, _, x), (_, _, z) in zip(leaves_x, leaves_z, strict=True)] + if any(x.shape != z.shape for x, z in pairs): + return "records must have the same shape for each field", None + return None, [(x.base, z.base) for x, z in pairs] + + +def _udt_cast_error(dtype, to_dtype): + """Return why an element of UDT ``dtype`` cannot be stored as UDT ``to_dtype``, or ``None``. + + The layouts must correspond (:func:`_corresponding_leaves`), and each pair + of elements must be built-in numbers, or the same type. + """ + reason, pairs = _corresponding_leaves(dtype.np_type, to_dtype.np_type) + for x, z in pairs or (): + if x != z and not (x.kind in "biufc" and z.kind in "biufc"): + return f"elements of {x} do not cast to {z}" + return reason + + +def _leaf_dtypes(np_type): + """Yield the base dtype of each leaf of ``np_type``: a record's fields, an array's element.""" + if np_type.names is not None: + for _py, _c, leaf in _iter_record_leaves(np_type): + yield leaf.base + else: + yield np_type.base + + +def _casts_safely_into(literal_type, dtype): + """Whether each leaf of ``literal_type`` casts safely (``np.can_cast``) into UDT ``dtype``'s. + + A scalar type fills every leaf; otherwise the layouts must correspond + (:func:`_corresponding_leaves`). + """ + if literal_type == dtype: + return True + np_from, np_to = literal_type.np_type, dtype.np_type + if np_from.names is None and np_from.subdtype is None: + return all(np.can_cast(np_from, leaf, "safe") for leaf in _leaf_dtypes(np_to)) + reason, pairs = _corresponding_leaves(np_from, np_to) + return reason is None and all(np.can_cast(x, z, "safe") for x, z in pairs) + + +def _store_cast_op(result_dtype, output_dtype): + """Return the UnaryOp that converts a result for an object of ``output_dtype``, or ``None``. + + GraphBLAS casts between built-in types on store (C's casts, with a float + saturating into an integer type), but it cannot cast to or from a UDT; + left to GraphBLAS, that is a bare GrB_DOMAIN_MISMATCH or, for a single + element, bytes read as the wrong type. So when a UDT is involved, the + result is stored by applying the op returned here, which casts each + element or record leaf the way GraphBLAS casts built-in types, between + UDTs whose layouts correspond (:func:`_udt_cast_error`). Any other + pair raises ``DomainMismatch``. ``None`` means GraphBLAS stores the result + as it is. This is the one place that decides what may be stored where. + """ + if result_dtype is output_dtype or not (result_dtype._is_udt or output_dtype._is_udt): + return None + reason = "UDTs do not cast" + if _has_numba and result_dtype._is_udt and output_dtype._is_udt: + if (reason := _udt_cast_error(result_dtype, output_dtype)) is None: + from .operator.unary import _udt_cast_op + + return _udt_cast_op(result_dtype, output_dtype) + elif _has_numba: + reason = "a UDT casts only to another UDT" + raise DomainMismatch( + f"cannot store a result of type {result_dtype} in an object of type " + f"{output_dtype}: {reason}. Use .new() for an object of the result's type, or " + "compute the result in the type you want to store." + ) + + +def _fused_store_expr(expr, output_dtype): + """Return ``expr`` computing straight into ``output_dtype``, or raise ``DomainMismatch``. + + A result of another UDT type is stored by the op that computes it, which + casts each element as it writes it (``udt_codegen.fused_store_op``), so no + temporary of the result's own type is made: data GraphBLAS works on is + often too large to hold twice. That works for an element-wise + computation by a built-in op that lifts to UDTs (``apply``, + ``ewise_mult``, ``ewise_union``). ``ewise_add`` copies an entry present in + only one input without the op, and a semiring or reduction would add in + the narrower type, so those and any other computation raise. + """ + from .operator.udt_codegen import fused_store_op # needs Numba + + op, cfunc_name = expr.op, expr.cfunc_name + if op is not None and op.opclass == "Monoid": + # Element-wise, a Monoid is its BinaryOp. + op, cfunc_name = op.binaryop, cfunc_name.replace("_Monoid", "_BinaryOp") + fused = None + if any(k in cfunc_name for k in ("_apply", "eWiseMult", "eWiseUnion")): + fused = fused_store_op(op, output_dtype) + if fused is None: + method = "extraction" if expr.method_name == "__getitem__" else expr.method_name + what = method if op is None else f"{method} with {op.name}" + raise DomainMismatch( + f"cannot store a result of type {expr.dtype} in an object of type " + f"{output_dtype}: GraphBLAS cannot cast a UDT, and only a built-in op that lifts " + "to UDTs, in apply, ewise_mult or ewise_union, can write another type as it " + f"computes ({what} cannot), so this would need a temporary copy of the whole " + "result. Make it explicitly by storing the computed result, as in " + "`obj << expr.new()`, which casts it in one more pass." + ) + rv = copy.copy(expr) + rv.op, rv.cfunc_name, rv.dtype = fused, cfunc_name, output_dtype + return rv + + +def _scalar_accum_op(scalar, accum): + """Return the op to accumulate into ``scalar`` with, as ``scalar.ewise_add`` runs it. + + ``__call__`` types ``accum`` on the object's dtype. For a UDT Scalar, that + would convert a literal into the UDT (``s += 0.5``), where a Vector + computes ``s + 0.5`` in its own type and casts as it stores; the untyped + op does the same here. + """ + return accum.parent if scalar.dtype._is_udt else accum + + +def _cast_for_store(value, dtype): + """Return Scalar, Vector or Matrix ``value``, or its expression, to assign as ``dtype``. + + GraphBLAS's assign and setElement take a UDT value as raw bytes, so one of + another type would be read as the wrong type; :func:`_store_cast_op` + casts it or raises. Assign cannot cast as it writes, so a value of + another UDT type needs a converted copy: one element for a Scalar, which + is made here, but the whole value for a Vector or Matrix, which raises. + """ + if (cast_op := _store_cast_op(value.dtype, dtype)) is None: + return value + from .base import BaseType + from .scalar import Scalar + + if output_type(value) is not Scalar: + raise DomainMismatch( + f"cannot assign a value of type {value.dtype} into an object of type {dtype}: " + "assign cannot cast a UDT as it writes, so this would need a converted copy of " + "the whole value. Make it explicitly with .dup(dtype=...) first." + ) + if not isinstance(value, BaseType): + value = value.new(name="") + return value.apply(cast_op).new(name=value.name) + + +# ewise_add, which copies an entry present in only one input without the op. + + +def _is_lifted_op(op): + """Whether ``op`` is a built-in BinaryOp that lifts to UDTs, or a Monoid of one.""" + from .operator.binary import BinaryOp + from .operator.monoid import Monoid + + if isinstance(op, Monoid): + op = op.binaryop + return type(op) is BinaryOp and not op._anonymous and op.name in BUILTIN_UDT_BINARY_OPS + + +def _ewise_add_needs_cast(op, left_dtype, right_dtype): + """Whether ``ewise_add`` with ``op`` would have GraphBLAS cast to or from a UDT. + + eWiseAdd copies an entry present in only one input into the result, cast + to the op's output type, and GraphBLAS checks that cast before it looks at + the data. A UDT casts only to itself, so such a pair is a bare + GrB_DOMAIN_MISMATCH even when every entry has a partner. + """ + out = op.return_type + if not (out._is_udt or left_dtype._is_udt or right_dtype._is_udt): + return False + return left_dtype != out or right_dtype != out + + +def _ewise_add_cast_error(op, left_dtype, right_dtype): + """Return the ``DomainMismatch`` for an ``ewise_add`` that would need a UDT cast.""" + return DomainMismatch( + f"ewise_add cannot use {op.parent!r} on {left_dtype} and {right_dtype}: an entry " + f"present in only one input is copied into the result as {op.return_type}, and " + "GraphBLAS cannot cast a UDT, so this would need converted copies of the operands. " + "Use ewise_union, which takes a default for each side (infix + does, with zeros), or " + "ewise_mult; or convert an operand first with .dup(dtype=...)." + ) + + +def _ewise_add_casts_udt(left, right, op): + """Whether ``left.ewise_add(right, op)`` would need GraphBLAS to cast a UDT.""" + try: + left_dtype, right_dtype = left.dtype, right.dtype + except TypeError: + # An infix operand that has no dtype here (``x & y`` of non-BOOL); + # ewise_add reports it as before. + return False + if not (left_dtype._is_udt or right_dtype._is_udt): + return False + from .operator import get_typed_op + + typed_op = get_typed_op(op, left_dtype, right_dtype, kind="binary") + return _ewise_add_needs_cast(typed_op, left_dtype, right_dtype) + + +def _scalar_ewise_add_converted(left, right, op): + """Return ``left``, ``right`` and ``op`` for an ``ewise_add`` of two Scalars that needs a cast. + + ``op`` is typed on the two, and :func:`_ewise_add_needs_cast` holds: + GraphBLAS would cast an unpaired entry to the result type, which a UDT + cannot be. For a built-in op that lifts to UDTs, both are converted to the + result type first, which changes no value, and the op applies in that one + type: an int8 UDT plus 0.5 is float64, and an empty Scalar still adds as 0. + A Scalar is one element, so the converted copies cost nothing; Vector and + Matrix raise instead, as any other op does here (:func:`_ewise_add_cast_error`). + """ + if not _is_lifted_op(op.parent): + raise _ewise_add_cast_error(op, left.dtype, right.dtype) + if left.dtype != op.return_type: + left = _converted_to(left, op.return_type, name="") + if right.dtype != op.return_type: + right = _converted_to(right, op.return_type, name="") + from .operator import get_typed_op + + return left, right, get_typed_op(op.parent, left.dtype, right.dtype, kind="binary") + + +def _plus_zero(dtype): + """Return a Scalar of ``dtype`` that ``plus`` adds to any value without changing it. + + Its float and complex leaves are ``-0.0``, because ``-0.0 + x`` is ``x`` + for every float, where ``0.0 + -0.0`` is ``0.0``. + """ + from .scalar import Scalar + + np_type = dtype.np_type + if np_type.subdtype is not None: + zero = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) + else: + zero = np.zeros(1, dtype=np_type) + for leaf in _leaf_arrays(zero): + if leaf.dtype.kind in "fc": + leaf[...] = complex(-0.0, -0.0) if leaf.dtype.kind == "c" else -0.0 + return Scalar.from_value(zero if np_type.subdtype is not None else zero[0], dtype, name="") + + +def _converted_to(operand, dtype, **new_kwargs): + """Return ``operand`` converted to UDT ``dtype``, which holds its elements, value for value. + + It adds ``dtype``'s zero on the left with the lifted ``plus``. With the + zero on the left, the result is ``dtype`` itself even when the operand has + the same elements in another layout (a packed and an aligned record), + since a left operand that holds the promoted elements is the result type. + """ + from .. import binary + + rv = operand.apply(binary.plus, left=_plus_zero(dtype)).new(**new_kwargs) + if rv.dtype != dtype: # pragma: no cover (safety) + raise DomainMismatch(f"cannot convert {operand.dtype} to {dtype}") + return rv + + +# isequal, and == without Numba. + + +def _element_shapes_differ(dtype1, dtype2): + """Whether elements of ``dtype1`` and ``dtype2`` differ in shape, for ``isequal``. + + ``==`` on array UDTs broadcasts, as numpy's does, but ``isequal`` asks + whether two values are the same array, as ``np.array_equal`` does: the + shapes must match, apart from leading axes of length 1 (the rule a store + uses), and an element of a built-in dtype has shape ``()``. Records pair + field by field instead (:func:`_check_udt_pair`). + """ + if dtype1 is dtype2 or not (dtype1._is_udt or dtype2._is_udt): + return False + shapes = [] + for np_type in (dtype1.np_type, dtype2.np_type): + if np_type.names is not None: + return False + shapes.append(_strip_leading_ones(np_type.subdtype[1]) if np_type.subdtype else ()) + return shapes[0] != shapes[1] + + +def _literal_shape_differs(dtype, value): + """Whether literal ``value`` differs in shape from an element of array UDT ``dtype``. + + As :func:`_element_shapes_differ`, for ``Scalar.isequal`` with a literal: a + number has shape ``()``, so it equals an ``FP64[1]`` element but not an + ``FP64[3]`` one, as in ``np.array_equal``. ``False`` for a record UDT. + """ + np_type = dtype.np_type + if np_type.subdtype is None or isinstance(value, dict): + return False + try: + shape = np.shape(value) + except ValueError: # ragged; Scalar.from_value reports it + return False + return _strip_leading_ones(shape) != _strip_leading_ones(np_type.subdtype[1]) + + +def _leaves_equal(x, x_type, y, y_type): + """Whether values ``x`` and ``y`` are equal in every leaf, as ``binary.eq`` compares them. + + For ``==`` on a UDT Scalar without Numba. Records pair their fields in + order (the caller checks, by ``_check_udt_pair``, that two records pair as + for ``binary.eq``), and a value that is not a record stands in every + field. Each pair of leaves compares in the type their dtypes promote to, + whatever the values (numpy 1 types a scalar by its value), and broadcasts. + """ + if x_type.names is not None or y_type.names is not None: + x_names = x_type.names or [None] * len(y_type.names) + y_names = y_type.names or [None] * len(x_type.names) + return all( + _leaves_equal( + x if x_name is None else x[x_name], + x_type if x_name is None else x_type.fields[x_name][0], + y if y_name is None else y[y_name], + y_type if y_name is None else y_type.fields[y_name][0], + ) + for x_name, y_name in zip(x_names, y_names, strict=True) + ) + x = np.asarray(x, dtype=x_type.base) + y = np.asarray(y, dtype=y_type.base) + common = np.result_type(x.dtype, y.dtype) + return bool(np.all(x.astype(common) == y.astype(common))) + + +def _udt_eq_without_numba(scalar, other): + """Return ``scalar == other`` as a bool, in numpy, for a UDT when Numba is not installed. + + ``binary.eq`` on a UDT needs Numba. ``scalar`` is nonempty, as is ``other`` + when it is a Scalar, and one of them is a UDT. A literal ``other`` is typed + as ``binary.eq`` types it, and the values compare leaf by leaf as + ``binary.eq`` compares them (:func:`_leaves_equal`). + """ + from .. import binary + from .literals import _literal_operand + from .scalar import Scalar + + if type(other) is not Scalar: + value, dtype = _literal_operand(scalar.dtype, other, binary.eq) + other = Scalar.from_value(value, dtype, is_cscalar=None, name="s_eq_other") + # Two UDTs pair as for binary.eq, or raise its KeyError (records with other + # field names, array shapes that do not broadcast, a record and an array). + info = _get_udt_info(scalar.dtype), _get_udt_info(other.dtype) + _check_udt_pair("eq", scalar.dtype, other.dtype, *info) + return _leaves_equal(scalar.value, scalar.dtype.np_type, other.value, other.dtype.np_type) + + +# An element's value as bytes and back, for Scalar.value. + + +def _element_bytes(dtype, value): + """Return ``value`` as the bytes of an element of UDT ``dtype``, padding zeroed. + + ``value`` is anything numpy assigns to one element: a tuple or dict for a + record, an array or sequence for an array UDT, or a number for every leaf. + """ + np_type = dtype.np_type + # ``_udt_bytes`` zeros a record's alignment padding, which numpy may leave + # as junk however it writes the record. + if np_type.subdtype is None: + arr = np.zeros(1, dtype=np_type) + if isinstance(value, dict) and np_type.names is not None: + value = _dict_to_record(np_type, value) + else: + arr = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) + value = _view_if_same_layout(value, arr.dtype) + if _has_narrow_float_leaf(np_type) and np.geterr()["over"] == "warn": + # A float too large for a float32 leaf is infinity, and numpy warns; + # warn at the caller's line instead, as a built-in FP32 does, so the + # default filter shows each place it happens. + try: + with np.errstate(over="raise"): + arr[:] = value + except FloatingPointError: + utils._warn_from_caller( + f"overflow encountered in cast: {value!r} as {dtype}", RuntimeWarning + ) + with np.errstate(over="ignore"): + arr[:] = value + else: + arr[:] = value + return _udt_bytes(np_type, arr) + + +def _element_value(dtype, data): + """Return the element of UDT ``dtype`` that cffi ``data`` holds: a record or an array.""" + np_type = dtype.np_type + rv = np.array(ffi.buffer(data[0 : np_type.itemsize])) + if np_type.subdtype is None: + return rv.view(np_type)[0] + base, shape = np_type.subdtype + return rv.view(base).reshape(shape) + + +def _dict_to_record(np_type, d): + """Converts e.g. ``{"x": 1, "y": 2.3}`` to ``(1, 2.3)``.""" + rv = [] + for name, (dtype, _) in np_type.fields.items(): + val = d[name] + if dtype.names is not None and isinstance(val, dict): + rv.append(_dict_to_record(dtype, val)) + else: + rv.append(val) + return tuple(rv) + + +def _leaf_arrays(arr): + """Yield the array of each leaf of a structured array, in field order.""" + if arr.dtype.names is None: + yield arr + else: + for name in arr.dtype.names: + yield from _leaf_arrays(arr[name]) + + +@lru_cache +def _has_narrow_float_leaf(np_type): + """Whether ``np_type`` has a float leaf narrower than float64, which a float may overflow.""" + if np_type.subdtype is not None: + return _has_narrow_float_leaf(np_type.subdtype[0]) + if np_type.names is not None: + return any(_has_narrow_float_leaf(np_type.fields[name][0]) for name in np_type.names) + kind, size = np_type.kind, np_type.itemsize + return (kind == "f" and size < 8) or (kind == "c" and size < 16) + + +@lru_cache +def _padding_ranges(np_type): + """The ``(start, stop)`` byte ranges of ``np_type`` that are alignment padding. + + Walks records, nested records, and array dtypes, and returns the gaps + between the field data. + """ + data = [] + _collect_data_ranges(np_type, 0, data) + data.sort() + ranges = [] + pos = 0 + for start, stop in data: + if start > pos: + ranges.append((pos, start)) + pos = max(pos, stop) + if pos < np_type.itemsize: + ranges.append((pos, np_type.itemsize)) + return tuple(ranges) + + +def _collect_data_ranges(np_type, offset, out): + """Append the ``(start, stop)`` byte ranges of ``np_type`` that hold field data.""" + if np_type.subdtype is not None: + base = np_type.subdtype[0] + for start in range(offset, offset + np_type.itemsize, base.itemsize): + _collect_data_ranges(base, start, out) + elif np_type.names is not None: + for field_type, field_offset in (np_type.fields[name][:2] for name in np_type.names): + _collect_data_ranges(field_type, offset + field_offset, out) + else: + out.append((offset, offset + np_type.itemsize)) + + +def _udt_bytes(np_type, arr): + """The bytes of a UDT buffer, with any alignment padding zeroed. + + ``tobytes()`` flattens any-rank numpy buffers (in particular, multi-dim + array UDTs like ``FP64[2, 3]``) to a 1-D byte buffer that cffi can copy + into GrB_Scalar storage. + + numpy promises nothing about the padding bytes it leaves behind, and which + way it goes has changed across releases: numpy 2.5 made ``arr[:] = record`` + a whole-item copy out of an uninitialized temporary, so even a buffer from + ``np.zeros`` comes back with junk between the fields. A caller's own array + can carry junk padding in too. Since the whole record is copied into the + GrB_Scalar and read back out of ``Scalar.value``, anything that compares or + hashes those raw bytes would otherwise see an equal-valued scalar as + unequal. + """ + raw = arr.tobytes() + ranges = _padding_ranges(np_type) + if not ranges: + return raw + buf = bytearray(raw) + for start, stop in ranges: + buf[start:stop] = bytes(stop - start) + return bytes(buf) diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index c8aaeb1e6..0a6087a18 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -9,6 +9,7 @@ from .base import BaseExpression, BaseType, _check_mask, call from .descriptor import lookup as descriptor_lookup from .expr import _ALL_INDICES, AmbiguousAssignOrExtract, IndexerResolver, InfixExprBase, Updater +from .literals import _literal_operand, _literal_store_dtype, _weak_builtin_literal_dtype from .mask import Mask, StructuralMask, ValueMask from .operator import ( UNKNOWN_OPCLASS, @@ -26,16 +27,14 @@ ScalarExpression, ScalarIndexExpr, _as_scalar, + _scalar_index, + _union_default, +) +from .udt import ( _cast_for_store, - _check_literal_fits, - _check_scalar_fits, _element_shapes_differ, _ewise_add_cast_error, _ewise_add_needs_cast, - _literal_operand, - _literal_store_dtype, - _scalar_index, - _weak_builtin_literal_dtype, ) from .utils import ( _CArray, @@ -1234,52 +1233,8 @@ def _ewise_union(self, other, op, left_default, right_default, is_infix=False): ) temp_op = get_typed_op(op, self.dtype, other.dtype, kind="binary") - left_dtype = temp_op.type - dtype = left_dtype if left_dtype._is_udt else None - if type(left_default) is not Scalar: - if dtype is not None: - _check_literal_fits(dtype, left_default) - else: - left_default, dtype = _literal_operand(left_dtype, left_default, op) - try: - left = Scalar.from_value( - left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar - ) - except TypeError: - left = self._expect_type( - left_default, - Scalar, - within=method_name, - keyword_name="left_default", - extra_message="Literal scalars also accepted.", - op=op, - ) - else: - _check_scalar_fits(dtype, left_default) - left = _as_scalar(left_default, dtype, is_cscalar=False) # pragma: is_grbscalar - right_dtype = temp_op.type2 - dtype = right_dtype if right_dtype._is_udt else None - if type(right_default) is not Scalar: - if dtype is not None: - _check_literal_fits(dtype, right_default) - else: - right_default, dtype = _literal_operand(right_dtype, right_default, op) - try: - right = Scalar.from_value( - right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar - ) - except TypeError: - right = self._expect_type( - right_default, - Scalar, - within=method_name, - keyword_name="right_default", - extra_message="Literal scalars also accepted.", - op=op, - ) - else: - _check_scalar_fits(dtype, right_default) - right = _as_scalar(right_default, dtype, is_cscalar=False) # pragma: is_grbscalar + left = _union_default(self, left_default, temp_op.type, op, "left_default") + right = _union_default(self, right_default, temp_op.type2, op, "right_default") if is_infix: op1 = _get_typed_op_from_exprs(op, self, right, kind="binary") diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index c1a8cd7ea..b4c353767 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -10,7 +10,7 @@ from graphblas import core, dtypes from graphblas.core import _supports_udfs as supports_udfs from graphblas.core import lib -from graphblas.core.operator.udt_utils import _has_jit_set +from graphblas.core.operator.udt_codegen import _has_jit_set from graphblas.core.utils import _NP2, values_to_numpy_buffer from graphblas.dtypes import lookup_dtype diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 4ed64a54c..4f69baf32 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2240,7 +2240,7 @@ def test_udt_unary_ainv_abs_array(array_udt): def test_udt_jit_typedef(): """Registering a UDT sets GxB_JIT_C_NAME and GxB_JIT_C_DEFINITION.""" from graphblas.core import ffi, lib - from graphblas.core.operator.udt_utils import _has_jit_set + from graphblas.core.operator.udt_codegen import _has_jit_set if not _has_jit_set: pytest.skip("JIT not available") @@ -2279,7 +2279,7 @@ def test_udt_jit_typedef(): def test_udt_jit_op_definitions(): """Auto-compiled UDT ops carry the JIT C name and source.""" from graphblas.core import ffi, lib - from graphblas.core.operator.udt_utils import _has_jit_set + from graphblas.core.operator.udt_codegen import _has_jit_set if not _has_jit_set: pytest.skip("JIT not available") @@ -2434,7 +2434,7 @@ def test_udt_array_wrapper_stays_within_element(shape): import numba - from graphblas.core.operator.base import _get_udt_wrapper + from graphblas.core.operator.udf import _get_udt_wrapper udt = dtypes.register_anonymous(np.dtype((np.float64, shape)), f"_ArrWrapPin{len(shape)}D") @@ -2813,7 +2813,7 @@ def test_udt_broadcast_matches_numba_slice_assign(): """ import numba - from graphblas.core.operator.base import _fits_by_broadcast + from graphblas.core.operator.udf import _fits_by_broadcast @numba.njit def _assign(z, src): # pragma: no cover (numba) @@ -2862,16 +2862,16 @@ def test_udt_udf_shape_check_runs_udf_only_for_built_arrays(monkeypatch): UDF builds, typed as a plain ``Array``, needs the probe; so does a list, whose type has no length (``test_udt_record_leaf_sequence_return``). """ - from graphblas.core.operator import base as _base + from graphblas.core.operator import udf as _udf probed = [] - real_probe = _base._run_udf_probe + real_probe = _udf._run_udf_probe def spy(numba_func, arg_dtypes): probed.append(numba_func.py_func.__name__) return real_probe(numba_func, arg_dtypes) - monkeypatch.setattr(_base, "_run_udf_probe", spy) + monkeypatch.setattr(_udf, "_run_udf_probe", spy) udt = dtypes.register_anonymous(np.dtype((np.float32, (18,))), "_ProbeWhen18") rec = dtypes.register_anonymous( np.dtype( @@ -3057,7 +3057,7 @@ def test_udt_array_udf_known_extents_never_match_by_rank(monkeypatch): """ import numba - from graphblas.core.operator import base as _base + from graphblas.core.operator import udf as _udf # FP32 keeps this UDT under that 128-byte cap, so it registers everywhere. udt21 = dtypes.register_anonymous(np.dtype((np.float32, (21,))), "_RankGuard21") @@ -3066,16 +3066,16 @@ def test_udt_array_udf_known_extents_never_match_by_rank(monkeypatch): assert ret.dtype == udt21.numba_type.dtype assert ret.ndim == udt21.numba_type.ndim - real_lookup = _base.lookup_dtype + real_lookup = _udf.lookup_dtype def refuse_nested(key, value=None): if isinstance(key, numba.core.types.NestedArray): raise ValueError("simulated: SuiteSparse refused to register this UDT") return real_lookup(key, value) - monkeypatch.setattr(_base, "lookup_dtype", refuse_nested) + monkeypatch.setattr(_udf, "lookup_dtype", refuse_nested) with pytest.raises(UdfParseError, match="matches no input array UDT"): - _base._resolve_udt_return_type(ret, udt21) + _udf._resolve_udt_return_type(ret, udt21) @pytest.mark.skipif("not supports_udfs") @@ -5332,7 +5332,7 @@ def test_udt_literal_rules_edge_cases(): from decimal import Decimal from fractions import Fraction - from graphblas.core.operator.udt_utils import _weak_literal_udt + from graphblas.core.literals import _weak_literal_udt i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_EdgeI8") f4s = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_EdgeF4") @@ -5610,7 +5610,7 @@ def test_udt_jit_kernel_only_for_a_same_type_result(): 64 entries keep SuiteSparse off its iso and short-vector shortcuts, so a JIT kernel would run. """ - from graphblas.core.operator.udt_utils import _has_jit_set + from graphblas.core.operator.udt_codegen import _has_jit_set ints = dtypes.register_anonymous(np.dtype((np.int32, (11,))), "_JitDivI11") N = 64 @@ -6064,7 +6064,7 @@ def test_compile_codegen_helper(): """ import linecache - from graphblas.core.operator.udt_utils import _compile_codegen + from graphblas.core.operator.udt_codegen import _compile_codegen fn = _compile_codegen( "def _op(x, y):\n return x + y\n", diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 097fc6866..678f852fa 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -25,7 +25,7 @@ ) from graphblas.core import _supports_udfs as supports_udfs from graphblas.core.operator.indexbinary import _has_idxbinop -from graphblas.core.operator.udt_utils import _has_jit_set +from graphblas.core.operator.udt_codegen import _has_jit_set from graphblas.core.ss import _IS_SSGB7 from graphblas.core.ss.jit_config import _enable_jit_for_udt diff --git a/pyproject.toml b/pyproject.toml index 9bd988a10..9c9701dde 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -388,7 +388,7 @@ ignore = [ "graphblas/core/operator/__init__.py" = ["A005"] "graphblas/io/__init__.py" = ["A005"] # shadows a standard-library module "graphblas/core/operator/base.py" = ["S102"] # exec is used for UDF -"graphblas/core/operator/udt_utils.py" = ["S102"] # exec is used for UDT op codegen +"graphblas/core/operator/udt_codegen.py" = ["S102"] # exec is used for UDT op codegen "graphblas/monoid/numpy.py" = ["PLW0108"] # lambda is needed for numba.njit "graphblas/core/ss/matrix.py" = [ "NPY002", # numba doesn't support rng generator yet From 54d3b778042c5a6e6af1aa7f0796612743877104 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 00:54:04 -0500 Subject: [PATCH 12/54] Pass CODECOV_TOKEN to the coverage upload Codecov takes tokenless uploads only from forks. Every upload from this PR's CI runs was refused with "Token required because branch is protected", so the PR had no coverage report, and the PRs stacked on it had no base to compare against. The repository now has the secret. --- .github/workflows/test_and_build.yml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.github/workflows/test_and_build.yml b/.github/workflows/test_and_build.yml index 574bcd3d8..8dec76f13 100644 --- a/.github/workflows/test_and_build.yml +++ b/.github/workflows/test_and_build.yml @@ -358,6 +358,11 @@ jobs: coverage report --show-missing - name: codecov uses: codecov/codecov-action@v5 + with: + # Codecov takes tokenless uploads only from forks. A PR from a branch of this + # repository, or a push to main, is refused with "Token required because + # branch is protected". + token: ${{ secrets.CODECOV_TOKEN }} - name: Notebooks Execution check if: matrix.slowtask == 'notebooks' run: | From 67a70331d17970652c559456d76fb59b491930c9 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 01:33:59 -0500 Subject: [PATCH 13/54] Cover how core/literals.py types literal operands Tests for the literal-typing paths CI never ran: - Beside built-in dtypes: a comparison op given as a string with an int out of range (ewise_union with "<"), an op string that names no BinaryOp (the op reports it), an int beyond INT64 beside a bool (OverflowError, as numpy's bools & 2**70), a user's BinaryOp with an int out of range, and a complex literal on a backend without complex types (vanilla only). - Beside UDTs: an infinite or NaN Decimal, literals not laid out like an element (ragged, too long or too short for a field, not broadcastable) or not numbers, which numpy then rejects, a number beside a bytes field, a nested record that is not the last field, and the conversions that lose a number (an imaginary part, a float32 overflow, a fraction in a record's array field) or keep a bytes field as it is. - Scalar.isequal with a ragged literal on an array UDT (ValueError). --- graphblas/tests/test_op.py | 59 +++++++++++++++++++++++++++++++++- graphblas/tests/test_scalar.py | 3 ++ graphblas/tests/test_vector.py | 18 ++++++++++- 3 files changed, 78 insertions(+), 2 deletions(-) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 4f69baf32..30c6cb9f3 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4524,12 +4524,23 @@ def test_udt_literal_converts_only_when_exact(): rec = dtypes.register_anonymous( np.dtype([("ex_i", np.int16), ("ex_f", np.float32)], align=True), "_ExactRec" ) + floats = dtypes.register_anonymous(np.dtype((np.float32, (2,))), "_ExactF32") + nested = dtypes.register_anonymous( + np.dtype([("ex_v", np.int8, (2,)), ("ex_w", np.int8)]), "_ExactNested" + ) + text = dtypes.register_anonymous(np.dtype([("ex_s", "S4"), ("ex_n", np.int8)]), "_ExactText") v = Vector(arr, size=3) v[0] = [1, 2, 3] w = Vector(arr, size=3) w[1] = [4, 5, 6] r = Vector(rec, size=1) r[0] = (1, 1.5) + f = Vector(floats, size=1) + f[0] = [1, 2] + n = Vector(nested, size=1) + n[0] = ([1, 2], 3) + t = Vector(text, size=1) + t[0] = (b"ab", 1) udf = BinaryOp.register_anonymous(lambda x, y: x + y, "_exact_plus", is_udt=True) second = BinaryOp.register_anonymous(lambda x, y: y, "_exact_second", is_udt=True) # Literals of the UDT's type convert, a float rounding into a float32 field. @@ -4557,20 +4568,32 @@ def test_udt_literal_converts_only_when_exact(): lambda: gb.Scalar.from_value([1, 2, 3], arr).ewise_union(gb.Scalar(arr), udf, 0, 0.5), "_ExactArr", ), + # Converting would lose a number: an imaginary part, a finite float that + # float32 makes infinite, or a fraction in an array field of a record. + (lambda: f.apply(udf, right=1 + 2j), "_ExactF32"), + (lambda: f.apply(udf, right=np.complex128(1 + 2j)), "_ExactF32"), + (lambda: f.apply(udf, right=np.float64(1e300)), "_ExactF32: it is typed as FP64,"), + (lambda: n.apply(second, right=([1, 2.5], 3)), "_ExactNested"), ]: with pytest.raises(ValueError, match=re.escape(f"does not fit {typed_as}")): expr() # One that does not fit by type but converts exactly still converts, as it - # did on main, with a DeprecationWarning. + # did on main, with a DeprecationWarning. A field that is not a number, such + # as bytes, is converted as it is. + text_literal = np.array((b"cd", 2), dtype=[("ex_s", "S4"), ("ex_n", np.int64)])[()] for expr, typed_as in [ (lambda: v.apply(udf, right=[2.0, 2.0, 2.0]), "_ExactArr"), (lambda: v.apply(udf, right=np.ones(3, dtype=np.int64)), "_ExactArr"), (lambda: v.apply(udf, right=np.int64(1)), "_ExactArr: it is typed as INT64,"), (lambda: v.apply(udf, right=(np.int8(1), 2, 3)), "_ExactArr"), (lambda: r.apply(second, right=np.int32(2)), "_ExactRec: it is typed as INT32,"), + (lambda: n.apply(second, right=([1, 2.0], 3)), "_ExactNested"), + (lambda: t.apply(second, right=text_literal), "_ExactText"), ]: with pytest.warns(DeprecationWarning, match=re.escape(f"does not fit {typed_as}")): expr() + with pytest.warns(DeprecationWarning, match="does not fit _ExactText"): + assert t.apply(second, right=text_literal).new()[0].new().value.tolist() == (b"cd", 2) with pytest.raises(OverflowError, match="300 out of bounds for int8"): v.apply(udf, right=(1, 2, 300)) # Defaults that fit still work, and fill in for the missing side. @@ -4668,6 +4691,12 @@ def test_udt_sequence_and_array_literals_type_like_numbers(): dtype=[("sq_a", np.int8), ("sq_in", [("sq_b", np.float64), ("sq_c", np.int8)])], ) assert (n + strong).new()[0].new().value.tolist() == (2, (2.5, 4.5)) + first = dtypes.register_anonymous( + np.dtype([("sq_fin", [("sq_fb", np.int8), ("sq_fc", np.float32)]), ("sq_fa", np.int16)]) + ) + m = Vector(first, size=1) + m[0] = ((2, 3.5), 1) + assert binary.plus(m, ((0.5, 1), 1)).new()[0].new().value.tolist() == ((2.5, 4.5), 2) # eq and ne take the literal as it is, like numpy, instead of raising. assert (v == (0.5, 4, 5)).new().isequal(Vector.from_coo([0, 1], [False, False])) assert (v == (2, 4, 5)).new().isequal(Vector.from_coo([0, 1], [True, False])) @@ -5351,9 +5380,32 @@ def test_udt_literal_rules_edge_cases(): result = (v + value).new() assert result.dtype == i8s assert result[0].new().value.tolist() == [3, 4, 5] + # An infinite or NaN Decimal is the float it stands for. + for value in [Decimal("Infinity"), Decimal("NaN")]: + result = (v + value).new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + np.testing.assert_array_equal(result[0].new().value, [float(value)] * 3) udf = BinaryOp.register_anonymous(lambda x, y: x + y, "_edge_plus", is_udt=True) with pytest.raises(ValueError, match=r"Fraction\(1, 2\) does not fit _EdgeI8"): v.apply(udf, right=Fraction(1, 2)) + # A literal not laid out like an element (ragged, too long or too short for + # a field, not broadcastable), or not a number, has no type of its own, so + # it is converted into the UDT, and numpy says why it cannot be. + rec = dtypes.register_anonymous( + np.dtype([("edge_v", np.float64, (3,)), ("edge_n", np.int8)]), "_EdgeRec" + ) + r = Vector(rec, size=1) + r[0] = ([1, 2, 3], 4) + for expr, reason in [ + (lambda: binary.plus(v, [1, [2, 3]]), "with a sequence"), + (lambda: v + np.array([1, 2]), "could not broadcast"), + (lambda: v + "abc", "invalid literal"), + (lambda: v.apply(udf, right="abc"), "invalid literal"), + (lambda: binary.plus(r, (1, 2, 3)), "tuple of length 3"), + (lambda: binary.plus(r, ((1, 2), 3)), "could not broadcast"), + ]: + with pytest.raises(ValueError, match=reason): + expr() c = Vector(c8s, size=1) c[0] = [1, 2, 3] assert c.apply(binary.plus, right=(1 + 2j, 0, 0)).new()[0].new().value[0] == 2 + 2j @@ -5538,6 +5590,8 @@ def vec(dtype, value): assert not u.isequal(gb.Scalar.from_value([1.0], f1)) assert u.isequal((1, 1, 1)) assert u.isequal(np.ones((1, 3), dtype=np.int8)) + with pytest.raises(ValueError, match="with a sequence"): # ragged: it has no shape + u.isequal([1, [1, 1]]) assert gb.Scalar.from_value([5], f1).isequal(5) # A literal is compared as given, not converted into the UDT first. assert zeros != 0.5 @@ -5649,6 +5703,9 @@ def test_udt_eq_compile_failure_is_a_udfparseerror(): binary.eq(v & w).new() with pytest.raises(KeyError, match=r"elements of \|S4 and float64 have no common type"): binary.plus(v & w).new() + # A number beside the bytes field is an int64 there, which pairs no better. + with pytest.raises(KeyError, match=r"elements of \|S4 and int64 have no common type"): + v + 1 @pytest.mark.skipif("not supports_udfs") diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index 6c14e70df..81223942c 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -675,6 +675,9 @@ def test_ewise_mult_add(s): s.ewise_add(v) with pytest.raises(TypeError, match="Literal scalars also"): s.ewise_mult(v) + # A literal is typed beside a string op that names no BinaryOp, so the op reports it. + with pytest.raises(ValueError, match="Unknown binary string: 'plus_times'"): + s.ewise_mult(1, "plus_times") def test_select(s): diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 34d1ae353..701d2cc31 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -1533,12 +1533,23 @@ def test_python_number_is_weak_beside_builtin_dtypes(): assert (f32 * 0.5).dtype == dtypes.FP32 if dtypes._supports_complex: assert (f32 + 1j).dtype == dtypes.FC32 + else: + # Without complex types, a complex number keeps its own type, which plus does not take. + with pytest.raises(KeyError, match="plus does not work with"): + f32 + 1j s8 = Scalar.from_value(100, dtype=dtypes.INT8) assert (s8 + 1).new().dtype == (1 + s8).new().dtype == dtypes.INT8 assert (Scalar.from_value(0.5, dtype=dtypes.FP32) * 0.5).new().dtype == dtypes.FP32 A = Matrix.from_coo([0], [0], [100], dtype=dtypes.INT8) assert (A * 2).dtype == dtypes.INT8 - exprs = [lambda: i8 + 300, lambda: u8 + -1, lambda: s8 * 300] + # Beside a bool, an int is INT64, and must fit it, as in numpy's ``bools & 2**70``. + exprs = [ + lambda: i8 + 300, + lambda: u8 + -1, + lambda: s8 * 300, + lambda: bools + 2**70, + lambda: bools.apply(binary.band, right=2**70), + ] if supports_udfs: # integer floordiv is a Numba op exprs.append(lambda: i8 // 300) for expr in exprs: @@ -1576,6 +1587,7 @@ def test_python_number_is_weak_beside_builtin_dtypes(): w8 = Vector.from_coo([1], [1], dtype=dtypes.INT8, size=2) assert i8.ewise_union(w8, binary.plus, 0, 0).dtype == dtypes.INT8 assert i8.ewise_union(w8, binary.plus, 0.5, 0).dtype == dtypes.FP64 + assert i8.ewise_union(w8, "<", 0, 300).new().to_coo()[1].tolist() == [True, True] sparse = Vector.from_coo([0], [5], dtype=dtypes.INT8, size=2) assert sparse.to_dense(fill_value=0).dtype == np.int8 assert sparse.to_dense(fill_value=300).tolist() == [5, 300] # widened, not lost @@ -1590,6 +1602,10 @@ def test_python_number_is_weak_beside_builtin_dtypes(): if supports_udfs: above = indexunary.register_anonymous(lambda x, i, j, thunk: x > thunk) assert i8.apply(above, 300).new().to_coo()[1].tolist() == [False, False] + # A user's BinaryOp computes in the collection's type, as plus does. + plus = binary.register_anonymous(lambda x, y: x + y) + with pytest.raises(OverflowError, match="out of bounds"): + i8.apply(plus, right=300) # Typed scalars promote as numpy 2 promotes them, on numpy 1 too. assert dtypes.unify(dtypes.INT8, dtypes.INT64, is_right_scalar=True) == dtypes.INT64 From 70d9bac0c72ccbfda1e5487f13f74e9069fff668 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 01:34:09 -0500 Subject: [PATCH 14/54] Cover select's rewrite of value thunks of another type Tests for the branches of operator/select.py's _value_select_in_own_type that CI never ran: a complex thunk beside integers and floats (only valueeq and valuene, which have an answer), a NaN thunk of another float type, an infinite thunk beside integers, and an empty thunk Scalar, which GraphBLAS rejects. --- graphblas/tests/test_vector.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 701d2cc31..81e36d851 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -1683,6 +1683,9 @@ def test_select_value_thunk_of_another_type(): (i8, select.valuegt, -300, [0, 1, 2, 3]), (i8, select.valuegt, np.nan, []), (i8, select.valuene, np.nan, [0, 1, 2, 3]), + (i8, select.valuelt, np.inf, [0, 1, 2, 3]), + (i8, select.valuege, np.inf, []), + (i8, select.valuele, -np.inf, []), (i8, indexunary.valueeq, np.float64(2.0), [2]), (i8, select.valueeq, Scalar.from_value(1.0), [1]), (i8, select.valuelt, Scalar.from_value(300, dtype=dtypes.INT64), [0, 1, 2, 3]), @@ -1694,10 +1697,23 @@ def test_select_value_thunk_of_another_type(): (f32, select.valueeq, np.float64(0.1), []), (f32, select.valuegt, np.float64(0.1), [0, 1]), (f32, select.valuene, np.float64(0.1), [0, 1, 2]), + (f32, select.valuelt, np.float64(np.nan), []), + (f32, select.valuene, np.float64(np.nan), [0, 1, 2]), # A Python float beside FP32 is FP32 (NEP 50), as numpy compares it. (f32, select.valueeq, 0.1, [0]), ]: assert x.select(op, thunk).new().to_coo()[0].tolist() == kept, (x.dtype, op, thunk) + if dtypes._supports_complex: + # A real value equals a complex one only when its imaginary part is 0. + for x, op, thunk, kept in [ + (i8, select.valueeq, 1 + 0j, [1]), + (i8, select.valueeq, 1 + 2j, []), + (i8, select.valuene, 1 + 2j, [0, 1, 2, 3]), + (f32, select.valueeq, np.complex64(1), [1]), + ]: + assert x.select(op, thunk).new().to_coo()[0].tolist() == kept, (x.dtype, op, thunk) + with pytest.raises(EmptyObject): + i8.select(select.valueeq, Scalar(dtypes.INT64)).new() A = Matrix.from_coo([0, 0, 1], [0, 1, 1], [1, 100, -5], dtype=dtypes.INT8) assert A.select(select.valueeq, 1.0).new().to_coo()[2].tolist() == [1] assert A.select(select.valuelt, 300).new().nvals == 3 From 9513655b2de7885e4f3d6e70394ec731c0b2732d Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 01:34:18 -0500 Subject: [PATCH 15/54] Cover UDT store casts, ewise_add casts and fused stores Tests for the paths in core/udt.py, Matrix and udt_codegen that CI never ran: - store casts between records whose field shapes differ or whose bytes fields differ in width (DomainMismatch saying why), a Scalar expression assigned to many elements of another UDT type, one element extracted from a Matrix into a Scalar of another type, and Matrix.power with a semiring whose result is another type; - without Numba no UDT cast compiles, so every such store raises (monkeypatched, so it runs everywhere); - Matrix.ewise_add with operands of another type than the result's (DomainMismatch, and DimensionMismatch first for other shapes), and Scalar.ewise_add with a Monoid and with a non-lifted op; - a fused store with a built-in Scalar operand, and fused stores that cannot be made (first, a user's op); - (x | y) + z of non-BOOL vectors, which reports the infix error. --- graphblas/tests/test_infix.py | 2 ++ graphblas/tests/test_op.py | 57 ++++++++++++++++++++++++++++++++++- 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/graphblas/tests/test_infix.py b/graphblas/tests/test_infix.py index a1a59fd69..94d2555ae 100644 --- a/graphblas/tests/test_infix.py +++ b/graphblas/tests/test_infix.py @@ -497,6 +497,8 @@ def test_multi_infix_vector(): v1.ewise_mult(v2 | v3) with pytest.raises(TypeError, match="Automatic computation"): (v1 | v2).ewise_mult(v3) + with pytest.raises(TypeError, match="Automatic computation"): + (v1 | v2) + v3 @autocompute diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 30c6cb9f3..f946f4dd8 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4448,9 +4448,23 @@ def test_udt_ewise_add_needs_operands_of_the_result_type(): A[0, 0] = [1, 2, 3] B = Matrix(fp32s, nrows=1, ncols=2) B[0, 1] = [0.5, 0.5, 0.5] - for expr in [lambda: v.ewise_add(w), lambda: w.ewise_add(v), lambda: binary.plus(v | w)]: + for expr in [ + lambda: v.ewise_add(w), + lambda: w.ewise_add(v), + lambda: binary.plus(v | w), + lambda: A.ewise_add(B), + lambda: A.ewise_add(w[:2].new()), + ]: with pytest.raises(DomainMismatch, match="ewise_add cannot use .* converted copies"): expr() + with pytest.raises(DimensionMismatch, match="Matrix shapes must match"): + A.ewise_add(Matrix(fp32s, nrows=2, ncols=2)) + # A Scalar converts only for an op that lifts to UDTs, so ``first`` raises too. + s2 = gb.Scalar.from_value([0.5, 0.5, 0.5], dtype=fp32s) + result = s.ewise_union(s2, monoid.plus, 0, 0).new() # a Monoid lifts as its BinaryOp + np.testing.assert_array_equal(result.value, [1.5, 2.5, 3.5]) + with pytest.raises(DomainMismatch, match="ewise_add cannot use binary.first"): + s.ewise_add(s2, binary.first) # Infix ``+`` takes the union with each operand's zero for a missing value # instead, as ``-`` takes 0. for left, right in [(v, w), (w, v)]: @@ -4913,6 +4927,11 @@ def test_udt_store_casts_each_element_as_builtin_types_do(): C << A.mxm(A, semiring.plus_times) with pytest.raises(DomainMismatch, match="would need a temporary copy"): w << v.ewise_add(v, binary.plus) + # Only a built-in op that lifts to UDTs can, so not first or a user's op. + udf = BinaryOp.register_anonymous(lambda x, y: x, "_store_cast_first", is_udt=True) + for not_lifted in [binary.first, udf]: + with pytest.raises(DomainMismatch, match="would need a temporary copy"): + w << v.ewise_mult(v, not_lifted) with pytest.raises(DomainMismatch, match="extraction cannot"): w << v[[0, 1, 2]] # The computed result, stored as an existing object, casts in one pass. @@ -4922,6 +4941,12 @@ def test_udt_store_casts_each_element_as_builtin_types_do(): C << A.power(3) C << A.power(1) assert C.to_coo()[2].tolist() == [first] + # A semiring whose result is another type (int truediv is float) stores that + # type as it is. + P = Matrix.from_coo([0], [0], [[[1, -1], [2, 4], [-8, 3]]], dtype=i8s, nrows=1, ncols=1) + Q = Matrix(f64s, nrows=1, ncols=1) + Q << P.power(2, semiring.plus_truediv) + assert Q.isequal(semiring.plus_truediv(P @ P).new()) # Leading axes of length 1 do not change the layout. lead = dtypes.register_anonymous(np.dtype((np.float32, (1, 3, 2))), "_StoreCastLead") u = Vector(lead, size=3) @@ -4934,6 +4959,17 @@ def test_udt_store_casts_each_element_as_builtin_types_do(): renamed = np.dtype([(f"sr{i}", dst) for i, (_src, dst) in enumerate(pairs)]) with pytest.raises(DomainMismatch, match="same field names"): Vector(dtypes.register_anonymous(renamed), size=7) << x + # A record field keeps its shape, and one that is not a number casts only to + # its own type, as for array UDTs and built-in types. + narrow = dtypes.register_anonymous(np.dtype([("sf_v", np.int8, (2,)), ("sf_s", "S4")])) + wide = dtypes.register_anonymous(np.dtype([("sf_v", np.int8, (3,)), ("sf_s", "S4")])) + longer = dtypes.register_anonymous(np.dtype([("sf_v", np.int8, (2,)), ("sf_s", "S8")])) + y = Vector(narrow, size=1) + y[0] = ([1, 2], b"ab") + with pytest.raises(DomainMismatch, match="same shape for each field"): + Vector(wide, size=1) << y + with pytest.raises(DomainMismatch, match=r"elements of \|S4 do not cast to \|S8"): + Vector(longer, size=1) << y raw = Vector(dtypes.register_anonymous(np.dtype("S48"), "_StoreCastBytes"), size=3) with pytest.raises(DomainMismatch, match="only record and array UDTs cast"): raw << v @@ -4946,6 +4982,18 @@ def test_udt_store_casts_each_element_as_builtin_types_do(): expr() +def test_udt_store_casts_need_numba(monkeypatch): + """Without Numba no UDT cast compiles, so every store across types with a UDT raises.""" + i8s = dtypes.register_anonymous(np.dtype((np.int8, (2,))), "_NoNumbaCastI8") + f64s = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_NoNumbaCastF64") + v = Vector(i8s, size=1) + v[0] = [1, 2] + monkeypatch.setattr(gb.core.udt, "_has_numba", False) + for out in [Vector(f64s, size=1), Vector(dtypes.FP64, size=1)]: + with pytest.raises(DomainMismatch, match=r"_NoNumbaCastI8 in an object .*: UDTs do not"): + out << v + + @pytest.mark.skipif("not supports_udfs") def test_udt_extracted_element_of_another_type_is_cast(): """One element extracted into another UDT type is cast, not copied as raw bytes. @@ -4964,6 +5012,8 @@ def test_udt_extracted_element_of_another_type_is_cast(): s = gb.Scalar(i16s) s << v[0] assert s.value.tolist() == [1, -300] + s << A[1, 2] + assert s.value.tolist() == [2, -32768] # saturates assert v[0].new(dtype=i16s).value.tolist() == [1, -300] assert A.T[2, 1].new(dtype=i16s).value.tolist() == [2, -32768] # saturates w = Vector(i16s, size=3) @@ -5023,6 +5073,7 @@ def test_udt_fused_store_matches_compute_then_cast(): (lambda: f * f, i16s), (lambda: f.ewise_mult(f, monoid.plus), i16s), (lambda: r * 1.5, rec_i), + (lambda: r.apply(binary.minus, left=gb.Scalar.from_value(0.5)), rec_i), (lambda: r.apply(unary.abs), rec_i), ] mask = Vector.from_coo([0, 2], [True, True], size=4) @@ -5433,6 +5484,10 @@ def test_udt_store_casts_reach_every_store_path(): s = gb.Scalar(f8s) s << v.reduce(aggregator) np.testing.assert_array_equal(s.value, expected) + # A Scalar expression assigned to many elements is computed, then cast. + u = Vector(f8s, size=3) + u[[0, 2]] << unary.ainv(gb.Scalar.from_value([1, 2, 3], dtype=i8s)) + assert u.to_coo()[1].tolist() == [[-1, -2, -3]] * 2 with pytest.raises(DomainMismatch, match="a UDT casts only to another UDT"): gb.Scalar(FP64) << v.reduce(agg.sum) w = Vector(f8s, size=3) From 97469285b71fbd6618fe9b5cb56eaaf4b50d7a42 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 01:34:27 -0500 Subject: [PATCH 16/54] Cover UDF return-type errors and the Numba error summary Tests for operator/udf.py: a UDF on a UDT that returns None, or an array with no array UDT input, or a tuple beside a UDT that is neither a record nor an array, raises UdfParseError naming the type; built-in ops on such a UDT raise their KeyError. The one-line summary of a Numba error falls back to the first non-empty line, then to a fixed sentence. --- graphblas/tests/test_op.py | 45 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index f946f4dd8..82d091c62 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2050,6 +2050,17 @@ def _three(x, y): # pragma: no cover (numba; raises before execution) with pytest.raises(UdfParseError, match="tuple of length 3.*expected 2"): op_three[udt] + # Neither a scalar, nor a tuple, nor an array that fits an array UDT input. + def _nothing(x, y): # pragma: no cover (numba; raises before execution) + return None + + def _array(x, y): # pragma: no cover (numba; raises before execution) + return np.zeros(2) + + for func in [_nothing, _array]: + with pytest.raises(UdfParseError, match="returned an unsupported type"): + BinaryOp.register_anonymous(func, is_udt=True)[udt] + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow @@ -2067,6 +2078,19 @@ def _bad_tuple(x, y): # pragma: no cover (numba; raises before execution) with pytest.raises(UdfParseError, match="array UDTs of shape.*numpy array"): op[audt] + # A UDT that is neither a record nor an array takes no tuple either, and no + # built-in op lifts to it. + raw = dtypes.register_anonymous(np.dtype("S8"), "_RetErrBytesUDT") + + def _pair(x, y): # pragma: no cover (numba; raises before execution) + return (1, 2) + + with pytest.raises(UdfParseError, match="returned an unsupported type"): + BinaryOp.register_anonymous(_pair, is_udt=True)[raw] + for lifted in [binary.plus, unary.ainv]: + with pytest.raises(KeyError, match="require a record dtype .* or an array dtype"): + lifted[raw] + @pytest.fixture(scope="module") def record_udt(): @@ -5763,6 +5787,27 @@ def test_udt_eq_compile_failure_is_a_udfparseerror(): v + 1 +@pytest.mark.skipif("not supports_udfs") +def test_udt_numba_error_summary_is_never_empty(): + """A Numba error's one-line summary falls back to its first line, then to a sentence. + + The summary is Numba's diagnostic line when one is there, with the + signature from the line below when the diagnostic ends in a colon. A + message without one gives its first non-empty line, and an empty message + a fixed sentence, so a UdfParseError always says something. + """ + from graphblas.core.operator.udf import _summarize_numba_typing_error as summarize + + line = "No implementation of function f found for signature:" + colored = f"\x1b[1m{line}\x1b[0m" + assert summarize(ValueError(f"{colored}\n\n >>> f(int64)")) == f"{line} >>> f(int64)" + assert summarize(ValueError(f"Failed in nopython mode\n{colored}\n\n")) == line + assert summarize(ValueError("\n \nSomething else went wrong\nat line 2")) == ( + "Something else went wrong" + ) + assert summarize(ValueError("")) == "Numba could not compile the function for these input types" + + @pytest.mark.skipif("not supports_udfs") def test_udt_pair_is_never_unified(): """Two different UDTs reach the op as they are, never unified to one type. From 1abf25030c1c6ef8953ce2f4b0eb719395cf9550 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 01:34:32 -0500 Subject: [PATCH 17/54] Cover the JIT layout checks for records a C struct cannot match A record with explicit offsets can have a C struct's size with a field elsewhere, as can such a record nested in one whose own offsets match, and a nested record with an array field has no C typedef. Each skips the JIT and still computes with the cfunc. --- graphblas/tests/test_ssjit.py | 39 +++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 678f852fa..bc74fb03d 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -1618,6 +1618,45 @@ def test_packed_nested_record_layout_skips_jit(): assert udt.jit_c_definition is None +@pytest.mark.skipif("not supports_udfs") +def test_record_layouts_c_cannot_declare_skip_jit(): + """A record whose fields a C struct cannot lay out the same way skips the JIT. + + Explicit offsets can give a record a C struct's size with a field + somewhere else, as can such a record nested in one whose own offsets + match. A nested record with an array field has no C struct here, so the + record holding it has none either. Each computes with the cfunc. + """ + if _IS_SSGB7: + pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + + # Same size as the C struct (8 bytes), but the int32 at offset 1, not 4. + offset = np.dtype( + { + "names": ["ol_a", "ol_b"], + "formats": [np.int8, np.int32], + "offsets": [0, 1], + "itemsize": 8, + } + ) + outer = np.dtype([("ol_in", offset), ("ol_z", np.float64)]) + with_array = np.dtype( + [("ol_x", np.float64), ("ol_rec", [("ol_v", np.float64, (2,)), ("ol_w", np.float64)])], + align=True, + ) + for np_type, value, doubled in [ + (offset, (1, 400), (2, 800)), + (outer, ((1, 400), 2.5), ((2, 800), 5.0)), + (with_array, (1.5, ([2, 3], 4)), (3.0, ([4, 6], 8))), + ]: + udt = dtypes.register_anonymous(np_type) + assert udt.jit_c_name is None + assert udt.jit_c_definition is None + v = gb.Vector(udt, 1) + v[0] = value + assert ((v + v).new().to_coo()[1] == np.array([doubled], dtype=np_type)).all() + + @pytest.mark.slow @pytest.mark.skipif("not supports_udfs") def test_nested_record_udt_jits(): From 8e977df95ec927cf6c0ed0694ad48da231409e2f Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 01:34:51 -0500 Subject: [PATCH 18/54] Remove branches no caller can reach in the UDT code Behavior-neutral; each guard tests a state its callers never produce: - literals._check_literal_fits: every caller passes a UDT, never None. - select._value_select_in_own_type: an int too large for any float never reaches the float branch (beside a float type it is typed as that type, or float() raises first, as in numpy). - udt_codegen._c_aligned_version is only called with records. - udt_codegen._c_expr_unary: both callers pass the field's dtype. - udt_codegen._make_jit_c_comparison_definition: _udt_c_typedef has already returned None for a UDT that is neither array nor record. - udt_codegen.set_jit_c_on_op and set_jit_c_comparison_on_op: callers check _has_jit_set first, since the setter they pass does not exist without it. --- graphblas/core/literals.py | 5 +--- graphblas/core/operator/select.py | 5 +--- graphblas/core/operator/udt_codegen.py | 41 +++++++++++--------------- 3 files changed, 19 insertions(+), 32 deletions(-) diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index 9392cbca6..c92aa41cc 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -266,11 +266,8 @@ def _check_literal_fits(dtype, value): fits ``INT8[3]`` and ``np.int64(1)`` does not. ``Scalar.from_value`` then converts it, which changes no value. A literal that does not fit by type but converts exactly still converts, with a DeprecationWarning - (:func:`_refuse_unless_exact`). A ``dtype`` of ``None``, a built-in type, - takes any literal. + (:func:`_refuse_unless_exact`). """ - if dtype is None: - return literal_type = _literal_type(dtype, value) if literal_type is not None and not _casts_safely_into(literal_type, dtype): _refuse_unless_exact( diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 802219b53..6c2c97576 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -67,10 +67,7 @@ def _value_select_in_own_type(op, dtype, thunk, given): if value != value: # NaN: only ``!=`` holds, for every element new_name, new_value = name, nan else: - try: - near = ftype(value) - except OverflowError: # an int too large for any float - near = ftype(math.copysign(math.inf, value)) + near = ftype(value) above = near if float(near) >= value else np.nextafter(near, ftype(math.inf)) below = near if float(near) <= value else np.nextafter(near, ftype(-math.inf)) new_name = name diff --git a/graphblas/core/operator/udt_codegen.py b/graphblas/core/operator/udt_codegen.py index 09bed96fa..567389660 100644 --- a/graphblas/core/operator/udt_codegen.py +++ b/graphblas/core/operator/udt_codegen.py @@ -405,8 +405,6 @@ def _c_aligned_version(np_type): would emit, we need to rebuild the inner as ``align=True`` first, then use that rebuilt type as the outer's field type. """ - if np_type.names is None: - return np_type fields = [] for name in np_type.names: f = np_type.fields[name][0] @@ -1334,7 +1332,7 @@ def _c_float_floordiv_stmt(target, lhs, rhs, field_dtype): ) -def _c_expr_unary(py_op, operand, field_dtype=None): +def _c_expr_unary(py_op, operand, field_dtype): """Return a C expression for a unary op: e.g., ``-(x->a)``. ``field_dtype`` is only consulted for ``abs``: @@ -1349,13 +1347,12 @@ def _c_expr_unary(py_op, operand, field_dtype=None): ``abs`` of signed int wraps on INT_MIN, both matching the cfunc. """ if py_op == "abs": - if field_dtype is not None: - if field_dtype.kind == "f": - fn = "fabsf" if field_dtype.itemsize == 4 else "fabs" - return f"{fn}({operand})" - if field_dtype.kind == "c": - fn = "cabsf" if field_dtype.itemsize == 8 else "cabs" - return f"{fn}({operand})" + if field_dtype.kind == "f": + fn = "fabsf" if field_dtype.itemsize == 4 else "fabs" + return f"{fn}({operand})" + if field_dtype.kind == "c": + fn = "cabsf" if field_dtype.itemsize == 8 else "cabs" + return f"{fn}({operand})" return f"(({operand}) < 0 ? -({operand}) : ({operand}))" if py_op == "-": return f"-({operand})" @@ -1468,11 +1465,9 @@ def _make_jit_c_comparison_definition(op_name, dtype, *, is_eq): f"_Bool r = {int(is_eq)} ; for (int64_t i = 0 ; i < {size} ; i++) " f"{{ r {accumulate} ((x->v[i]) {op} (y->v[i])) ; }} *z = r ;" ) - elif np_type.names is not None: + else: terms = [f"((x->{c}) {op} (y->{c}))" for _py, c, _d in _iter_record_leaves(np_type)] body = f"*z = {join.join(terms) if terms else int(is_eq)} ;" - else: - return None params = f"_Bool *z, const {type_name} *x, const {type_name} *y" return c_name, f"void {c_name} ({params}) {{ {body} }}" @@ -1513,12 +1508,11 @@ def set_jit_c_comparison_on_op(gb_obj, op_name, dtype, set_string_func, *, is_eq Sibling of :func:`set_jit_c_on_op` for the structurally different comparison kernel (BOOL output, chained leaf comparisons). Returns - ``(c_name, c_defn)`` on success, ``None`` when JIT isn't available or - the dtype can't be expressed in C; callers cache the returned strings - on the op for introspection (see ``TypedUserBinaryOp.jit_c_source``). + ``(c_name, c_defn)`` on success, ``None`` when the dtype can't be + expressed in C; callers cache the returned strings on the op for + introspection (see ``TypedUserBinaryOp.jit_c_source``). Callers check + ``_has_jit_set`` first, as for :func:`set_jit_c_on_op`. """ - if not _has_jit_set: - return None result = _make_jit_c_comparison_definition(op_name, dtype, is_eq=is_eq) if result is None: return None @@ -1529,13 +1523,12 @@ def set_jit_c_comparison_on_op(gb_obj, op_name, dtype, set_string_func, *, is_eq def set_jit_c_on_op(gb_obj, op_name, py_op, dtype, set_string_func, arity=2): """Generate and set JIT C name+definition on a GrB operator, if possible. - Returns ``(c_name, c_defn)`` when the JIT setters are available and the - dtype is expressible in C, ``None`` otherwise. Callers may cache the - returned strings for introspection (see ``TypedUserUnaryOp.jit_c_source`` - and ``TypedUserBinaryOp.jit_c_source``). + Returns ``(c_name, c_defn)`` when the dtype is expressible in C, ``None`` + otherwise. Callers may cache the returned strings for introspection (see + ``TypedUserUnaryOp.jit_c_source`` and ``TypedUserBinaryOp.jit_c_source``). + Callers check ``_has_jit_set`` first: without it, SuiteSparse has no + ``set_string_func`` to pass. """ - if not _has_jit_set: - return None result = _make_jit_c_definition(op_name, py_op, dtype, arity) if result is None: return None From 95999ec85c3b7f0e4b457d49e7fb43fd22ee717f Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 01:35:01 -0500 Subject: [PATCH 19/54] Exclude Numba-compiled bodies and unreachable guards from coverage - numba: _complex_times and _complex_truediv run only inside compiled code, as do two UDFs in test_op.py. - safety: guards no public call reaches: a record leaf lookup_dtype refuses (a void field fails to register first), a Numba error in the UDF probe (the wrapper has the same operand types), a tuple whose length differs from the record _resolve_udt_return_type matched, a C expression for an op or dtype no lifted op gives, and the frame walk in _warn_from_caller running off the stack. - only on failure: the udt_op_path fixture's demotion check, and the re-raise for a large UDT off Windows. --- graphblas/core/operator/udf.py | 11 +++++++---- graphblas/core/operator/udt_codegen.py | 8 ++++---- graphblas/core/utils.py | 2 +- graphblas/tests/test_op.py | 10 ++++++---- 4 files changed, 18 insertions(+), 13 deletions(-) diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index 41e1b7c9d..55a78adfe 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -170,8 +170,8 @@ def _leaves(d): leaf_dtypes = [leaf for _py, _c, leaf in _leaves(d)] try: expected = [lookup_dtype(t).numba_type for t in leaf_dtypes] - except (ValueError, TypeError): - continue + except (ValueError, TypeError): # pragma: no cover (safety) + continue # no record UDT has such a leaf: a void field fails to register if all(et == tt for et, tt in zip(expected, numba_ret_type.types, strict=True)): return d # No perfect match; fall back to the first arity-compatible UDT. @@ -447,7 +447,8 @@ def _run_udf_probe(numba_func, arg_dtypes): args = [_udf_probe_value(d) for d in arg_dtypes] try: return (numba_func(*args),) - except NumbaError as exc: + 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 @@ -536,7 +537,9 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ # The tuple is flat over the leaves, one element per leaf. leaves = list(_iter_record_leaves(return_type.np_type)) if len(numba_ret_type.types) != len(leaves): - return # a wrong field count, not a wrong shape + # _resolve_udt_return_type matched the record to the tuple's length, + # so this would be a wrong field count, not a wrong shape. + return # pragma: no cover (safety) slots = [ (path, numba_ret_type.types[i], leaf.shape, i) for i, (path, _c, leaf) in enumerate(leaves) diff --git a/graphblas/core/operator/udt_codegen.py b/graphblas/core/operator/udt_codegen.py index 567389660..81d78f73f 100644 --- a/graphblas/core/operator/udt_codegen.py +++ b/graphblas/core/operator/udt_codegen.py @@ -485,7 +485,7 @@ def _expr_binary(py_op, x_expr, y_expr, x_dtype, y_dtype): return f"{x_expr} {py_op} {y_expr}" @numba.njit(error_model="numpy") - def _complex_times(x, y): + def _complex_times(x, y): # pragma: no cover (numba) """Multiply two complex numbers by the textbook formula, as numpy does. Numba's own ``*`` is CPython's, which from Python 3.14 (Numba 0.68) @@ -503,7 +503,7 @@ def _complex_times(x, y): return complex(ar * br - ai * bi, ar * bi + ai * br) @numba.njit(error_model="numpy") - def _complex_truediv(x, y): + def _complex_truediv(x, y): # pragma: no cover (numba) """Divide two complex numbers by Smith's method, as the C kernel does. This is CPython's ``_Py_c_quot`` before Python 3.14 (scale by the @@ -1195,7 +1195,7 @@ def _c_truediv_expr(lhs, rhs, field_dtype): Complex division is a statement (:func:`_c_complex_truediv_stmt`). """ if not _is_float(field_dtype): - raise ValueError(f"no C expression divides {field_dtype}") + raise ValueError(f"no C expression divides {field_dtype}") # pragma: no cover (safety) return f"({lhs}) / ({rhs})" @@ -1356,7 +1356,7 @@ def _c_expr_unary(py_op, operand, field_dtype): return f"(({operand}) < 0 ? -({operand}) : ({operand}))" if py_op == "-": return f"-({operand})" - raise ValueError(f"Unknown unary C op: {py_op}") + raise ValueError(f"Unknown unary C op: {py_op}") # pragma: no cover (safety) def _make_jit_c_definition(op_name, py_op, dtype, arity): diff --git a/graphblas/core/utils.py b/graphblas/core/utils.py index 86ec61a8a..6b8ec12ed 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -64,7 +64,7 @@ def _warn_from_caller(message, category): tests = package / "tests" frame = sys._getframe(1) stacklevel = 2 - while frame is not None: + while frame is not None: # pragma: no branch (safety) filename = Path(frame.f_code.co_filename) if not filename.is_relative_to(package) or filename.is_relative_to(tests): break diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 82d091c62..fdbdd5c5e 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1493,7 +1493,9 @@ def udt_op_path(request): demoted = path == "jit" and gb.ss.config["jit_c_control"] != "on" gb.ss.config["jit_c_control"] = previous if demoted: - pytest.fail("SuiteSparse demoted jit_c_control; the JIT path did not run") + pytest.fail( # pragma: no cover (only on failure) + "SuiteSparse demoted jit_c_control; the JIT path did not run" + ) def _udt_vectors(udt, xs, ys=None): @@ -4141,7 +4143,7 @@ def test_udt_array_ops_on_large_elements(): T = dtypes.register_anonymous(np.dtype((np.float64, (64, 64))), "_LoopBig") except InvalidValue: if sys.platform != "win32": - raise + raise # pragma: no cover (only on failure) # SuiteSparse:GraphBLAS built by MSVC (no variable-length arrays) refuses a # UDT larger than 1024 bytes (GB_VLA_MAXSIZE); this one is 32768. pytest.skip("SuiteSparse:GraphBLAS limits a UDT to 1024 bytes on Windows") @@ -4365,10 +4367,10 @@ def test_udt_user_op_named_like_a_lifted_op(): """ udt = dtypes.register_anonymous(np.dtype((np.int64, (2, 3))), "_NamedLikeLifted") - def scaled_plus(x, y): + def scaled_plus(x, y): # pragma: no cover (numba) return x * 10 + y - def difference(x, y): + def difference(x, y): # pragma: no cover (numba) return x - y plus_like = BinaryOp.register_anonymous(scaled_plus, "plus", is_udt=True) From 3f8ed9f9f3fe82470f480484f70ea42c50ddd4a7 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 01:35:06 -0500 Subject: [PATCH 20/54] Drop test conditions that cannot change and a needless Numba skip - test_dtype: the JIT name test is skipped unless _has_jit_set, which needs the suitesparse backend, so it always has gb.ss. - test_ssjit: the two complex JIT tests are skipped without Numba, so their own "numba is None" skips never ran. - test_op: test_compile_codegen_helper compiles plain Python, so its Numba skip was needless; it now runs without Numba too, which covers _compile_codegen's no-Numba namespace in CI's no-Numba jobs. --- graphblas/tests/test_dtype.py | 15 ++++++--------- graphblas/tests/test_op.py | 1 - graphblas/tests/test_ssjit.py | 4 ---- 3 files changed, 6 insertions(+), 14 deletions(-) diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index b4c353767..57358ec73 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -554,9 +554,8 @@ def test_udt_layouts_under_one_name_get_their_own_jit_c_names(): assert small.jit_c_name == "_SameJitName" assert big.jit_c_name != small.jit_c_name assert big.jit_c_name in big.jit_c_definition - jit_c_control = gb.ss.config["jit_c_control"] if gb.backend == "suitesparse" else None - if jit_c_control is not None: - gb.ss.config["jit_c_control"] = "on" + jit_c_control = gb.ss.config["jit_c_control"] + gb.ss.config["jit_c_control"] = "on" try: for udt in [small, big]: shape = udt.np_type.shape @@ -569,12 +568,10 @@ def test_udt_layouts_under_one_name_get_their_own_jit_c_names(): values = [v[i].new().value.tolist() for i in range(3)] assert values == [one.tolist(), one.tolist(), two.tolist()] finally: - if jit_c_control is not None: - gb.ss.config["jit_c_control"] = jit_c_control - if gb.backend == "suitesparse": - # Nor may a typedef registered in C take a name a UDT already has. - with pytest.raises(ValueError, match="unavailable"): - dtypes.ss.register_new("_SameJitName", "typedef struct { int8_t q ; } _SameJitName ;") + gb.ss.config["jit_c_control"] = jit_c_control + # Nor may a typedef registered in C take a name a UDT already has. + with pytest.raises(ValueError, match="unavailable"): + dtypes.ss.register_new("_SameJitName", "typedef struct { int8_t q ; } _SameJitName ;") def test_dtype_to_from_string(): diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index fdbdd5c5e..a4cd03582 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -6210,7 +6210,6 @@ def test_ops_have_ss(): mod.ss -@pytest.mark.skipif("not supports_udfs") def test_compile_codegen_helper(): """The ``_compile_codegen`` helper validates source and surfaces typos clearly. diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index bc74fb03d..2b83a4f85 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -1195,8 +1195,6 @@ def test_complex_truediv_udt_jit_matches_cfunc_at_extremes(): if not _has_jit_set: pytest.skip("jit_c_source introspection requires SuiteSparse:GraphBLAS >= 9") _require_jit_on() - if numba is None: - pytest.skip("numba required for the cfunc baseline") udt = dtypes.register_anonymous( np.dtype([("cxd_s", np.complex64), ("cxd_d", np.complex128)], align=True), "_CxDivJit" @@ -1249,8 +1247,6 @@ def test_complex_times_udt_jit_matches_cfunc_with_infinities(): if not _has_jit_set: pytest.skip("jit_c_source introspection requires SuiteSparse:GraphBLAS >= 9") _require_jit_on() - if numba is None: - pytest.skip("numba required for the cfunc baseline") rec = dtypes.register_anonymous( np.dtype([("cxt_s", np.complex64), ("cxt_d", np.complex128)], align=True), "_CxTimesJit" From 8890be907f99d772fb79208a9cd5a97d1ec9048b Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 01:47:45 -0500 Subject: [PATCH 21/54] Fix select on complex values with a thunk of another type Behavior fix: wrong results before, on main too (not a #591 regression). SuiteSparse:GraphBLAS reads the thunk of a value select as the input's type when the op is typed for another type, and the workaround from 69173110 converted the thunk only for bool, integer and float inputs. On an FC32 vector of [1, 2+1j, 3, nan], select(valueeq, np.complex128(1)) kept nothing (now [0]), valueeq with np.float64(3) or np.int64(3) kept nothing (now [2]), and valuene with np.complex128(2+1j) kept everything (now [0, 2, 3]); apply with the same ops was right. (Main also got a Python 3 wrong there; #591's typing of a weak int as FC32 already fixed that.) A complex input now gets the thunk in its own type when that converts exactly; otherwise (or with a NaN part) the thunk equals no element, so valueeq keeps nothing and valuene keeps everything. "Exactly" is as for float inputs: the thunk in the input's type must equal the thunk, part by part, in Python's int and float comparison, so np.int64(2**53 + 1) equals no complex64 element, as it equals no float32 one, where apply compares in complex128 and matches 2**53. The ordered value ops still raise KeyError on complex values, which have no order. A real input with a complex thunk keeps its rule: an imaginary part of 0 converts, any other equals nothing. --- graphblas/core/operator/select.py | 25 +++++++++++----- graphblas/tests/test_vector.py | 48 ++++++++++++++++++++++++++++++- 2 files changed, 65 insertions(+), 8 deletions(-) diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 6c2c97576..136f362dd 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -25,9 +25,11 @@ def _value_select_in_own_type(op, dtype, thunk, given): same elements: for integers, ``x < 2.5`` is ``x < 3``, and a thunk no element can equal, or that every element passes, selects nothing or everything; for floats, the thunk is rounded toward the side that keeps - the comparison exact. ``given`` is the thunk as the caller gave it, whose - value is exact where ``thunk`` may have rounded it (an int beyond 64 bits - is FP64). Other ops and types are returned unchanged. + the comparison exact; for complex numbers, which have no order, a thunk + that does not convert exactly equals no element. ``given`` is the thunk as + the caller gave it, whose value is exact where ``thunk`` may have rounded + it (an int beyond 64 bits is FP64). Other ops and types are returned + unchanged. """ # MAINT 2026-10-07: works around the SuiteSparse:GraphBLAS bug above (seen in # 8.0.2 through 10.6.0; fix proposed in DrTimothyAldenDavis/GraphBLAS#463). @@ -47,22 +49,31 @@ def _value_select_in_own_type(op, dtype, thunk, given): if ( parent is not getattr(select, name) and parent is not getattr(indexunary, name) - or dtype.np_type.kind not in "biuf" + or dtype.np_type.kind not in "biufc" or thunk._is_empty ): return op, thunk value = given.value if type(given) is Scalar else given if isinstance(value, (np.generic, np.ndarray)): value = value.item() - if isinstance(value, complex): + np_type = dtype.np_type + if isinstance(value, complex) and np_type.kind != "c": if value.imag != 0: if name not in {"valueeq", "valuene"}: # pragma: no cover (no complex order) return op, thunk value = math.nan # equal to nothing, as a complex number is to a real one else: value = value.real - np_type = dtype.np_type - if np_type.kind == "f": + if np_type.kind == "c": + # Only valueeq and valuene take complex values. The thunk is compared + # exactly, as for floats: one that does not convert exactly (2**53 + 1 + # beside complex64, or a NaN part) equals no element, as NaN does. + with np.errstate(over="ignore", invalid="ignore"): + near = np_type.type(value) + real, imag = (value.real, value.imag) if isinstance(value, complex) else (value, 0) + exact = float(near.real) == real and float(near.imag) == imag + new_name, new_value = name, near if exact else np_type.type(complex(math.nan, math.nan)) + elif np_type.kind == "f": nan, ftype = np_type.type(np.nan), np_type.type if value != value: # NaN: only ``!=`` holds, for every element new_name, new_value = name, nan diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 81e36d851..67be00169 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -1704,14 +1704,60 @@ def test_select_value_thunk_of_another_type(): ]: assert x.select(op, thunk).new().to_coo()[0].tolist() == kept, (x.dtype, op, thunk) if dtypes._supports_complex: - # A real value equals a complex one only when its imaginary part is 0. + # A real value equals a complex one only when its imaginary part is 0, and + # a complex value equals a thunk of another type only when that converts + # exactly. Complex values have no order, so only valueeq and valuene apply. + fc32 = Vector.from_coo( + [0, 1, 2, 3, 4, 5, 6], + [1, 2 + 1j, 3, np.nan, np.complex64(0.1), np.inf, np.complex64(1 + 0.1j)], + dtype=dtypes.FC32, + ) + fc64 = Vector.from_coo([0, 1, 2], [1, 2 + 1j, 0.1], dtype=dtypes.FC64) + # Compared exactly, as floats are: 2**53 + 1 is not the float it rounds to. + big32 = Vector.from_coo([0, 1], [2**53, 2**24], dtype=dtypes.FC32) + bigf32 = Vector.from_coo([0, 1], [2**53, 2**24], dtype=dtypes.FP32) for x, op, thunk, kept in [ (i8, select.valueeq, 1 + 0j, [1]), (i8, select.valueeq, 1 + 2j, []), (i8, select.valuene, 1 + 2j, [0, 1, 2, 3]), (f32, select.valueeq, np.complex64(1), [1]), + (f32, select.valueeq, np.complex128(0.1), []), # float32(0.1) is not 0.1 + (f32, select.valuene, np.complex128(1 + 2j), [0, 1, 2]), + (fc32, select.valueeq, np.complex128(1), [0]), + (fc32, select.valueeq, np.complex128(2 + 1j), [1]), + (fc32, select.valueeq, np.complex64(2 + 1j), [1]), + (fc32, select.valueeq, 3, [2]), + (fc32, select.valuene, np.complex128(2 + 1j), [0, 2, 3, 4, 5, 6]), + (fc32, select.valueeq, np.float64(3), [2]), + (fc32, select.valueeq, np.int64(3), [2]), + (fc32, select.valuene, np.int64(3), [0, 1, 3, 4, 5, 6]), + (fc32, select.valueeq, np.complex64(0.1), [4]), + (fc32, select.valueeq, np.complex128(0.1), []), # complex64(0.1) is not 0.1 + (fc32, select.valuene, np.complex128(0.1), [0, 1, 2, 3, 4, 5, 6]), + (fc32, select.valueeq, np.complex128(complex(np.nan, 0)), []), + (fc32, select.valueeq, np.float64(np.inf), [5]), + (fc32, select.valueeq, np.complex128(1e300), []), # inf as complex64, not 1e300 + (fc32, select.valuene, np.complex128(1e300), [0, 1, 2, 3, 4, 5, 6]), + (fc32, select.valueeq, np.complex64(1 + 0.1j), [6]), + (fc32, select.valueeq, np.complex128(1 + 0.1j), []), # in the imaginary part + (fc32, indexunary.valueeq, np.complex128(2 + 1j), [1]), + (fc32, select.valueeq, Scalar.from_value(np.complex128(2 + 1j)), [1]), + (big32, select.valueeq, np.int64(2**53), [0]), + (big32, select.valueeq, np.int64(2**53 + 1), []), + (bigf32, select.valueeq, np.int64(2**53 + 1), []), + (fc64, select.valueeq, np.complex128(2 + 1j), [1]), + (fc64, select.valueeq, np.complex64(0.1), []), # complex64(0.1) is not 0.1 + (fc64, select.valueeq, np.float64(0.1), [2]), + (fc64, select.valuene, np.int64(1), [1, 2]), + (fc64, select.valueeq, 1, [0]), ]: assert x.select(op, thunk).new().to_coo()[0].tolist() == kept, (x.dtype, op, thunk) + with pytest.raises(KeyError, match="valuelt does not work with FC64"): + fc32.select(select.valuelt, np.complex128(1)) + C = Matrix.from_coo([0, 0, 1], [0, 1, 1], [1, 2 + 1j, 3], dtype=dtypes.FC32) + equal = C.select(select.valueeq, np.complex128(2 + 1j)).new() + assert equal.to_coo()[2].tolist() == [2 + 1j] + assert C.select(select.valuene, np.int64(3)).new().nvals == 2 with pytest.raises(EmptyObject): i8.select(select.valueeq, Scalar(dtypes.INT64)).new() A = Matrix.from_coo([0, 0, 1], [0, 1, 1], [1, 100, -5], dtype=dtypes.INT8) From dd44ae8d9ffedff5a36dfb68452d14a37126f695 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 02:02:11 -0500 Subject: [PATCH 22/54] Compare UDTs that are neither records nor arrays in eq and ne Behavior fix. A UDT of bytes, dates or the like holds one value per element. binary.eq and binary.ne on it raised TypeError ('NoneType' object is not subscriptable), and so did Vector and Matrix isequal and ==: on main too. A Scalar's == raised it as well, where main compared the values (a #591 regression, from 69173110's Scalar == through binary.eq), as did == with a literal (b"ab"), which the literal checks read as an array UDT. On the vanilla backend a complex literal is such a UDT, so int8_vec == 1j raised the same TypeError. eq and ne now compare the two values as numpy does, as a record field of that type compares (bytes by value, NaT equal to nothing), and only two elements of the one UDT compare: any other pair raises KeyError, from _check_udt_pair, with or without Numba. So int8_vec == 1j on vanilla is a KeyError, and so is a Scalar's == with another type (another bytes UDT, or a number), where main compared the values (True or False). A literal beside such a UDT is checked as one value. A Scalar's == on such a UDT compares in numpy, as without Numba, so it needs no Numba type for the element: a float16 UDT, which Numba cannot compile, compares as on main. Vector and Matrix == on float16 still need Numba and raise its NotImplementedError, as main raised there too. --- graphblas/core/literals.py | 10 +++--- graphblas/core/operator/binary.py | 26 ++++++++++---- graphblas/core/scalar.py | 8 ++++- graphblas/core/udt.py | 17 ++++++--- graphblas/tests/test_op.py | 59 +++++++++++++++++++++++++++++++ graphblas/tests/test_scalar.py | 7 ++++ graphblas/tests/test_vector.py | 8 +++-- 7 files changed, 116 insertions(+), 19 deletions(-) diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index c92aa41cc..6b025e2d5 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -289,7 +289,9 @@ def _check_literal_fits(dtype, value): if isinstance(value, dict): value = _dict_to_record(np_type, value) else: - target = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) + # An array's elements, or the one value of a UDT that is neither (bytes). + base, shape = np_type.subdtype or (np_type, ()) + target = np.zeros(shape, dtype=base) with warnings.catch_warnings(): warnings.simplefilter("error", RuntimeWarning) try: @@ -358,14 +360,14 @@ def _converts_exactly(dtype, value): given = np.zeros(1, dtype=_object_layout(np_type)) converted = np.zeros(1, dtype=np_type) else: - base, shape = np_type.subdtype + base, shape = np_type.subdtype or (np_type, ()) given = np.empty(shape, dtype=object) converted = np.zeros(shape, dtype=base) - given[:] = value + given[...] = value with warnings.catch_warnings(): # numpy warns for some lossy conversions; the comparison below finds each. warnings.simplefilter("ignore") - converted[:] = _view_if_same_layout(value, converted.dtype) + converted[...] = _view_if_same_layout(value, converted.dtype) except (TypeError, ValueError, OverflowError): return False for given_leaf, converted_leaf in zip( diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 966a49047..c592fac98 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -61,8 +61,9 @@ def _make_udt_comparison(dtype, dtype2, *, is_eq): Compares per-leaf with IEEE semantics (NaN compares unequal even to itself); scalar broadcasts to every leaf when paired with a - UDT. Mismatched-shape UDT pairs raise ``KeyError``. Returns - ``(wrapper_func, wrapper_sig)``. + UDT. Mismatched-shape UDT pairs raise ``KeyError``. A UDT that is + neither a record nor an array (bytes, say) compares only with its + own type, as one value. Returns ``(wrapper_func, wrapper_sig)``. """ from ..udt import _check_udt_pair, _get_udt_info, _iter_record_leaves from .udt_codegen import _array_operand_source @@ -72,15 +73,28 @@ def _make_udt_comparison(dtype, dtype2, *, is_eq): op_name = "eq" if is_eq else "ne" _check_udt_pair(op_name, dtype, dtype2, info_x, info_y) + nt = numba.types + cmp_op = "==" if is_eq else "!=" + join = " and " if is_eq else " or " + + if info_x is None and info_y is None: + # Bytes, a datetime, or a complex number where the backend has no + # complex types: an element is one value, compared as numpy compares + # it, as a record field of that type is. ``_check_udt_pair`` has made + # sure both are the one UDT. + ptr_type = nt.CPointer(dtype.numba_type) + wrapper = _compile_codegen( + f"def wrapper(z_ptr, x_ptr, y_ptr):\n z_ptr[0] = x_ptr[0] {cmp_op} y_ptr[0]\n", + func_name="wrapper", + source_label=f"", + ) + return wrapper, nt.void(nt.CPointer(INT8.numba_type), ptr_type, ptr_type) + x_is_scalar = info_x is None y_is_scalar = info_y is None udt_dtype = dtype2 if x_is_scalar else dtype udt_info = info_y if x_is_scalar else info_x - nt = numba.types - cmp_op = "==" if is_eq else "!=" - join = " and " if is_eq else " or " - if udt_info[0] == "array": # Array UDT: a scalar side is read through its pointer, with no # ``numba.carray``. The wrapper sees an array side as a flat diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 109067250..84c9fb946 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -22,6 +22,7 @@ _element_shapes_differ, _element_value, _ewise_add_needs_cast, + _get_udt_info, _is_lifted_op, _literal_shape_differs, _scalar_ewise_add_converted, @@ -1163,7 +1164,12 @@ def _scalar_eq(scalar, other): with np.errstate(over="ignore"): return (convert(scalar.value) == convert(value)).item() return (convert(scalar.value) == convert(value)).item() - if not _has_numba and (scalar.dtype._is_udt or (type(other) is Scalar and other.dtype._is_udt)): + udts = [scalar.dtype, other.dtype] if type(other) is Scalar else [scalar.dtype] + udts = [dtype for dtype in udts if dtype._is_udt] + if udts and (not _has_numba or any(_get_udt_info(dtype) is None for dtype in udts)): + # binary.eq needs Numba. A UDT that is neither a record nor an array (bytes, + # float16) holds one value, which numpy compares as binary.eq does, and with + # no Numba type for it (float16 has none) and no GraphBLAS call. return _udt_eq_without_numba(scalar, other) return bool(scalar.ewise_mult(other, binary.eq).new(name="s_eq").value) diff --git a/graphblas/core/udt.py b/graphblas/core/udt.py index f24a3ab11..13acf752c 100644 --- a/graphblas/core/udt.py +++ b/graphblas/core/udt.py @@ -110,12 +110,19 @@ def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): """Raise ``KeyError`` if two UDT operands disagree on shape. ``info_x`` and ``info_y`` are ``_get_udt_info`` results. When either is - ``None`` (one side is a scalar broadcast), no check fires. Without this - gate the generated wrapper code raises a cryptic Numba ``TypingError`` - on the first field access, or, for array fields that do not broadcast, - fails inside the cfunc without telling the caller. Element dtypes may - differ; :func:`_udt_op_types` promotes them. + ``None`` (a built-in scalar side broadcasts), no further check fires. A + UDT that is neither a record nor an array (bytes, say) holds one value, + which pairs only with its own type, on either side. Without this gate the + generated wrapper code raises a cryptic Numba ``TypingError`` on the + first field access, or, for array fields that do not broadcast, fails + inside the cfunc without telling the caller. Element dtypes may differ; + :func:`_udt_op_types` promotes them. """ + if dtype != dtype2 and (info_x is None and dtype._is_udt or info_y is None and dtype2._is_udt): + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): a UDT that is " + "neither a record nor an array pairs only with its own type." + ) if info_x is None or info_y is None: return kind_x, detail_x = info_x diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index a4cd03582..62550ee21 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -5789,6 +5789,65 @@ def test_udt_eq_compile_failure_is_a_udfparseerror(): v + 1 +@pytest.mark.skipif("not supports_udfs") +def test_udt_eq_compares_a_udt_that_is_neither_record_nor_array(): + """``==``, ``!=`` and ``isequal`` compare UDT elements of bytes or dates as numpy does. + + Such a UDT holds one value per element, which ``binary.eq`` and + ``binary.ne`` compare as a record field of that type compares: bytes by + value, and NaT equal to nothing. Only two elements of the one type + compare; any other pair raises binary.eq's KeyError. + """ + text = dtypes.register_anonymous(np.dtype("S4"), "_EqOpaqueBytes") + v = Vector(text, size=3) + v[0] = b"ab" + v[1] = b"cd" + w = Vector(text, size=3) + w[0] = b"ab" + w[1] = b"ce" + w[2] = b"x" + assert binary.eq(v & w).new().isequal(Vector.from_coo([0, 1], [True, False], size=3)) + assert (v != w).new().isequal(Vector.from_coo([0, 1], [False, True], size=3)) + assert (v == b"ab").new().isequal(Vector.from_coo([0, 1], [True, False], size=3)) + assert v.isequal(v.dup()) + assert not v.isequal(w) + A = Matrix(text, nrows=2, ncols=2) + A[0, 1] = b"ab" + A[1, 0] = b"cd" + B = A.dup() + assert A.isequal(B) + B[1, 0] = b"cc" + assert not A.isequal(B) + s = gb.Scalar.from_value(b"ab", text) + assert s == gb.Scalar.from_value(b"ab", text) + assert s != gb.Scalar.from_value(b"cd", text) + assert s == b"ab" + assert s != b"cd" + # A Scalar's == compares in numpy, so an element Numba has no type for compares too. + half = dtypes.register_anonymous(np.dtype(np.float16), "_EqOpaqueHalf") + h = gb.Scalar.from_value(1.5, half) + assert h == gb.Scalar.from_value(1.5, half) + assert h == 1.5 + assert gb.Scalar.from_value(np.nan, half) != gb.Scalar.from_value(np.nan, half) + dates = dtypes.register_anonymous(np.dtype("M8[s]"), "_EqOpaqueDates") + d = Vector(dates, size=2) + d[0] = np.datetime64("NaT", "s") + d[1] = np.datetime64(0, "s") + assert (d == d).new().to_coo()[1].tolist() == [False, True] + # Nor with a record or an array UDT, for a Scalar as for a Vector. + pair = dtypes.register_anonymous(np.dtype((np.int8, (5,))), "_EqOpaqueArr5") + rec = dtypes.register_anonymous(np.dtype([("eqo_a", np.int8)]), "_EqOpaqueRec") + for expr in [ + lambda: v == 1, + lambda: binary.eq(v & Vector.from_coo([0], [1], size=3)), + lambda: binary.eq(v & Vector(rec, size=3)), + lambda: gb.Scalar.from_value([1, 2, 3, 4, 5], pair) == h, + lambda: gb.Scalar.from_value((1,), rec) == s, + ]: + with pytest.raises(KeyError, match="pairs only with its own type"): + expr() + + @pytest.mark.skipif("not supports_udfs") def test_udt_numba_error_summary_is_never_empty(): """A Numba error's one-line summary falls back to its first line, then to a sentence. diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index 81223942c..cfb1361c6 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -723,6 +723,8 @@ def test_udt_eq_without_numba_matches_binary_eq(monkeypatch): r = Scalar.from_value((1, 2.5), rec) a = Scalar.from_value([1, 1, 1], arr) f = Scalar.from_value([0.1, 0.1], f32) + text = dtypes.register_anonymous(np.dtype("S6")) + t = Scalar.from_value(b"ab", text) cases = [ (r, (1, 2.5)), (r, (1, 2)), @@ -737,6 +739,9 @@ def test_udt_eq_without_numba_matches_binary_eq(monkeypatch): (a, Scalar.from_value([1, 1, 1], arr)), (f, 0.1), (f, np.float64(0.1)), + (t, b"ab"), + (t, b"cd"), + (t, Scalar.from_value(b"ab", text)), ] # UDTs that binary.eq does not pair raise its KeyError both ways. other_names = dtypes.register_anonymous( @@ -747,6 +752,8 @@ def test_udt_eq_without_numba_matches_binary_eq(monkeypatch): (r, Scalar.from_value((1, 2.5), other_names)), (a, Scalar.from_value([1, 1], arr2)), (r, a), + (t, Scalar.from_value(b"ab", dtypes.register_anonymous(np.dtype("S7")))), + (t, 1), ] def compare_all(): diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 67be00169..212e64d0c 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -1534,9 +1534,11 @@ def test_python_number_is_weak_beside_builtin_dtypes(): if dtypes._supports_complex: assert (f32 + 1j).dtype == dtypes.FC32 else: - # Without complex types, a complex number keeps its own type, which plus does not take. - with pytest.raises(KeyError, match="plus does not work with"): - f32 + 1j + # Without complex types, a complex number keeps its own type, which plus does not + # take, and which compares only with itself. + for expr in [lambda: f32 + 1j, lambda: f32 == 1j]: + with pytest.raises(KeyError, match="does not work with"): + expr() s8 = Scalar.from_value(100, dtype=dtypes.INT8) assert (s8 + 1).new().dtype == (1 + s8).new().dtype == dtypes.INT8 assert (Scalar.from_value(0.5, dtype=dtypes.FP32) * 0.5).new().dtype == dtypes.FP32 From bee1f0fe57e445cd6bc9b67b5d393db570da6bef Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 02:04:12 -0500 Subject: [PATCH 23/54] Type each product of Matrix.power for its UDT operands Behavior fix of a #591 regression: main, which typed UDT ops otherwise, computed int8 results here. power types the semiring for A, and GraphBLAS casts a built-in operand of another type to it, but not a UDT. With a semiring whose result is another type, such as plus_truediv on an int8 array UDT (float64 results), power(2) worked, while from n = 3 on the products of float64 squares with A or with each other raised a bare GrB_DOMAIN_MISMATCH from GrB_mxm. Each product of UDT operands that are not of the types the semiring was typed for now takes the semiring typed for its own operands, as any mxm of them does. power(n) into a float64 array object now equals the power of a float64 copy of A, as for built-in INT8, for every n. Unchanged: A.power(n, plus_truediv).new() makes an object of A's int8 type, and an mxm cannot cast a UDT as it stores, so it raises DomainMismatch by #591's store rule, where main computed int8 results (built-in INT8 truncates as it stores). --- graphblas/core/matrix.py | 21 +++++++++++++++------ graphblas/tests/test_op.py | 18 ++++++++++++------ 2 files changed, 27 insertions(+), 12 deletions(-) diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index d27e507cc..37707f3b2 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -111,6 +111,15 @@ def _reposition(updater, indices, chunk): def _power(updater, A, n, op): opts = updater.opts + + def times(left, right): + # ``op`` is typed for A, and GraphBLAS casts a built-in operand of another + # type to it, but not a UDT: an int8 UDT squares to float64 under truediv. + # So a product of UDT operands of other types takes ``op`` typed for them. + if A.dtype._is_udt and (left.dtype != op.type or right.dtype != op.type2): + return get_typed_op(op.parent, left.dtype, right.dtype, kind="semiring")(left @ right) + return op(left @ right) + if n == 0: v = Vector.from_scalar(op.binaryop.monoid.identity, A._nrows, A.dtype, name="v_diag") updater << v.diag(name="M_diag") @@ -123,7 +132,7 @@ def _power(updater, A, n, op): ): # The last product is stored by mxm, which cannot cast a UDT as it writes; # raise before computing the others. - expr = op(A @ A) + expr = times(A, A) if _store_cast_op(expr.dtype, out_dtype) is not None: _fused_store_expr(expr, out_dtype) # Use repeated squaring: compute A^2, A^4, A^8, etc., and combine terms as needed. @@ -158,14 +167,14 @@ def _power(updater, A, n, op): result = square.dup(name="Power", **opts) elif n == 0: # All done! No more terms to compute - updater << op(result @ square) + updater << times(result, square) return elif result is A: # Now we need to create a new matrix for the intermediate result - result = op(result @ square).new(name="Power", **opts) + result = times(result, square).new(name="Power", **opts) else: # Main branch: multiply `square` into `result` - result(**opts) << op(result @ square) + result(**opts) << times(result, square) n, bit = divmod(n, 2) if square_expr is not None: # We need to perform another squaring, so evaluate current `square_expr` first @@ -177,10 +186,10 @@ def _power(updater, A, n, op): square << square_expr if square is None: # First iteration! Create expression for first square - square_expr = op(A @ A) + square_expr = times(A, A) else: # Expression for repeated squaring - square_expr = op(square @ square) + square_expr = times(square, square) class Matrix(BaseType): diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 62550ee21..c0badf8b4 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4967,12 +4967,18 @@ def test_udt_store_casts_each_element_as_builtin_types_do(): C << A.power(3) C << A.power(1) assert C.to_coo()[2].tolist() == [first] - # A semiring whose result is another type (int truediv is float) stores that - # type as it is. - P = Matrix.from_coo([0], [0], [[[1, -1], [2, 4], [-8, 3]]], dtype=i8s, nrows=1, ncols=1) - Q = Matrix(f64s, nrows=1, ncols=1) - Q << P.power(2, semiring.plus_truediv) - assert Q.isequal(semiring.plus_truediv(P @ P).new()) + # A semiring whose result is another type (int truediv is float) gives that + # type, which stores as it is, and each later product takes the semiring + # typed for its operands, so the powers are those of a float64 copy, as for + # built-in types. + P = Matrix(i8s, nrows=2, ncols=2) + P[0, 0] = np.full((3, 2), 2) + P[0, 1] = np.full((3, 2), 1) + P[1, 1] = [[4, 2], [1, -1], [3, 5]] + Q = Matrix(f64s, nrows=2, ncols=2) + for n in [1, 2, 3, 4, 5]: + Q << P.power(n, semiring.plus_truediv) + assert Q.isequal(P.dup(dtype=f64s).power(n, semiring.plus_truediv).new()), n # Leading axes of length 1 do not change the layout. lead = dtypes.register_anonymous(np.dtype((np.float32, (1, 3, 2))), "_StoreCastLead") u = Vector(lead, size=3) From e64fd32626957f98bb8415f86b64797b9f3e1e88 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 02:13:00 -0500 Subject: [PATCH 24/54] Select with a float64 thunk too large for FP32 without a warning Behavior fix, small: f32_vec.select(valueeq, np.float64(1e300)) gave the right result but numpy warned "overflow encountered in cast" from inside python-graphblas (an error under -W error), as the select workaround from 69173110 converted the thunk to float32. A #591 regression: main had no such conversion (it let SuiteSparse read the thunk's bytes as float32, which was wrong in another way). numpy's own f32_array == np.float64(1e300) does not warn. The conversion now ignores the overflow: the thunk is infinity in float32, which compares with every element as 1e300 does, so the selected elements are unchanged. --- graphblas/core/operator/select.py | 5 ++++- graphblas/tests/test_vector.py | 12 ++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 136f362dd..0fca32ad4 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -78,7 +78,10 @@ def _value_select_in_own_type(op, dtype, thunk, given): if value != value: # NaN: only ``!=`` holds, for every element new_name, new_value = name, nan else: - near = ftype(value) + # A finite value too large for ftype is infinity there, which still + # compares right with every element, so numpy need not warn. + with np.errstate(over="ignore"): + near = ftype(value) above = near if float(near) >= value else np.nextafter(near, ftype(math.inf)) below = near if float(near) <= value else np.nextafter(near, ftype(-math.inf)) new_name = name diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 212e64d0c..87f05dad2 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -1760,6 +1760,18 @@ def test_select_value_thunk_of_another_type(): equal = C.select(select.valueeq, np.complex128(2 + 1j)).new() assert equal.to_coo()[2].tolist() == [2 + 1j] assert C.select(select.valuene, np.int64(3)).new().nvals == 2 + # A float64 too large for FP32 is compared as it is, without numpy's overflow warning. + with warnings.catch_warnings(): + warnings.simplefilter("error") + for op, thunk, kept in [ + (select.valueeq, 1e300, []), + (select.valuene, 1e300, [0, 1, 2]), + (select.valuelt, 1e300, [0, 1]), + (select.valuegt, 1e300, []), + (select.valuegt, -1e300, [0, 1]), + ]: + selected = f32.select(op, np.float64(thunk)).new().to_coo()[0].tolist() + assert selected == kept, (op, thunk) with pytest.raises(EmptyObject): i8.select(select.valueeq, Scalar(dtypes.INT64)).new() A = Matrix.from_coo([0, 0, 1], [0, 1, 1], [1, 100, -5], dtype=dtypes.INT8) From 75f73b87e9eac8337bb333932152298bfb9b0b5f Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 02:52:45 -0500 Subject: [PATCH 25/54] Refuse values that a bytes, date or duration leaf of a UDT would not keep Behavior fix. A literal that must become an element of a UDT and does not fit it by type converts, with a DeprecationWarning, only when its values convert exactly (#591's MAINT rule). The check skipped every leaf that holds no number, so whatever was given for a bytes, str, date or duration leaf counted as exact. Under a user's op, (1.5, 1) for a record of a bytes field and an int8 field became (b"1.5", 1), and 1/3 became b"0.33", with a warning that said "Its values convert exactly"; np.timedelta64(5, "ms") for a timedelta of seconds became 0 s, and a structured literal's six bytes became four, under the same warning. - main: converted silently (1.5 became b"1.5"). - 0e1ad03c: a record field converted with that warning; a bare UDT of bytes or durations raised "does not fit" (the check failed on its layout). - the commit "Compare UDTs that are neither records nor arrays in eq and ne" let a bare UDT convert with the warning too. - now: such a leaf keeps a value only if it converts back to it, so a Python number, a numpy bool, a string cut short or a duration in a coarser unit raises ValueError "does not fit". bytes or str that fit still convert, as does np.timedelta64(5000, "ms") as 5 s (with the warning, since ms to s is not a safe cast). Only a literal whose type does not fit reaches this check. A numpy number that numpy casts safely to the leaf's type fits by type: np.uint8(5) beside a 3-byte bytes UDT is b"5" without a warning, by numpy's rule. bytes or str alone beside a bare bytes UDT has no type, so numpy still cuts b"abcdef" to b"abc" there. A record literal holding a numpy number, text or a date gets no type at all and so no check; the next commit types it. --- graphblas/core/literals.py | 29 ++++++++++++++++++++++++++++- graphblas/tests/test_op.py | 25 +++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index 6b025e2d5..243c34547 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -14,6 +14,7 @@ """ import math +import numbers import struct import warnings from functools import lru_cache @@ -350,7 +351,10 @@ def _converts_exactly(dtype, value): """Whether converting ``value`` into UDT ``dtype`` keeps every number in it. A float may round to a narrower float, as numpy 2 rounds a Python number - beside float32, but may not become infinite or lose an imaginary part. + beside float32, but may not become infinite or lose an imaginary part. A + leaf that holds no number (bytes, a date, a duration) keeps a value only + if it converts back to it: not a number (``1.5`` would become ``b"1.5"``), + a string cut short, or a duration in a coarser unit (5 ms is 0 s). """ np_type = dtype.np_type try: @@ -375,6 +379,9 @@ def _converts_exactly(dtype, value): ): kind = converted_leaf.dtype.kind if kind not in "biufc": + for x, y in zip(given_leaf.ravel().tolist(), converted_leaf.ravel(), strict=True): + if _is_number(x) or not _converts_back(x, y): + return False continue for x, y in zip(given_leaf.ravel().tolist(), converted_leaf.ravel().tolist(), strict=True): if kind in "fc": @@ -386,6 +393,26 @@ def _converts_exactly(dtype, value): return True +def _is_number(x): + """A Python or numpy number or bool; not a timedelta, which numpy counts as an integer.""" + return isinstance(x, (numbers.Number, np.bool_)) and not isinstance(x, np.timedelta64) + + +def _converts_back(given, converted): + """Whether ``converted``, ``given`` as text, a date or a duration, converts back to ``given``. + + numpy compares dates and durations across units exactly, so 0 s is not 5 ms. + """ + given = np.asarray(given) + try: + back = np.asarray(converted).astype(given.dtype) + except (TypeError, ValueError): + return False + if given.dtype.kind in "mM" and np.isnat(given): + return bool(np.isnat(back)) + return bool(np.all(back == given)) + + def _object_layout(np_type): """Return ``np_type`` with every leaf an object, keeping field names, nesting and shapes.""" if np_type.names is not None: diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index c0badf8b4..c2d47a109 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4569,6 +4569,8 @@ def test_udt_literal_converts_only_when_exact(): np.dtype([("ex_v", np.int8, (2,)), ("ex_w", np.int8)]), "_ExactNested" ) text = dtypes.register_anonymous(np.dtype([("ex_s", "S4"), ("ex_n", np.int8)]), "_ExactText") + bare = dtypes.register_anonymous(np.dtype("S3"), "_ExactBareText") + seconds = dtypes.register_anonymous(np.dtype("m8[s]"), "_ExactBareSeconds") v = Vector(arr, size=3) v[0] = [1, 2, 3] w = Vector(arr, size=3) @@ -4581,8 +4583,14 @@ def test_udt_literal_converts_only_when_exact(): n[0] = ([1, 2], 3) t = Vector(text, size=1) t[0] = (b"ab", 1) + b = Vector(bare, size=1) + b[0] = b"0.5" + d = Vector(seconds, size=1) + d[0] = np.timedelta64(5, "s") + cut_text = np.array((b"abcdef", 1), dtype=[("ex_s", "S6"), ("ex_n", np.int64)])[()] udf = BinaryOp.register_anonymous(lambda x, y: x + y, "_exact_plus", is_udt=True) second = BinaryOp.register_anonymous(lambda x, y: y, "_exact_second", is_udt=True) + same = BinaryOp.register_anonymous(lambda x, y: x == y, "_exact_same", is_udt=True) # Literals of the UDT's type convert, a float rounding into a float32 field. assert v.apply(udf, right=(1, 2, 3)).new()[0].new().value.tolist() == [2, 4, 6] assert v.apply(udf, right=[True, 2, 3]).new().dtype == arr @@ -4614,6 +4622,13 @@ def test_udt_literal_converts_only_when_exact(): (lambda: f.apply(udf, right=np.complex128(1 + 2j)), "_ExactF32"), (lambda: f.apply(udf, right=np.float64(1e300)), "_ExactF32: it is typed as FP64,"), (lambda: n.apply(second, right=([1, 2.5], 3)), "_ExactNested"), + # A number does not become text exactly (0.5 would be b"0.5"), nor a string + # cut short, nor a duration in a coarser unit (5 ms is 0 s). + (lambda: t.apply(second, right=(1.5, 1)), "_ExactText"), + (lambda: b.apply(same, right=0.5), "_ExactBareText: it is typed as FP64,"), + (lambda: b.apply(same, right=np.True_), "_ExactBareText: it is typed as BOOL,"), + (lambda: t.apply(second, right=cut_text), "_ExactText"), + (lambda: d.apply(same, right=np.timedelta64(5, "ms")), "_ExactBareSeconds"), ]: with pytest.raises(ValueError, match=re.escape(f"does not fit {typed_as}")): expr() @@ -4634,6 +4649,16 @@ def test_udt_literal_converts_only_when_exact(): expr() with pytest.warns(DeprecationWarning, match="does not fit _ExactText"): assert t.apply(second, right=text_literal).new()[0].new().value.tolist() == (b"cd", 2) + with pytest.warns(DeprecationWarning, match="does not fit _ExactBareSeconds"): + assert d.apply(same, right=np.timedelta64(5000, "ms")).new()[0].new().value + with pytest.warns(DeprecationWarning, match="does not fit _ExactBareSeconds"): + # NaT stays NaT in any unit, and equals nothing. + assert not d.apply(same, right=np.timedelta64("NaT", "ms")).new()[0].new().value + # Text for a text leaf converts, bytes or str. + for literal in [b"cd", "cd"]: + assert t.apply(second, right=(literal, 1)).new()[0].new().value.tolist() == (b"cd", 1) + for literal in [b"0.5", "0.5"]: + assert b.apply(same, right=literal).new()[0].new().value with pytest.raises(OverflowError, match="300 out of bounds for int8"): v.apply(udf, right=(1, 2, 300)) # Defaults that fit still work, and fill in for the missing side. From 4fa58e61d205ed557530cf3f3e5064b65db94b14 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 06:13:27 -0500 Subject: [PATCH 26/54] Type every leaf of a record literal, numpy numbers, text and dates too Behavior fix, in #591's literal typing (older than its last commits; main had it too). A tuple or dict literal for a record UDT got no type at all when a leaf held anything but Python numbers: a numpy number, text or a date. Then nothing checked it, and numpy converted it into the record, losing value where a Python number alone raises. In an int8 record, under a user's op, (np.float64(1.5), 1) became (1, 1) while (1.5, 1) raised "does not fit"; under a lifted op, int8_rec + (np.float64(1.5), 1) computed with (1, 1), and int8_rec == (np.float64(1.5), 2) was True. A record of a bytes and an int8 field holding (b"cd", 1) had (b"cd", 1.5) == it True, and a user's op stored (b"cd", 1.5) as (b"cd", 1); a record with a date field the same. - main and 0e1ad03c: all of these as above, silently. - now every leaf is typed. A numpy number or bool is strong, as alone or beside a built-in type: the leaf takes its dtype, promoted with any Python numbers in it, and a numpy number may stand for a whole nested record. Text or a date in a text or date field takes the field's own type when it converts there and back unchanged, so ops still pair the field with itself (rec + (b"ab", 1) is (b"abab", 2), and (np.timedelta64(5, "m"), 1) adds 300 s to a field of seconds, as on main); otherwise it keeps its own type, which does not fit: text at its own length, a numpy date or duration in its unit. Other dates (datetime, str) take the field's type. So the user's op raises "does not fit" for np.float64(1.5) or 1.5 in an int8 field next to text or a date, for text longer than its field (it was cut short) and for a date or duration in a finer unit that loses value (5 ms in seconds); np.int64(2) converts with the DeprecationWarning (exact, but not a safe cast) and np.int16(2) fits an int16 field; the lifted op gives a float64 field (2.5, 3), and + with 5 ms in a field of seconds raises KeyError (main added 0 s); and the comparisons are False. In a comparison, an int beyond its integer leaf's range now becomes infinity in that leaf only, not in the whole literal, so a text leaf beside it still pairs: rec == (b"ab", 300) is False (it raised). A text date is compared as a date in the exactness check. A sequence for an array UDT keeps its rule, the array numpy makes of it. --- graphblas/core/literals.py | 116 ++++++++++++++++++++++++++++++------- graphblas/tests/test_op.py | 64 ++++++++++++++++++++ 2 files changed, 158 insertions(+), 22 deletions(-) diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index 243c34547..c6162c207 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -13,6 +13,7 @@ a UDT without fitting it is deprecated in :func:`_refuse_unless_exact` (MAINT). """ +import datetime import math import numbers import struct @@ -107,10 +108,9 @@ def _literal_operand(dtype, value, op): if is_comparison and not isinstance(value, (np.generic, np.ndarray)): np_type = dtype.np_type parts = _dict_to_record(np_type, value) if isinstance(value, dict) else value - beyond = _int_beyond_leaf(np_type, parts) - if beyond is not None: - # As infinity in every leaf, it still equals no element. - value = math.inf if beyond > 0 else -math.inf + if _int_beyond_leaf(np_type, parts) is not None: + # As infinity in its leaf, it still equals no element. + value = _beyond_as_infinity(np_type, parts) literal_type = _literal_type(dtype, value) return value, literal_type _check_literal_fits(dtype, value) @@ -398,16 +398,29 @@ def _is_number(x): return isinstance(x, (numbers.Number, np.bool_)) and not isinstance(x, np.timedelta64) +def _keeps_value(value, dtype): + """Whether ``value`` converts into ``dtype`` and back unchanged (:func:`_converts_back`).""" + try: + converted = np.array(value, dtype=dtype)[()] + except (TypeError, ValueError, OverflowError): + return False + return _converts_back(value, converted) + + def _converts_back(given, converted): """Whether ``converted``, ``given`` as text, a date or a duration, converts back to ``given``. numpy compares dates and durations across units exactly, so 0 s is not 5 ms. + A date given as text is compared as a date. """ - given = np.asarray(given) + given, converted = np.asarray(given), np.asarray(converted) try: - back = np.asarray(converted).astype(given.dtype) - except (TypeError, ValueError): - return False + back = converted.astype(given.dtype) + except (TypeError, ValueError, RuntimeError): + try: + back, given = converted, given.astype(converted.dtype) + except (TypeError, ValueError, RuntimeError): + return False if given.dtype.kind in "mM" and np.isnat(given): return bool(np.isnat(back)) return bool(np.all(back == given)) @@ -514,12 +527,13 @@ def _as_python_number(value): return value -def _literal_leaf_values(value, shape): +def _literal_leaf_values(value, shape, *, numpy_scalars=False): """Return the Python numbers in literal ``value`` read as an array of ``shape``, or ``None``. ``shape`` is ``()`` for a scalar leaf. A number fills any shape, as numpy broadcasts it. A sequence that does not fill the shape, or holds anything - that is not a Python number, gives ``None``. + that is not a Python number (or, with ``numpy_scalars``, a numpy number or + bool), gives ``None``. """ if shape: arr = np.empty(shape, dtype=object) @@ -530,7 +544,10 @@ def _literal_leaf_values(value, shape): values = [_as_python_number(x) for x in arr.ravel().tolist()] else: values = [_as_python_number(value)] - return values if all(type(x) in _LITERAL_KINDS for x in values) else None + numpy_types = (np.number, np.bool_) if numpy_scalars else () + if all(type(x) in _LITERAL_KINDS or isinstance(x, numpy_types) for x in values): + return values + return None def _weak_literal_udt(dtype, value, *, exact=False): @@ -547,24 +564,54 @@ def _weak_literal_udt(dtype, value, *, exact=False): number (bytes, datetime) takes the literal's default dtype and fails to pair. With ``exact``, for ``eq`` and ``ne``, an int out of a leaf's range takes a type that holds it instead of raising (:func:`_weak_literal_leaf`), - so ``int8_udt == 300`` is False. A sequence not laid out like an element - gives ``None``. + so ``int8_udt == 300`` is False. In a record, a numpy number or bool is + strong, as beside a built-in type: ``int8_record + (np.float64(1.5), 1)`` + has a float64 field. (A sequence for an array UDT is the array numpy makes + of it instead: :func:`_literal_type`.) A sequence not laid out like an + element gives ``None``. """ - def leaf_dtype(leaf, part): + def leaf_dtype(leaf, part, *, in_record=False): base, shape = leaf.subdtype or (leaf, ()) - values = _literal_leaf_values(part, shape) + if in_record and not shape and base.kind in "SUMm": + # Text or a date in such a field is typed too, so the number fields + # beside it are. It takes the field's own type when it converts into + # the field and back unchanged, so ops pair the field with itself; + # otherwise its own type, which does not fit: text as the field's + # kind at its own length, a numpy date or duration in its unit. + # Other dates (datetime, str) take the field's type. + if isinstance(part, (str, bytes)) and base.kind in "SU": + own = np.dtype(f"{base.kind}{max(len(part), 1)}") + elif isinstance(part, (np.datetime64, np.timedelta64)): + own = part.dtype + if own.kind != base.kind: + return own + elif isinstance(part, (str, bytes, datetime.date, datetime.timedelta)): + return base + else: + own = None + if own is not None: + return base if _keeps_value(part, base) else own + values = _literal_leaf_values(part, shape, numpy_scalars=in_record) if values is None: raise _LiteralDoesNotFit - kind = max((_LITERAL_KINDS[type(x)] for x in values), key="bifc".index) - if base.kind in "biufc": - elem = _weak_literal_leaf(values, kind, base, exact=exact) - else: - elem = _WEAK_LITERAL_DEFAULT[kind] + weak = [x for x in values if type(x) in _LITERAL_KINDS] + strong = [x.dtype for x in values if type(x) not in _LITERAL_KINDS] + elem = None + if weak: + kind = max((_LITERAL_KINDS[type(x)] for x in weak), key="bifc".index) + if base.kind in "biufc": + elem = _weak_literal_leaf(weak, kind, base, exact=exact) + else: + elem = _WEAK_LITERAL_DEFAULT[kind] + if strong: + elem = np.result_type(*strong, *([] if elem is None else [elem])) return np.dtype((elem, shape)) if shape else elem def record_leaves(np_type, part): - if type(_as_python_number(part)) in _LITERAL_KINDS: + if type(_as_python_number(part)) in _LITERAL_KINDS or isinstance( + part, (np.number, np.bool_) + ): parts = [part] * len(np_type.names) elif isinstance(part, (tuple, list)) and len(part) == len(np_type.names): parts = part @@ -575,7 +622,7 @@ def record_leaves(np_type, part): if field.names is not None: yield from record_leaves(field, item) else: - yield leaf_dtype(field, item) + yield leaf_dtype(field, item, in_record=True) np_type = dtype.np_type try: @@ -611,6 +658,31 @@ def _int_beyond_leaf(np_type, value): return None +def _beyond_as_infinity(np_type, value): + """Return literal ``value`` with each int beyond its integer leaf's range as infinity. + + ``value`` is laid out as :func:`_int_beyond_leaf` reads it. Infinity of the + int's sign equals no element, as the int does; every other value is kept, + so a text leaf still pairs with its field. + """ + if np_type.names is not None: + parts = value if isinstance(value, (tuple, list)) else [value] * len(np_type.names) + return tuple( + _beyond_as_infinity(np_type.fields[name][0], part) + for name, part in zip(np_type.names, parts, strict=True) + ) + base, shape = np_type.subdtype or (np_type, ()) + values = _literal_leaf_values(value, shape) if base.kind in "iu" else None + if values is None: + return value + lo, hi = _int_bounds(base) + values = [ + (math.inf if x > 0 else -math.inf) if type(x) is int and not lo <= x <= hi else x + for x in values + ] + return np.array(values, dtype=object).reshape(shape).tolist() if shape else values[0] + + def _array_literal_udt(dtype, value): """Return the UDT a numpy array ``value`` is beside UDT ``dtype``, or ``None`` to convert it. diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index c2d47a109..8db184d14 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4558,7 +4558,9 @@ def test_udt_literal_converts_only_when_exact(): float64; a numpy value is strong and must cast safely. A Monoid in ``apply`` is typed as its BinaryOp, so its literal promotes instead. """ + import datetime import re + import warnings arr = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_ExactArr") rec = dtypes.register_anonymous( @@ -4571,6 +4573,9 @@ def test_udt_literal_converts_only_when_exact(): text = dtypes.register_anonymous(np.dtype([("ex_s", "S4"), ("ex_n", np.int8)]), "_ExactText") bare = dtypes.register_anonymous(np.dtype("S3"), "_ExactBareText") seconds = dtypes.register_anonymous(np.dtype("m8[s]"), "_ExactBareSeconds") + dates = dtypes.register_anonymous( + np.dtype([("ex_d", "M8[s]"), ("ex_m", np.int8)]), "_ExactDate" + ) v = Vector(arr, size=3) v[0] = [1, 2, 3] w = Vector(arr, size=3) @@ -4588,6 +4593,8 @@ def test_udt_literal_converts_only_when_exact(): d = Vector(seconds, size=1) d[0] = np.timedelta64(5, "s") cut_text = np.array((b"abcdef", 1), dtype=[("ex_s", "S6"), ("ex_n", np.int64)])[()] + when = Vector(dates, size=1) + when[0] = (np.datetime64(0, "s"), 1) udf = BinaryOp.register_anonymous(lambda x, y: x + y, "_exact_plus", is_udt=True) second = BinaryOp.register_anonymous(lambda x, y: y, "_exact_second", is_udt=True) same = BinaryOp.register_anonymous(lambda x, y: x == y, "_exact_same", is_udt=True) @@ -4599,6 +4606,8 @@ def test_udt_literal_converts_only_when_exact(): assert r.apply(second, right=(1, 0.1)).new()[0].new().value.tolist() == (1, np.float32(0.1)) assert r.apply(second, right={"ex_i": 2, "ex_f": 0.5}).new()[0].new().value.tolist() == (2, 0.5) assert r.apply(second, right=np.int8(2)).new()[0].new().value.tolist() == (2, 2.0) + # A numpy number in a record is strong, as alone: np.int16 fits the int16 field. + assert r.apply(second, right=(np.int16(2), 0.5)).new()[0].new().value.tolist() == (2, 0.5) # A sequence holding numpy values is strong, as the array numpy makes of it # (int64, or int32 on Windows with numpy 1). made = np.asarray((np.int8(1), 2, 3)).dtype @@ -4610,6 +4619,18 @@ def test_udt_literal_converts_only_when_exact(): (lambda: v.apply(udf, left=np.array([0.5, 1, 1])), "_ExactArr"), (lambda: v.apply(udf, right=(np.float64(0.5), 1, 1)), "_ExactArr"), (lambda: r.apply(second, left=(0.5, 1)), "_ExactRec"), + (lambda: r.apply(second, right=(np.float64(1.5), 1)), "_ExactRec"), + (lambda: r.apply(second, right={"ex_i": np.float64(1.5), "ex_f": 1}), "_ExactRec"), + (lambda: n.apply(second, right=((np.float64(1.5), 2), 3)), "_ExactNested"), + # Text or a date in a record is typed too, so the number beside it is checked, + # and text longer than its field does not fit. + (lambda: t.apply(second, right=(b"ab", 1.5)), "_ExactText"), + (lambda: t.apply(second, right=(b"ab", np.float64(1.5))), "_ExactText"), + (lambda: t.apply(second, right=(b"abcdef", 1)), "_ExactText"), + (lambda: when.apply(second, right=(np.datetime64(0, "s"), 1.5)), "_ExactDate"), + (lambda: when.apply(second, right=(datetime.date(1970, 1, 1), 1.5)), "_ExactDate"), + (lambda: when.apply(second, right=("1970-01-01", 1.5)), "_ExactDate"), + (lambda: when.apply(second, right=(np.datetime64(5, "ms"), 1)), "_ExactDate"), (lambda: v.ewise_union(w, udf, 0.5, 0), "_ExactArr"), (lambda: v.ewise_union(w, udf, 0, 0.5), "_ExactArr"), ( @@ -4642,6 +4663,7 @@ def test_udt_literal_converts_only_when_exact(): (lambda: v.apply(udf, right=np.int64(1)), "_ExactArr: it is typed as INT64,"), (lambda: v.apply(udf, right=(np.int8(1), 2, 3)), "_ExactArr"), (lambda: r.apply(second, right=np.int32(2)), "_ExactRec: it is typed as INT32,"), + (lambda: r.apply(second, right=(np.int64(2), 0.5)), "_ExactRec"), (lambda: n.apply(second, right=([1, 2.0], 3)), "_ExactNested"), (lambda: t.apply(second, right=text_literal), "_ExactText"), ]: @@ -4657,6 +4679,20 @@ def test_udt_literal_converts_only_when_exact(): # Text for a text leaf converts, bytes or str. for literal in [b"cd", "cd"]: assert t.apply(second, right=(literal, 1)).new()[0].new().value.tolist() == (b"cd", 1) + assert t.apply(second, right=(b"c", 1)).new()[0].new().value.tolist() == (b"c", 1) + # Text or a date that converts into its field and back fits it, with no warning: + # trailing NULs, a finer date that is whole seconds, NaT. + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert t.apply(second, right=(b"ab\0\0\0", 1)).new()[0].new().value.tolist() == (b"ab", 1) + fine = when.apply(second, right=(np.datetime64(5000, "ms"), 1)).new()[0].new().value + assert fine["ex_d"] == np.datetime64(5, "s") + nat = when.apply(second, right=(np.datetime64("NaT", "ms"), 1)).new()[0].new().value + assert np.isnat(nat["ex_d"]) + assert (t == (b"ab", 1)).new()[0].new().value + assert not (t == (b"ab", 1.5)).new()[0].new().value + assert (when == (np.datetime64(0, "s"), 1)).new()[0].new().value + assert not (when == (np.datetime64(0, "s"), 1.5)).new()[0].new().value for literal in [b"0.5", "0.5"]: assert b.apply(same, right=literal).new()[0].new().value with pytest.raises(OverflowError, match="300 out of bounds for int8"): @@ -4756,6 +4792,34 @@ def test_udt_sequence_and_array_literals_type_like_numbers(): dtype=[("sq_a", np.int8), ("sq_in", [("sq_b", np.float64), ("sq_c", np.int8)])], ) assert (n + strong).new()[0].new().value.tolist() == (2, (2.5, 4.5)) + # A numpy number in a record literal keeps its dtype, as a 0-d array does. + result = binary.plus(n, (np.float64(0.5), (1, 1))).new() + assert result.dtype.np_type["sq_a"] == np.float64 + assert result[0].new().value.tolist() == (1.5, (3, 4.5)) + assert (n == (np.float64(1), (2, 3.5))).new()[0].new().value + assert not (n == (np.float64(1.5), (2, 3.5))).new()[0].new().value + result = binary.plus(n, (1, np.float64(1.5))).new() # one number for the nested record + assert result.dtype.np_type["sq_in"]["sq_b"] == np.float64 + assert result[0].new().value.tolist() == (2, (3.5, 5.0)) + # Text or a duration that converts into its field and back takes the field's type, + # so a lifted op pairs the record with itself; one that does not keeps its own. + text = dtypes.register_anonymous(np.dtype([("sq_tx", "S4"), ("sq_tn", np.int8)]), "_SeqText") + tv = Vector(text, size=1) + tv[0] = (b"ab", 1) + assert binary.plus(tv, (b"ab", 1)).new()[0].new().value.tolist() == (b"abab", 2) + assert tv.apply(binary.max, right=(b"ac", 0)).new()[0].new().value.tolist() == (b"ac", 1) + assert tv.apply(binary.min, right=(b"ac", 0)).new()[0].new().value.tolist() == (b"ab", 0) + assert not (tv == (b"ab", 300)).new()[0].new().value # 300 is no int8 + secs = dtypes.register_anonymous( + np.dtype([("sq_sd", "m8[s]"), ("sq_sn", np.int8)]), "_SeqSeconds" + ) + sv = Vector(secs, size=1) + sv[0] = (np.timedelta64(10, "s"), 1) + for delta, total in [(np.timedelta64(5, "m"), 310), (np.timedelta64(5000, "ms"), 15)]: + result = binary.plus(sv, (delta, 1)).new() + assert result[0].new().value["sq_sd"] == np.timedelta64(total, "s") + with pytest.raises(KeyError, match="no common type"): + binary.plus(sv, (np.timedelta64(5, "ms"), 1)) # 5 ms is no whole number of seconds first = dtypes.register_anonymous( np.dtype([("sq_fin", [("sq_fb", np.int8), ("sq_fc", np.float32)]), ("sq_fa", np.int16)]) ) From 31eccfdc906665c3953480a6066b26ea85ba7567 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 13:07:45 -0500 Subject: [PATCH 27/54] Refuse UDT ops on float16 elements by name; compare them in numpy for Scalar == Behavior fix. Numba has a float16 type but no code for it, so every op compiled on a UDT with a float16 leaf (a field, an array's element, or a bare float16 UDT) failed with a bare NotImplementedError: float16. - A record with a float16 field, + (1.0,) or == (1.5,): NotImplementedError on main, at 0e1ad03c and before this commit; now KeyError. - Its Scalar == (1.5,): True on main, NotImplementedError at 0e1ad03c and before; now True again. - A bare float16 UDT, Vector == or isequal: TypeError on main and at 0e1ad03c, NotImplementedError before; now KeyError. - An int8 record + (np.float16(1.5), 1) or == (np.float16(1.5), 2): main and 0e1ad03c cut 1.5 to 1 and answered (2, 3) and True, both wrong; before, KeyError ("no common type") and NotImplementedError; now KeyError. Every UDT op's _compile_udt (unary, binary, indexunary, select, indexbinary) now refuses a float16 leaf up front with the KeyError an op gives for a type it does not take: "binary.eq does not work with (...): Numba cannot compute on float16, which ... holds." A Scalar's == on a UDT with a float16 leaf compares in numpy, as for a UDT that is neither a record nor an array. A store that would cast to or from a UDT with a float16 leaf (another record or array UDT, or dup(dtype=...)), which raised the same NotImplementedError on main, at 0e1ad03c and before, now raises the DomainMismatch of a store that cannot cast, saying why: "... Numba cannot compute on float16, which ... holds." Widening float16 to float32 was not done: the result type would change, and results stored back as float16 would round. --- graphblas/core/operator/binary.py | 3 +++ graphblas/core/operator/indexbinary.py | 3 +++ graphblas/core/operator/indexunary.py | 3 +++ graphblas/core/operator/select.py | 3 +++ graphblas/core/operator/unary.py | 3 +++ graphblas/core/scalar.py | 10 +++++--- graphblas/core/udt.py | 35 ++++++++++++++++++++++++-- graphblas/tests/test_op.py | 29 +++++++++++++++++++++ 8 files changed, 84 insertions(+), 5 deletions(-) diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index c592fac98..1f4d4db38 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -625,6 +625,9 @@ def _compile_udt(self, dtype, dtype2): dtypes = (dtype, dtype2) if dtypes in self._udt_types: return self._udt_ops[dtypes] + from ..udt import _refuse_half + + _refuse_half(self, dtype, dtype2) # Built-in ops lift to UDTs by name. A UDF registered under one of those # names (``BinaryOp.register_anonymous(func, "plus", is_udt=True)``) is diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index c07968881..c987a64e6 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -374,6 +374,9 @@ def _compile_udt(self, dtype, dtype2): dtypes = (dtype, dtype2) if dtypes in self._udt_types: return self._udt_ops[dtypes] + from ..udt import _refuse_half + + _refuse_half(self, dtype, dtype2) if self._numba_func is None: raise KeyError(f"{self.name} does not work with {dtypes} types") diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index 25eecc0b0..0a7437def 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -207,6 +207,9 @@ def _compile_udt(self, dtype, dtype2): dtypes = (dtype, dtype2) if dtypes in self._udt_types: return self._udt_ops[dtypes] + from ..udt import _refuse_half + + _refuse_half(self, dtype, dtype2) if self._numba_func is None: raise KeyError(f"{self.name} does not work with {dtypes} types") diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 0fca32ad4..18e001fcc 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -259,6 +259,9 @@ def _compile_udt(self, dtype, dtype2): dtypes = (dtype, dtype2) if dtypes in self._udt_types: return self._udt_ops[dtypes] + from ..udt import _refuse_half + + _refuse_half(self, dtype, dtype2) if self._numba_func is None: raise KeyError(f"{self.name} does not work with {dtypes} types") diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index bffedeedd..c0bdf0434 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -256,6 +256,9 @@ def unary_wrapper(z, x): def _compile_udt(self, dtype, dtype2): if dtype in self._udt_types: return self._udt_ops[dtype] + from ..udt import _refuse_half + + _refuse_half(self, dtype) # See the matching note in ``BinaryOp._compile_udt``: plain-scalar # misses on UDT-dispatch-enabled ops get the legacy message. if not dtype._is_udt and self._numba_func is None: diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 84c9fb946..c44a97c40 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -23,6 +23,7 @@ _element_value, _ewise_add_needs_cast, _get_udt_info, + _has_half_leaf, _is_lifted_op, _literal_shape_differs, _scalar_ewise_add_converted, @@ -1166,10 +1167,13 @@ def _scalar_eq(scalar, other): return (convert(scalar.value) == convert(value)).item() udts = [scalar.dtype, other.dtype] if type(other) is Scalar else [scalar.dtype] udts = [dtype for dtype in udts if dtype._is_udt] - if udts and (not _has_numba or any(_get_udt_info(dtype) is None for dtype in udts)): + if udts and ( + not _has_numba + or any(_get_udt_info(dtype) is None or _has_half_leaf(dtype.np_type) for dtype in udts) + ): # binary.eq needs Numba. A UDT that is neither a record nor an array (bytes, - # float16) holds one value, which numpy compares as binary.eq does, and with - # no Numba type for it (float16 has none) and no GraphBLAS call. + # float16) holds one value, which numpy compares as binary.eq does, with no + # GraphBLAS call; and numpy compares float16, which Numba cannot. return _udt_eq_without_numba(scalar, other) return bool(scalar.ewise_mult(other, binary.eq).new(name="s_eq").value) diff --git a/graphblas/core/udt.py b/graphblas/core/udt.py index 13acf752c..d4befb350 100644 --- a/graphblas/core/udt.py +++ b/graphblas/core/udt.py @@ -341,13 +341,18 @@ def _corresponding_leaves(np_x, np_z): def _udt_cast_error(dtype, to_dtype): """Return why an element of UDT ``dtype`` cannot be stored as UDT ``to_dtype``, or ``None``. - The layouts must correspond (:func:`_corresponding_leaves`), and each pair - of elements must be built-in numbers, or the same type. + The layouts must correspond (:func:`_corresponding_leaves`), each pair of + elements must be built-in numbers, or the same type, and none float16, + which the Numba code that casts cannot compute on. """ reason, pairs = _corresponding_leaves(dtype.np_type, to_dtype.np_type) for x, z in pairs or (): if x != z and not (x.kind in "biufc" and z.kind in "biufc"): return f"elements of {x} do not cast to {z}" + if reason is None: + for udt in (dtype, to_dtype): + if _has_half_leaf(udt.np_type): + return f"Numba cannot compute on float16, which {udt} holds" return reason @@ -745,6 +750,32 @@ def _leaf_arrays(arr): yield from _leaf_arrays(arr[name]) +@lru_cache +def _has_half_leaf(np_type): + """Whether ``np_type`` has a float16 leaf: a record field, an array's element, or itself.""" + if np_type.subdtype is not None: + return _has_half_leaf(np_type.subdtype[0]) + if np_type.names is not None: + return any(_has_half_leaf(np_type.fields[name][0]) for name in np_type.names) + return np_type.kind == "f" and np_type.itemsize == 2 + + +def _refuse_half(op, *dtypes): + """Raise ``KeyError`` if a UDT among ``dtypes`` has a float16 leaf, which ``op`` cannot take. + + Numba has a float16 type but no code for it: compiling an op on one raises + a bare ``NotImplementedError``. So such an op is refused up front, as an op + that does not take a type is. + """ + for dtype in dtypes: + if dtype is not None and dtype._is_udt and _has_half_leaf(dtype.np_type): + names = ", ".join(str(d) for d in dtypes if d is not None) + raise KeyError( + f"{op._modname}.{op.name} does not work with ({names}): Numba cannot " + f"compute on float16, which {dtype} holds." + ) + + @lru_cache def _has_narrow_float_leaf(np_type): """Whether ``np_type`` has a float leaf narrower than float64, which a float may overflow.""" diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 8db184d14..5622428ec 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4798,6 +4798,13 @@ def test_udt_sequence_and_array_literals_type_like_numbers(): assert result[0].new().value.tolist() == (1.5, (3, 4.5)) assert (n == (np.float64(1), (2, 3.5))).new()[0].new().value assert not (n == (np.float64(1.5), (2, 3.5))).new()[0].new().value + # Numba cannot compute on float16, so a float16 leaf is refused by name. + for expr in [ + lambda: binary.plus(n, (np.float16(0.5), (1, 1))), + lambda: n == (np.float16(1), (2, 3.5)), + ]: + with pytest.raises(KeyError, match="Numba cannot compute on float16"): + expr() result = binary.plus(n, (1, np.float64(1.5))).new() # one number for the nested record assert result.dtype.np_type["sq_in"]["sq_b"] == np.float64 assert result[0].new().value.tolist() == (2, (3.5, 5.0)) @@ -5924,6 +5931,28 @@ def test_udt_eq_compares_a_udt_that_is_neither_record_nor_array(): assert h == gb.Scalar.from_value(1.5, half) assert h == 1.5 assert gb.Scalar.from_value(np.nan, half) != gb.Scalar.from_value(np.nan, half) + # Numba cannot compute on float16, so an op on it is refused by name, for a + # record with a float16 field too; a Scalar's == compares in numpy. + halves = Vector(half, size=1) + halves[0] = 1.5 + half_rec = dtypes.register_anonymous(np.dtype([("eqo_h", np.float16)]), "_EqOpaqueHalfRec") + assert gb.Scalar.from_value((1.5,), half_rec) == (1.5,) + for expr in [ + lambda: halves == halves, + lambda: halves.isequal(halves), + lambda: binary.plus(Vector(half_rec, size=1), (1.0,)), + lambda: -Vector(half_rec, size=1), + ]: + with pytest.raises(KeyError, match="Numba cannot compute on float16"): + expr() + # A store cast is Numba code too, so it raises the DomainMismatch of a store that cannot cast. + single_rec = dtypes.register_anonymous(np.dtype([("eqo_h", np.float32)]), "_EqOpaqueSingleRec") + for expr in [ + lambda: Vector(single_rec, size=1) << Vector(half_rec, size=1), + lambda: Vector(half_rec, size=1).dup(dtype=single_rec), + ]: + with pytest.raises(DomainMismatch, match="Numba cannot compute on float16"): + expr() dates = dtypes.register_anonymous(np.dtype("M8[s]"), "_EqOpaqueDates") d = Vector(dates, size=2) d[0] = np.datetime64("NaT", "s") From 6ab263ef250de71d80d311ef66c21319104aff17 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 13:10:13 -0500 Subject: [PATCH 28/54] Type a record's array field as an array UDT types its literal Behavior fix. A sequence for an array UDT is typed by the rule #591 gives it: Python numbers alone are weak, each taking the element's type when its kind holds them, and a numpy array, or a sequence holding a numpy value, is the array numpy makes of it. The same sequence for an array field of a record was typed element by element, numpy values strong and Python numbers weak beside the field, and a numpy array was read as Python numbers. So, with an int8[3] array UDT and a record with an int8[3] field: - (np.int8(1), 300, 3): the array UDT gives int64[3] and computes [2, 302, 6]; the record raised OverflowError (main and 0e1ad03c too). - np.array([1, 300, 3]): the array UDT gives int64[3]; the record raised OverflowError (0e1ad03c), and main wrapped 300 to 44 silently. - (np.float64(1.5), 2, 3): main and 0e1ad03c cut 1.5 to 1; the commit "Type every leaf of a record literal, numpy numbers, text and dates too" already gave a float64 field. A record's array field now takes the array UDT's rule: an array, or a sequence holding a numpy number or bool, is the array numpy makes of it, broadcast to the field's shape. So the record gives an int64[3] field and [2, 302, 6], as the array UDT does. Python numbers alone, and a numpy number for the whole field, are typed as before. A comparison keeps a Python int beyond the element's range exact, as it does everywhere (int8_vec == 300 is False), also inside such an array: the int is infinity in its place before numpy makes the array. So rec == ((np.int8(1), 2**70, 3), 4) and int8_array_udt == (np.int8(1), 2**70, 3) are False and != True. The array UDT raised numpy's OverflowError there (older than this commit), and the record field, False before, would have raised it under the array rule alone. A user's op with such a literal still raises OverflowError. --- graphblas/core/literals.py | 60 +++++++++++++++++++++++++++----------- graphblas/tests/test_op.py | 18 ++++++++++++ 2 files changed, 61 insertions(+), 17 deletions(-) diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index c6162c207..7f85ad6b8 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -103,15 +103,15 @@ def _literal_operand(dtype, value, op): # truediv gives float leaves whatever the int, so, as a comparison, it # types an int out of an integer leaf's range exactly instead of raising. exact = is_comparison or op.name == "truediv" + if is_comparison and not isinstance(value, (np.generic, np.ndarray)): + np_type = dtype.np_type + parts = _dict_to_record(np_type, value) if isinstance(value, dict) else value + if _int_beyond_leaf(np_type, parts) is not None: + # As infinity in its leaf, it still equals no element, and it + # stays out of an array numpy makes of numpy values beside it. + value = _beyond_as_infinity(np_type, parts) literal_type = _literal_type(dtype, value, exact=exact) if literal_type is not None: - if is_comparison and not isinstance(value, (np.generic, np.ndarray)): - np_type = dtype.np_type - parts = _dict_to_record(np_type, value) if isinstance(value, dict) else value - if _int_beyond_leaf(np_type, parts) is not None: - # As infinity in its leaf, it still equals no element. - value = _beyond_as_infinity(np_type, parts) - literal_type = _literal_type(dtype, value) return value, literal_type _check_literal_fits(dtype, value) return value, dtype @@ -592,20 +592,29 @@ def leaf_dtype(leaf, part, *, in_record=False): own = None if own is not None: return base if _keeps_value(part, base) else own + if in_record and shape and _holds_numpy_values(part): + # As for an array UDT, an array, or a sequence holding numpy values, is + # the array numpy makes of it (strong), in the field's shape. + try: + arr = np.asarray(part) + np.broadcast_to(arr, shape) + except (TypeError, ValueError): + raise _LiteralDoesNotFit from None + if arr.dtype.kind not in "biufc": + raise _LiteralDoesNotFit + return np.dtype((arr.dtype, shape)) values = _literal_leaf_values(part, shape, numpy_scalars=in_record) if values is None: raise _LiteralDoesNotFit - weak = [x for x in values if type(x) in _LITERAL_KINDS] - strong = [x.dtype for x in values if type(x) not in _LITERAL_KINDS] - elem = None - if weak: - kind = max((_LITERAL_KINDS[type(x)] for x in weak), key="bifc".index) + if type(values[0]) not in _LITERAL_KINDS: + # A numpy number or bool is strong: its dtype, in every place of the field. + elem = values[0].dtype + else: + kind = max((_LITERAL_KINDS[type(x)] for x in values), key="bifc".index) if base.kind in "biufc": - elem = _weak_literal_leaf(weak, kind, base, exact=exact) + elem = _weak_literal_leaf(values, kind, base, exact=exact) else: elem = _WEAK_LITERAL_DEFAULT[kind] - if strong: - elem = np.result_type(*strong, *([] if elem is None else [elem])) return np.dtype((elem, shape)) if shape else elem def record_leaves(np_type, part): @@ -635,6 +644,19 @@ def record_leaves(np_type, part): return dtype if literal_type == np_type else lookup_dtype(literal_type) +def _holds_numpy_values(value): + """Whether ``value`` is a numpy array, or a sequence holding a numpy number or bool.""" + if isinstance(value, np.ndarray): + return True + if not isinstance(value, (tuple, list)): + return False + try: + items = np.asarray(value, dtype=object).ravel().tolist() + except (TypeError, ValueError): # ragged + return False + return any(isinstance(x, (np.number, np.bool_)) for x in items) + + def _int_beyond_leaf(np_type, value): """Return a Python int in literal ``value`` beyond its integer leaf's range, or ``None``. @@ -644,6 +666,8 @@ def _int_beyond_leaf(np_type, value): """ if np_type.names is not None: parts = value if isinstance(value, (tuple, list)) else [value] * len(np_type.names) + if len(parts) != len(np_type.names): + return None # not laid out like an element; typing it reports that for name, part in zip(np_type.names, parts, strict=True): found = _int_beyond_leaf(np_type.fields[name][0], part) if found is not None: @@ -652,7 +676,7 @@ def _int_beyond_leaf(np_type, value): base, shape = np_type.subdtype or (np_type, ()) if base.kind in "iu": lo, hi = _int_bounds(base) - for x in _literal_leaf_values(value, shape) or (): + for x in _literal_leaf_values(value, shape, numpy_scalars=True) or (): if type(x) is int and not lo <= x <= hi: return x return None @@ -667,12 +691,14 @@ def _beyond_as_infinity(np_type, value): """ if np_type.names is not None: parts = value if isinstance(value, (tuple, list)) else [value] * len(np_type.names) + if len(parts) != len(np_type.names): + return value return tuple( _beyond_as_infinity(np_type.fields[name][0], part) for name, part in zip(np_type.names, parts, strict=True) ) base, shape = np_type.subdtype or (np_type, ()) - values = _literal_leaf_values(value, shape) if base.kind in "iu" else None + values = _literal_leaf_values(value, shape, numpy_scalars=True) if base.kind in "iu" else None if values is None: return value lo, hi = _int_bounds(base) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 5622428ec..60ac4ea7a 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4805,6 +4805,24 @@ def test_udt_sequence_and_array_literals_type_like_numbers(): ]: with pytest.raises(KeyError, match="Numba cannot compute on float16"): expr() + # A record's array field takes a sequence as an array UDT does: an array, or a + # sequence holding a numpy value, is the array numpy makes of it. + arr_rec = dtypes.register_anonymous( + np.dtype([("sq_av", np.int8, (3,)), ("sq_an", np.int8)]), "_SeqArrRec" + ) + ar = Vector(arr_rec, size=1) + ar[0] = ([1, 2, 3], 4) + for literal in [(np.int8(1), 300, 3), np.array([1, 300, 3])]: + made = np.dtype((np.asarray(literal).dtype, (3,))) + assert binary.plus(v, literal).new().dtype.np_type == made + result = binary.plus(ar, (literal, 1)).new() + assert result.dtype.np_type["sq_av"] == made + assert result[0].new().value["sq_av"].tolist() == [2, 302, 6] + # A comparison keeps a Python int beyond int8 exact, there too: it equals nothing. + big = (np.int8(1), 2**70, 3) + for x, literal in [(ar, (big, 4)), (v, big)]: + assert not (x == literal).new()[0].new().value + assert (x != literal).new()[0].new().value result = binary.plus(n, (1, np.float64(1.5))).new() # one number for the nested record assert result.dtype.np_type["sq_in"]["sq_b"] == np.float64 assert result[0].new().value.tolist() == (2, (3.5, 5.0)) From fbf7ba71235c75f5f492a273ff708a2a48bb4180 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 13:14:22 -0500 Subject: [PATCH 29/54] Refuse == between text, dates and numbers with eq's KeyError Behavior fix. binary.eq and ne compared each pair of record leaves with Numba, which cannot compare bytes with a number, so such a pair raised UdfParseError carrying Numba's typing error: - a record with a bytes and a float64 field, eq with a record of two float64 fields: GrB_DOMAIN_MISMATCH on main, UdfParseError at 0e1ad03c and before this commit; - the same record == 1 or != (1.0, 1.0), and its Scalar == (1.0, 1.0): main converted the number into the bytes field (1 became b"1") and compared the converted value; UdfParseError at 0e1ad03c and before; - a nested record == (b"x", np.float64(1)), whose number stands for a nested record holding a bytes field: False on main and at 0e1ad03c, UdfParseError before this commit. Text, a date, a duration and a number now compare only with their own kind; any other pair of leaves raises the KeyError binary.eq gives for a pair it does not take: "binary.eq does not work with (...): elements of |S4 and float64 do not compare." Scalar == raises it with or without Numba. This follows the rules #591 already has: plus finds no common type for such leaves (a KeyError), and a UDT that is neither a record nor an array pairs only with its own type. Answering False instead (a number for a text leaf can never be equal) would turn the record pair, an error since main, into an answer, and would differ from main anyway where main compared b"1" with 1 as equal. --- graphblas/core/operator/binary.py | 8 ++++++- graphblas/core/udt.py | 35 +++++++++++++++++++++++++++++++ graphblas/tests/test_op.py | 29 ++++++++++++++++++------- graphblas/tests/test_scalar.py | 2 ++ 4 files changed, 65 insertions(+), 9 deletions(-) diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 1f4d4db38..8ea60ad99 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -65,13 +65,19 @@ def _make_udt_comparison(dtype, dtype2, *, is_eq): neither a record nor an array (bytes, say) compares only with its own type, as one value. Returns ``(wrapper_func, wrapper_sig)``. """ - from ..udt import _check_udt_pair, _get_udt_info, _iter_record_leaves + from ..udt import ( + _check_leaves_compare, + _check_udt_pair, + _get_udt_info, + _iter_record_leaves, + ) from .udt_codegen import _array_operand_source info_x = _get_udt_info(dtype) info_y = _get_udt_info(dtype2) op_name = "eq" if is_eq else "ne" _check_udt_pair(op_name, dtype, dtype2, info_x, info_y) + _check_leaves_compare(op_name, dtype, dtype2) nt = numba.types cmp_op = "==" if is_eq else "!=" diff --git a/graphblas/core/udt.py b/graphblas/core/udt.py index d4befb350..dd5009fbc 100644 --- a/graphblas/core/udt.py +++ b/graphblas/core/udt.py @@ -303,6 +303,40 @@ def _promoted_record(records, result_leaves): return np.dtype(fields, align=any(r.isalignedstruct for r in records)) +def _value_kind(np_type): + """What a leaf of ``np_type`` holds, as ``eq`` compares it: a number, text, or its own kind.""" + if np_type.kind in "biufc": + return "number" + return "text" if np_type.kind in "SU" else np_type.kind + + +def _check_leaves_compare(op_name, dtype, dtype2): + """Raise ``KeyError`` unless each pair of leaves that ``eq`` compares holds one kind of value. + + Text, a date, a duration and a number compare only with their own kind, as + ``plus`` finds no common type for any other pair. ``dtype`` and ``dtype2`` + already pair (:func:`_check_udt_pair`); a side that is not a record stands + in every leaf of the other. + """ + + def leaves(np_type): + if np_type.names is not None: + return [leaf.base for _py, _c, leaf in _iter_record_leaves(np_type)] + return [np_type.base] + + xs, ys = leaves(dtype.np_type), leaves(dtype2.np_type) + if len(xs) == 1: + xs = xs * len(ys) + if len(ys) == 1: + ys = ys * len(xs) + for x, y in zip(xs, ys, strict=True): + if _value_kind(x) != _value_kind(y): + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): elements of " + f"{x} and {y} do not compare." + ) + + # Storing a value of one type in an object of another. @@ -679,6 +713,7 @@ def _udt_eq_without_numba(scalar, other): # field names, array shapes that do not broadcast, a record and an array). info = _get_udt_info(scalar.dtype), _get_udt_info(other.dtype) _check_udt_pair("eq", scalar.dtype, other.dtype, *info) + _check_leaves_compare("eq", scalar.dtype, other.dtype) return _leaves_equal(scalar.value, scalar.dtype.np_type, other.value, other.dtype.np_type) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 60ac4ea7a..2f548368e 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -5884,13 +5884,13 @@ def test_udt_jit_kernel_only_for_a_same_type_result(): @pytest.mark.skipif("not supports_udfs") -def test_udt_eq_compile_failure_is_a_udfparseerror(): - """``eq`` reports a pair the codegen cannot type as a UdfParseError, as ``plus`` does. +def test_udt_eq_on_leaves_of_other_kinds_is_a_keyerror(): + """``eq`` refuses leaves that hold different kinds of value with a KeyError, as ``plus`` does. - Its cfunc was compiled without the wrapper that turns Numba's errors into - a one-line UdfParseError, so a bytes field against a float field surfaced - as Numba's full TypingError. ``plus`` rejects the pair sooner: bytes and - float64 have no common dtype for the result field. + Text, a date, a duration and a number compare only with their own kind, so a + bytes field against a float field, or a number given for a bytes field, is + binary.eq's KeyError, as it is for ``plus``: bytes and float64 have no + common dtype for the result field. """ text = dtypes.register_anonymous(np.dtype([("ueq_a", "S4"), ("ueq_b", np.float64)]), "_UeqS") num = dtypes.register_anonymous( @@ -5900,13 +5900,26 @@ def test_udt_eq_compile_failure_is_a_udfparseerror(): v[0] = (b"ab", 1.0) w = Vector(num, size=1) w[0] = (1.0, 1.0) - with pytest.raises(UdfParseError, match="binary.eq does not work with"): - binary.eq(v & w).new() with pytest.raises(KeyError, match=r"elements of \|S4 and float64 have no common type"): binary.plus(v & w).new() # A number beside the bytes field is an int64 there, which pairs no better. with pytest.raises(KeyError, match=r"elements of \|S4 and int64 have no common type"): v + 1 + nested = dtypes.register_anonymous( + np.dtype([("ueq_x", "S2"), ("ueq_in", [("ueq_y", "S2"), ("ueq_n", np.int8)])]), "_UeqN" + ) + n = Vector(nested, size=1) + n[0] = (b"x", (b"y", 1)) + for expr in [ + lambda: binary.eq(v & w), + lambda: v == 1, + lambda: v != (1.0, 1.0), + lambda: gb.Scalar.from_value((b"ab", 1.0), text) == (1.0, 1.0), + lambda: n == (b"x", np.float64(1)), + ]: + with pytest.raises(KeyError, match=r"elements of \|S\d and \w+ do not compare"): + expr() + assert (n == (b"x", (b"y", 1))).new()[0].new().value @pytest.mark.skipif("not supports_udfs") diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index cfb1361c6..0845be5c0 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -754,6 +754,8 @@ def test_udt_eq_without_numba_matches_binary_eq(monkeypatch): (r, a), (t, Scalar.from_value(b"ab", dtypes.register_anonymous(np.dtype("S7")))), (t, 1), + # A number for a bytes field does not compare. + (Scalar.from_value((b"ab", 1), np.dtype([("tx", "S2"), ("tn", np.int8)])), (1, 1)), ] def compare_all(): From f8b10021deb2c93cd57affddbaee0123c2b61232 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 13:15:58 -0500 Subject: [PATCH 30/54] Test the text-date and duration-for-a-date branches of record literal typing Tests only. Two branches of the record literal typing had no row: - a date given as text shorter than numpy's own date string ("1970-01-01" for a field of seconds), which _converts_back compares as a date: with a number that converts exactly, the literal converts with the DeprecationWarning, and without the branch it raised "does not fit"; - a numpy duration given for a date field, which keeps its own type: numpy converts np.timedelta64(5, "s") to a date by value, so without the guard it took the date field's type and failed in numpy's own conversion; with it, it raises "does not fit". --- 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 2f548368e..8bd571160 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4631,6 +4631,7 @@ def test_udt_literal_converts_only_when_exact(): (lambda: when.apply(second, right=(datetime.date(1970, 1, 1), 1.5)), "_ExactDate"), (lambda: when.apply(second, right=("1970-01-01", 1.5)), "_ExactDate"), (lambda: when.apply(second, right=(np.datetime64(5, "ms"), 1)), "_ExactDate"), + (lambda: when.apply(second, right=(np.timedelta64(5, "s"), 1)), "_ExactDate"), (lambda: v.ewise_union(w, udf, 0.5, 0), "_ExactArr"), (lambda: v.ewise_union(w, udf, 0, 0.5), "_ExactArr"), ( @@ -4666,6 +4667,7 @@ def test_udt_literal_converts_only_when_exact(): (lambda: r.apply(second, right=(np.int64(2), 0.5)), "_ExactRec"), (lambda: n.apply(second, right=([1, 2.0], 3)), "_ExactNested"), (lambda: t.apply(second, right=text_literal), "_ExactText"), + (lambda: when.apply(second, right=("1970-01-01", 2.0)), "_ExactDate"), # a date as text ]: with pytest.warns(DeprecationWarning, match=re.escape(f"does not fit {typed_as}")): expr() From 44abf82d0490b546aea29e94da8a32ac8e714466 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 13:18:21 -0500 Subject: [PATCH 31/54] Let a UDF return an operand of a UDT that is neither a record nor an array Behavior fix. A user's op on a UDT of bytes, str or the like that returns an operand as it is (lambda x, y: x) failed to compile: Numba types the element as a char sequence, which no DataType is registered for, so the return type was not found. - main: UdfParseError "UDT UDF returned an unsupported type [char x 8]". - 0e1ad03c: TypeError 'NoneType' object is not subscriptable (the literal beside it was read as an array UDT). - before this commit: the UdfParseError again. - now: the return takes the operand's UDT when its Numba type is that UDT's, so the op computes; [b"abc"] for x, the literal for y. --- graphblas/core/operator/udf.py | 10 ++++++++++ graphblas/tests/test_op.py | 5 +++++ 2 files changed, 15 insertions(+) diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index 55a78adfe..9eeae225d 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -153,6 +153,16 @@ def _resolve_udt_return_type(numba_ret_type, *dtypes): return lookup_dtype(numba_ret_type) except (ValueError, TypeError): pass + for d in dtypes: + # Bytes and the like have a Numba type no DataType is registered for; + # an operand of that type, returned as it is, keeps its UDT. + if ( + d._is_udt + and d.np_type.names is None + and d.np_type.subdtype is None + and d.numba_type == numba_ret_type + ): + return d if isinstance(numba_ret_type, numba.core.types.BaseTuple): n = len(numba_ret_type.types) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 8bd571160..dce4589bd 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2089,6 +2089,11 @@ def _pair(x, y): # pragma: no cover (numba; raises before execution) with pytest.raises(UdfParseError, match="returned an unsupported type"): BinaryOp.register_anonymous(_pair, is_udt=True)[raw] + # An operand returned as it is keeps its type. + v = Vector(raw, size=2) + v[0] = b"abc" + first = BinaryOp.register_anonymous(lambda x, y: x, "_ret_err_first", is_udt=True) + assert v.apply(first, right=b"zz").new().to_coo()[1].tolist() == [b"abc"] for lifted in [binary.plus, unary.ainv]: with pytest.raises(KeyError, match="require a record dtype .* or an array dtype"): lifted[raw] From fc670c9856dc4b832eceee84b9ab74efa71042a6 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 13:21:42 -0500 Subject: [PATCH 32/54] Keep text that a bare text UDT would cut short out of it, and compare text widths Behavior fix. Text given beside a UDT of bytes or str (not a record) had no type, so numpy cut it to the UDT's width: with an S4 element b"abcd", b"abcdef" became b"abcd". - v == b"abcdef": TypeError on main and at 0e1ad03c; True before this commit; now False (!= is True). - Scalar == and Scalar.isequal with b"abcdef": True on main (isequal True at 0e1ad03c too), now False. - a user's op with b"abcdef": main and the commits before this one cut it silently (0e1ad03c raised TypeError); now ValueError "does not fit". Such text is now typed as a record's text field is: the UDT itself when it converts into it and back unchanged (b"ab", "ab" into S4, so bare_s4 == b"ab" works as before), else its own type at its own length, which does not fit. For eq and ne, that literal is text of the same kind at another width, so this commit also loosens, there and only there, the pairing rule of "Compare UDTs that are neither records nor arrays in eq and ne": two UDTs of text of the same kind (bytes with bytes, str with str) at different widths now pair under eq and ne and compare as text, as numpy compares them and as a record's S4 and S2 fields already do. An S4 and an S6 Vector compare element by element (they raised KeyError before; TypeError on main and at 0e1ad03c). Bytes against str, a number, float16 or any other kind still raise the pairing KeyError, as does every other op on such a pair. Stores are unchanged: v[0] = b"abcdef" still stores b"abcd", as every literal store converts by numpy's cast (int8_vec[0] = 1.5 stores 1). --- graphblas/core/literals.py | 10 ++++++++-- graphblas/core/operator/binary.py | 9 +++++---- graphblas/core/udt.py | 25 ++++++++++++++++++++----- graphblas/tests/test_op.py | 17 +++++++++++++++++ graphblas/tests/test_scalar.py | 4 +++- 5 files changed, 53 insertions(+), 12 deletions(-) diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index 7f85ad6b8..205d83d3d 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -229,13 +229,19 @@ def _literal_type(dtype, value, *, exact=False): out of a leaf's range takes a type that holds it. A numpy value is strong: its dtype, in the array layout of its shape for an array (:func:`_array_literal_udt`); a sequence holding numpy values is the - array numpy makes of it. Anything else has no type here, and - ``Scalar.from_value`` reports it. + array numpy makes of it. Text beside a UDT of text is typed as in a + record's text field: the UDT itself when it converts there and back + unchanged, else its own type at its own length, which does not fit. + Anything else has no type here, and ``Scalar.from_value`` reports it. """ value = _as_python_number(value) np_type = dtype.np_type if isinstance(value, dict) and np_type.names is not None: value = _dict_to_record(np_type, value) + if np_type.names is None and np_type.kind in "SU" and isinstance(value, (str, bytes)): + if _keeps_value(value, np_type): + return dtype + return lookup_dtype(np.dtype(f"{np_type.kind}{max(len(value), 1)}")) if type(value) in _LITERAL_KINDS or isinstance(value, (tuple, list)): literal_type = _weak_literal_udt(dtype, value, exact=exact) if literal_type is not None or type(value) in _LITERAL_KINDS: diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 8ea60ad99..4386587aa 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -87,14 +87,15 @@ def _make_udt_comparison(dtype, dtype2, *, is_eq): # Bytes, a datetime, or a complex number where the backend has no # complex types: an element is one value, compared as numpy compares # it, as a record field of that type is. ``_check_udt_pair`` has made - # sure both are the one UDT. - ptr_type = nt.CPointer(dtype.numba_type) + # sure both are the one UDT, or text of one kind. wrapper = _compile_codegen( f"def wrapper(z_ptr, x_ptr, y_ptr):\n z_ptr[0] = x_ptr[0] {cmp_op} y_ptr[0]\n", func_name="wrapper", - source_label=f"", + source_label=f"", ) - return wrapper, nt.void(nt.CPointer(INT8.numba_type), ptr_type, ptr_type) + x_ptr_type = nt.CPointer(dtype.numba_type) + y_ptr_type = nt.CPointer(dtype2.numba_type) + return wrapper, nt.void(nt.CPointer(INT8.numba_type), x_ptr_type, y_ptr_type) x_is_scalar = info_x is None y_is_scalar = info_y is None diff --git a/graphblas/core/udt.py b/graphblas/core/udt.py index dd5009fbc..6b5b852fe 100644 --- a/graphblas/core/udt.py +++ b/graphblas/core/udt.py @@ -112,13 +112,28 @@ def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): ``info_x`` and ``info_y`` are ``_get_udt_info`` results. When either is ``None`` (a built-in scalar side broadcasts), no further check fires. A UDT that is neither a record nor an array (bytes, say) holds one value, - which pairs only with its own type, on either side. Without this gate the - generated wrapper code raises a cryptic Numba ``TypingError`` on the - first field access, or, for array fields that do not broadcast, fails - inside the cfunc without telling the caller. Element dtypes may differ; + which pairs only with its own type, on either side; for ``eq`` and ``ne`` + only, text pairs with text of the same kind at another width, as numpy + and a record's text fields compare it. Without this gate the generated + wrapper code raises a cryptic Numba ``TypingError`` on the first field + access, or, for array fields that do not broadcast, fails inside the + cfunc without telling the caller. Element dtypes may differ; :func:`_udt_op_types` promotes them. """ - if dtype != dtype2 and (info_x is None and dtype._is_udt or info_y is None and dtype2._is_udt): + texts = ( + op_name in {"eq", "ne"} + and info_x is None + and info_y is None + and dtype._is_udt + and dtype2._is_udt + and dtype.np_type.kind == dtype2.np_type.kind + and dtype.np_type.kind in "SU" + ) + if ( + dtype != dtype2 + and not texts + and (info_x is None and dtype._is_udt or info_y is None and dtype2._is_udt) + ): raise KeyError( f"binary.{op_name} does not work with ({dtype}, {dtype2}): a UDT that is " "neither a record nor an array pairs only with its own type." diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index dce4589bd..39b5a1ccb 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4654,6 +4654,7 @@ def test_udt_literal_converts_only_when_exact(): (lambda: t.apply(second, right=(1.5, 1)), "_ExactText"), (lambda: b.apply(same, right=0.5), "_ExactBareText: it is typed as FP64,"), (lambda: b.apply(same, right=np.True_), "_ExactBareText: it is typed as BOOL,"), + (lambda: b.apply(same, right=b"0.5000"), "_ExactBareText"), # was cut to b"0.5" (lambda: t.apply(second, right=cut_text), "_ExactText"), (lambda: d.apply(same, right=np.timedelta64(5, "ms")), "_ExactBareSeconds"), ]: @@ -5963,6 +5964,21 @@ def test_udt_eq_compares_a_udt_that_is_neither_record_nor_array(): assert s != gb.Scalar.from_value(b"cd", text) assert s == b"ab" assert s != b"cd" + # Under eq and ne, text pairs with text of its kind at another width, as numpy + # compares it, so a literal longer than the element is unequal, not cut short. + wide = dtypes.register_anonymous(np.dtype("S6"), "_EqOpaqueWide") + u = Vector(wide, size=3) + u[0] = b"ab" + u[1] = b"cdefgh" + for left, right in [(v, u), (u, v)]: + assert binary.eq(left & right).new().to_coo()[1].tolist() == [True, False] + assert binary.ne(left & right).new().to_coo()[1].tolist() == [False, True] + four = Vector(text, size=1) + four[0] = b"abcd" + assert not (four == b"abcdef").new()[0].new().value + assert (four != b"abcdef").new()[0].new().value + assert gb.Scalar.from_value(b"abcd", text) != b"abcdef" + assert gb.Scalar.from_value(b"ab", text) == gb.Scalar.from_value(b"ab", wide) # A Scalar's == compares in numpy, so an element Numba has no type for compares too. half = dtypes.register_anonymous(np.dtype(np.float16), "_EqOpaqueHalf") h = gb.Scalar.from_value(1.5, half) @@ -6003,6 +6019,7 @@ def test_udt_eq_compares_a_udt_that_is_neither_record_nor_array(): lambda: v == 1, lambda: binary.eq(v & Vector.from_coo([0], [1], size=3)), lambda: binary.eq(v & Vector(rec, size=3)), + lambda: binary.eq(v & Vector(dtypes.register_anonymous(np.dtype("U4")), size=3)), lambda: gb.Scalar.from_value([1, 2, 3, 4, 5], pair) == h, lambda: gb.Scalar.from_value((1,), rec) == s, ]: diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index 0845be5c0..4a69521ac 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -742,6 +742,8 @@ def test_udt_eq_without_numba_matches_binary_eq(monkeypatch): (t, b"ab"), (t, b"cd"), (t, Scalar.from_value(b"ab", text)), + (t, Scalar.from_value(b"ab", dtypes.register_anonymous(np.dtype("S7")))), + (t, b"abcdefgh"), ] # UDTs that binary.eq does not pair raise its KeyError both ways. other_names = dtypes.register_anonymous( @@ -752,7 +754,7 @@ def test_udt_eq_without_numba_matches_binary_eq(monkeypatch): (r, Scalar.from_value((1, 2.5), other_names)), (a, Scalar.from_value([1, 1], arr2)), (r, a), - (t, Scalar.from_value(b"ab", dtypes.register_anonymous(np.dtype("S7")))), + (t, Scalar.from_value("ab", dtypes.register_anonymous(np.dtype("U7")))), (t, 1), # A number for a bytes field does not compare. (Scalar.from_value((b"ab", 1), np.dtype([("tx", "S2"), ("tn", np.int8)])), (1, 1)), From 9536335a0180626712fe654ddfc5be8d745d76f5 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 18:39:27 -0500 Subject: [PATCH 33/54] Keep a NaT duration NaT in another unit on numpy before 2.3 Behavior fix. numpy before 2.3 converts a NaT timedelta64 scalar into another unit as the integer NaT is stored as: NaT ms assigned into timedelta64[s] is 9223372036854774 s, and NaT minutes is 0 s. An array of durations (astype) and a date convert NaT correctly, and numpy 2.3 and later convert it correctly every way. On numpy 2.2 (numpy 1.26 converts NaT the same way), with a UDT of timedelta64[s] or a record with a timedelta64[s] field: - Scalar.from_value, v[0] = ..., Scalar.value = ..., Vector.from_coo, and the Scalar a literal becomes, given NaT ms or NaT minutes: stored 9223372036854774 s or 0 s, as on main; now NaT. - A NaT ms literal beside the bare UDT: ValueError "does not fit", since it converted to a number and so not exactly; now it converts with the DeprecationWarning, as on numpy 2.4. - (NaT ms, 1) beside the record: typed as timedelta64[ms], so binary.plus raised KeyError "no common type"; now it takes the field's type, as on numpy 2.4. The conversion into an element, and the checks of whether a literal converts exactly, now hand numpy a NaT of no unit in place of each NaT duration (alone, or in a tuple or list), which every numpy converts to NaT in any unit. numpy 2.4 gives the same results before and after. --- graphblas/core/dtypes.py | 32 +++++++++++++++++++++++++++ graphblas/core/literals.py | 8 ++++--- graphblas/core/udt.py | 4 ++-- graphblas/core/utils.py | 4 ++-- graphblas/tests/test_op.py | 10 ++++++++- graphblas/tests/test_scalar.py | 40 ++++++++++++++++++++++++++++++++++ 6 files changed, 90 insertions(+), 8 deletions(-) diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 71cce87e0..55fa6c35d 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -660,6 +660,38 @@ def _view_if_same_layout(values, np_type): return values +def _generic_nats(values, np_type): + """Return ``values`` with each NaT duration in it as a NaT of no unit, if ``np_type`` has one. + + numpy before 2.3 converts a NaT ``timedelta64`` scalar, alone or in a tuple + or list, into another unit as the integer it is stored as: NaT ms into + ``timedelta64[s]`` is 9223372036854774 s, and NaT s into ``timedelta64[ms]`` + is 0 ms. A NaT of no unit converts to NaT in every unit. An array of + durations, and a date, convert NaT correctly. + """ + if not _has_duration_leaf(np_type): + return values + return _with_generic_nats(values) + + +def _has_duration_leaf(np_type): + """Whether ``np_type`` has a timedelta64 leaf: a record field, an array's element, or itself.""" + if np_type.subdtype is not None: + return _has_duration_leaf(np_type.subdtype[0]) + if np_type.names is not None: + return any(_has_duration_leaf(np_type.fields[name][0]) for name in np_type.names) + return np_type.kind == "m" + + +def _with_generic_nats(values): + if isinstance(values, np.timedelta64): + return np.timedelta64("NaT") if np.isnat(values) else values + if isinstance(values, (tuple, list)): + items = [_with_generic_nats(x) for x in values] + return items if isinstance(values, list) else tuple(items) + return values + + def _dtype_to_string(dtype): """Convert a dtype to a string that can be safely evaluated to recreate the dtype. diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index 205d83d3d..2905cb032 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -24,7 +24,7 @@ from ..dtypes import FP64, lookup_dtype from . import utils -from .dtypes import _view_if_same_layout +from .dtypes import _generic_nats, _view_if_same_layout from .operator import IndexUnaryOp, SelectOp, TypedOpBase, binary_from_string from .operator.binary import BinaryOp from .operator.monoid import Monoid @@ -377,7 +377,9 @@ def _converts_exactly(dtype, value): with warnings.catch_warnings(): # numpy warns for some lossy conversions; the comparison below finds each. warnings.simplefilter("ignore") - converted[...] = _view_if_same_layout(value, converted.dtype) + converted[...] = _generic_nats( + _view_if_same_layout(value, converted.dtype), converted.dtype + ) except (TypeError, ValueError, OverflowError): return False for given_leaf, converted_leaf in zip( @@ -407,7 +409,7 @@ def _is_number(x): def _keeps_value(value, dtype): """Whether ``value`` converts into ``dtype`` and back unchanged (:func:`_converts_back`).""" try: - converted = np.array(value, dtype=dtype)[()] + converted = np.array(_generic_nats(value, dtype), dtype=dtype)[()] except (TypeError, ValueError, OverflowError): return False return _converts_back(value, converted) diff --git a/graphblas/core/udt.py b/graphblas/core/udt.py index 6b5b852fe..2262b82a8 100644 --- a/graphblas/core/udt.py +++ b/graphblas/core/udt.py @@ -31,7 +31,7 @@ from ..exceptions import DomainMismatch from . import _has_numba, ffi, utils -from .dtypes import _registry, _view_if_same_layout, lookup_dtype +from .dtypes import _generic_nats, _registry, _view_if_same_layout, lookup_dtype from .utils import output_type # The built-in ops that lift to UDTs, each with the Python operator or function @@ -750,7 +750,7 @@ def _element_bytes(dtype, value): value = _dict_to_record(np_type, value) else: arr = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) - value = _view_if_same_layout(value, arr.dtype) + value = _generic_nats(_view_if_same_layout(value, arr.dtype), arr.dtype) if _has_narrow_float_leaf(np_type) and np.geterr()["over"] == "warn": # A float too large for a float32 leaf is infinity, and numpy warns; # warn at the caller's line instead, as a built-in FP32 does, so the diff --git a/graphblas/core/utils.py b/graphblas/core/utils.py index 6b8ec12ed..37cc5e351 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -4,7 +4,7 @@ from ..dtypes import _INDEX, lookup_dtype from . import ffi, lib -from .dtypes import _view_if_same_layout +from .dtypes import _generic_nats, _view_if_same_layout _NP2 = np.__version__.startswith("2.") @@ -121,7 +121,7 @@ def values_to_numpy_buffer( np_type = _get_subdtype(dtype.np_type) # https://numpy.org/doc/stable/release/2.0.0-notes.html#new-copy-keyword-meaning-for-array-and-asarray-constructors array = np.array( - _view_if_same_layout(array, np_type), + _generic_nats(_view_if_same_layout(array, np_type), np_type), np_type, copy=copy or _NP2 and None, order=order, diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 39b5a1ccb..91c7b3430 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4681,8 +4681,10 @@ def test_udt_literal_converts_only_when_exact(): assert t.apply(second, right=text_literal).new()[0].new().value.tolist() == (b"cd", 2) with pytest.warns(DeprecationWarning, match="does not fit _ExactBareSeconds"): assert d.apply(same, right=np.timedelta64(5000, "ms")).new()[0].new().value + # NaT stays NaT in any unit, and equals nothing. + with pytest.warns(DeprecationWarning, match="does not fit _ExactBareSeconds"): + assert np.isnat(d.apply(second, right=np.timedelta64("NaT", "ms")).new()[0].new().value) with pytest.warns(DeprecationWarning, match="does not fit _ExactBareSeconds"): - # NaT stays NaT in any unit, and equals nothing. assert not d.apply(same, right=np.timedelta64("NaT", "ms")).new()[0].new().value # Text for a text leaf converts, bytes or str. for literal in [b"cd", "cd"]: @@ -4697,6 +4699,8 @@ def test_udt_literal_converts_only_when_exact(): assert fine["ex_d"] == np.datetime64(5, "s") nat = when.apply(second, right=(np.datetime64("NaT", "ms"), 1)).new()[0].new().value assert np.isnat(nat["ex_d"]) + # A duration in a coarser unit casts safely, so it fits by type. + assert np.isnat(d.apply(second, right=np.timedelta64("NaT", "m")).new()[0].new().value) assert (t == (b"ab", 1)).new()[0].new().value assert not (t == (b"ab", 1.5)).new()[0].new().value assert (when == (np.datetime64(0, "s"), 1)).new()[0].new().value @@ -4853,6 +4857,10 @@ def test_udt_sequence_and_array_literals_type_like_numbers(): assert result[0].new().value["sq_sd"] == np.timedelta64(total, "s") with pytest.raises(KeyError, match="no common type"): binary.plus(sv, (np.timedelta64(5, "ms"), 1)) # 5 ms is no whole number of seconds + for nat in [np.timedelta64("NaT", "ms"), np.timedelta64("NaT", "m")]: + result = binary.plus(sv, (nat, 1)).new() # NaT is NaT in seconds + assert result.dtype == secs + assert np.isnat(result[0].new().value["sq_sd"]) first = dtypes.register_anonymous( np.dtype([("sq_fin", [("sq_fb", np.int8), ("sq_fc", np.float32)]), ("sq_fa", np.int16)]) ) diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index 4a69521ac..bbaaaff96 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -279,6 +279,46 @@ def test_udt_scalar_nested_array_roundtrip(): np.testing.assert_array_equal(got, val) +def test_udt_nat_duration_in_another_unit_stays_nat(): + """A NaT duration stored into a duration of another unit is NaT, not a number. + + numpy before 2.3 converts a NaT ``timedelta64`` scalar into another unit as + the integer it is stored as: NaT ms became 9223372036854774 s, and NaT + minutes 0 s. A date or an array of durations was not affected. + """ + seconds = dtypes.register_anonymous(np.dtype("m8[s]"), "_NatSeconds") + rec = dtypes.register_anonymous(np.dtype([("nat_d", "m8[s]"), ("nat_n", np.int8)]), "_NatRec") + days = dtypes.register_anonymous(np.dtype([("nat_t", "M8[D]"), ("nat_n", np.int8)])) + for unit in ["ns", "ms", "m", "D"]: + nat = np.timedelta64("NaT", unit) + assert np.isnat(Scalar.from_value(nat, seconds).value) + s = Scalar(seconds) + s.value = nat + assert np.isnat(s.value) + v = Vector(seconds, size=2) + v[0] = nat + assert np.isnat(v[0].new().value) + v[:] = nat + assert np.isnat(v.to_coo()[1]).all() + r = Vector(rec, size=2) + r[0] = (nat, 1) + r[1] = {"nat_d": nat, "nat_n": 1} + assert np.isnat(r.to_coo()[1]["nat_d"]).all() + # Off SuiteSparse, ``build`` uses ``binary.any`` as its ``dup_op``, which a + # UDT gets only by compiling it with Numba. + if suitesparse or gb.core._supports_udfs: + v = Vector.from_coo([0, 1], [nat, nat], dtype=seconds) + assert np.isnat(v.to_coo()[1]).all() + r = Vector.from_coo([0, 1], [(nat, 1), (nat, 2)], dtype=rec) + assert np.isnat(r.to_coo()[1]["nat_d"]).all() + r = Vector(days, size=1) + r[0] = (np.datetime64("NaT", unit), 1) + assert np.isnat(r[0].new().value["nat_t"]) + # Only NaT changes how it converts; a value converts as numpy converts it. + assert Scalar.from_value(np.timedelta64(5000, "ms"), seconds).value == np.timedelta64(5, "s") + assert Scalar.from_value(np.timedelta64(2, "m"), seconds).value == np.timedelta64(120, "s") + + def test_nvals(s): assert s.nvals == 1 s.clear() From ff9c4334e0779ab36018e6d42ab3732ee14abd49 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 20:14:31 -0500 Subject: [PATCH 34/54] Keep the JIT layout test's record short enough to name on SuiteSparse 8 SuiteSparse before 9 has no GrB_NAME setter, so registration names a UDT by its numpy repr, which must fit in 128 bytes, and warns when it does not; pytest makes the warning an error. test_record_layouts_c_cannot_declare_skip_jit registered its record with an array field with align=True, which gives the dict form of the repr: 143 bytes. Its float64 fields have the same offsets packed or aligned, so it is now registered packed (a 71-byte repr); a C struct still cannot declare it, and the test runs the same lines. Every other UDT the tests register with a repr over 128 bytes is in a test that ignores that warning. --- graphblas/tests/test_ssjit.py | 1 - 1 file changed, 1 deletion(-) diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 2b83a4f85..f11d2766a 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -1638,7 +1638,6 @@ def test_record_layouts_c_cannot_declare_skip_jit(): outer = np.dtype([("ol_in", offset), ("ol_z", np.float64)]) with_array = np.dtype( [("ol_x", np.float64), ("ol_rec", [("ol_v", np.float64, (2,)), ("ol_w", np.float64)])], - align=True, ) for np_type, value, doubled in [ (offset, (1, 400), (2, 800)), From f6d912908cab05933642525da10d6f3e8e910a41 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 21 Oct 2026 01:17:31 -0500 Subject: [PATCH 35/54] Convert a NaT duration in its leaf's unit, not as a NaT of no unit "Keep a NaT duration NaT in another unit on numpy before 2.3" handed numpy a NaT of no unit, np.timedelta64("NaT"), in place of each NaT duration. numpy 2.5 deprecates that generic unit ("The 'generic' unit for NumPy timedelta is deprecated"), so on numpy 2.5 every store or literal of a NaT duration into a duration UDT or field warned with a DeprecationWarning, which the test suite's filterwarnings makes an error: test_udt_nat_duration_in_another_unit_stays_nat, test_udt_sequence_and_array_literals_type_like_numbers and test_udt_literal_converts_only_when_exact failed. Each NaT duration is now converted into the unit of the leaf it fills by converting it as an array (astype), which keeps NaT on every numpy, before numpy assigns it. The value is walked beside the dtype: a tuple of a record's length as its fields, a list (or a tuple for any other dtype) as a sequence of elements, and an array UDT by its element type. Nothing makes a generic unit. numpy 2.5.4: DeprecationWarning before, NaT with no warning now. numpy 1.26, 2.2 and 2.4: NaT in the target unit before and after. New rows: a nested record's duration field, and an array UDT of durations given a list, where Numba is absent (Numba cannot type an array of durations, so registering one fails with it). --- graphblas/core/dtypes.py | 33 +++++++++++++++++++-------------- graphblas/core/literals.py | 6 +++--- graphblas/core/udt.py | 4 ++-- graphblas/core/utils.py | 4 ++-- graphblas/tests/test_scalar.py | 10 ++++++++++ 5 files changed, 36 insertions(+), 21 deletions(-) diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 55fa6c35d..acc963c43 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -660,18 +660,32 @@ def _view_if_same_layout(values, np_type): return values -def _generic_nats(values, np_type): - """Return ``values`` with each NaT duration in it as a NaT of no unit, if ``np_type`` has one. +def _nats_in_leaf_units(values, np_type): + """Return ``values`` with each NaT duration in it in the unit of the leaf it fills. numpy before 2.3 converts a NaT ``timedelta64`` scalar, alone or in a tuple or list, into another unit as the integer it is stored as: NaT ms into ``timedelta64[s]`` is 9223372036854774 s, and NaT s into ``timedelta64[ms]`` - is 0 ms. A NaT of no unit converts to NaT in every unit. An array of - durations, and a date, convert NaT correctly. + is 0 ms. Converting an array (``astype``) keeps NaT, as does a date, so each + NaT scalar is converted that way first. ``values`` holds elements of + ``np_type``: a list is a sequence of them, and a tuple is a record's fields + (or, for any other dtype, a sequence too). """ if not _has_duration_leaf(np_type): return values - return _with_generic_nats(values) + while np_type.subdtype is not None: + np_type = np_type.subdtype[0] + if isinstance(values, list) or (isinstance(values, tuple) and np_type.names is None): + items = [_nats_in_leaf_units(x, np_type) for x in values] + return items if isinstance(values, list) else tuple(items) + if isinstance(values, tuple) and len(values) == len(np_type.names): + return tuple( + _nats_in_leaf_units(x, np_type.fields[name][0]) + for x, name in zip(values, np_type.names, strict=True) + ) + if np_type.kind == "m" and isinstance(values, np.timedelta64) and np.isnat(values): + return np.asarray(values).astype(np_type)[()] + return values def _has_duration_leaf(np_type): @@ -683,15 +697,6 @@ def _has_duration_leaf(np_type): return np_type.kind == "m" -def _with_generic_nats(values): - if isinstance(values, np.timedelta64): - return np.timedelta64("NaT") if np.isnat(values) else values - if isinstance(values, (tuple, list)): - items = [_with_generic_nats(x) for x in values] - return items if isinstance(values, list) else tuple(items) - return values - - def _dtype_to_string(dtype): """Convert a dtype to a string that can be safely evaluated to recreate the dtype. diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index 2905cb032..d7a24bc48 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -24,7 +24,7 @@ from ..dtypes import FP64, lookup_dtype from . import utils -from .dtypes import _generic_nats, _view_if_same_layout +from .dtypes import _nats_in_leaf_units, _view_if_same_layout from .operator import IndexUnaryOp, SelectOp, TypedOpBase, binary_from_string from .operator.binary import BinaryOp from .operator.monoid import Monoid @@ -377,7 +377,7 @@ def _converts_exactly(dtype, value): with warnings.catch_warnings(): # numpy warns for some lossy conversions; the comparison below finds each. warnings.simplefilter("ignore") - converted[...] = _generic_nats( + converted[...] = _nats_in_leaf_units( _view_if_same_layout(value, converted.dtype), converted.dtype ) except (TypeError, ValueError, OverflowError): @@ -409,7 +409,7 @@ def _is_number(x): def _keeps_value(value, dtype): """Whether ``value`` converts into ``dtype`` and back unchanged (:func:`_converts_back`).""" try: - converted = np.array(_generic_nats(value, dtype), dtype=dtype)[()] + converted = np.array(_nats_in_leaf_units(value, dtype), dtype=dtype)[()] except (TypeError, ValueError, OverflowError): return False return _converts_back(value, converted) diff --git a/graphblas/core/udt.py b/graphblas/core/udt.py index 2262b82a8..90dbcea47 100644 --- a/graphblas/core/udt.py +++ b/graphblas/core/udt.py @@ -31,7 +31,7 @@ from ..exceptions import DomainMismatch from . import _has_numba, ffi, utils -from .dtypes import _generic_nats, _registry, _view_if_same_layout, lookup_dtype +from .dtypes import _nats_in_leaf_units, _registry, _view_if_same_layout, lookup_dtype from .utils import output_type # The built-in ops that lift to UDTs, each with the Python operator or function @@ -750,7 +750,7 @@ def _element_bytes(dtype, value): value = _dict_to_record(np_type, value) else: arr = np.zeros(np_type.subdtype[1], dtype=np_type.subdtype[0]) - value = _generic_nats(_view_if_same_layout(value, arr.dtype), arr.dtype) + value = _nats_in_leaf_units(_view_if_same_layout(value, arr.dtype), arr.dtype) if _has_narrow_float_leaf(np_type) and np.geterr()["over"] == "warn": # A float too large for a float32 leaf is infinity, and numpy warns; # warn at the caller's line instead, as a built-in FP32 does, so the diff --git a/graphblas/core/utils.py b/graphblas/core/utils.py index 37cc5e351..14f4bbe5d 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -4,7 +4,7 @@ from ..dtypes import _INDEX, lookup_dtype from . import ffi, lib -from .dtypes import _generic_nats, _view_if_same_layout +from .dtypes import _nats_in_leaf_units, _view_if_same_layout _NP2 = np.__version__.startswith("2.") @@ -121,7 +121,7 @@ def values_to_numpy_buffer( np_type = _get_subdtype(dtype.np_type) # https://numpy.org/doc/stable/release/2.0.0-notes.html#new-copy-keyword-meaning-for-array-and-asarray-constructors array = np.array( - _generic_nats(_view_if_same_layout(array, np_type), np_type), + _nats_in_leaf_units(_view_if_same_layout(array, np_type), np_type), np_type, copy=copy or _NP2 and None, order=order, diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index bbaaaff96..686e4861a 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -289,6 +289,11 @@ def test_udt_nat_duration_in_another_unit_stays_nat(): seconds = dtypes.register_anonymous(np.dtype("m8[s]"), "_NatSeconds") rec = dtypes.register_anonymous(np.dtype([("nat_d", "m8[s]"), ("nat_n", np.int8)]), "_NatRec") days = dtypes.register_anonymous(np.dtype([("nat_t", "M8[D]"), ("nat_n", np.int8)])) + nested = dtypes.register_anonymous( + np.dtype([("nat_o", [("nat_d", "m8[s]")]), ("nat_n", np.int8)]), "_NatNested" + ) + # Numba cannot type an array of durations, so registering one needs it absent. + pair = None if gb.core._has_numba else dtypes.register_anonymous(np.dtype(("m8[s]", 2))) for unit in ["ns", "ms", "m", "D"]: nat = np.timedelta64("NaT", unit) assert np.isnat(Scalar.from_value(nat, seconds).value) @@ -311,6 +316,11 @@ def test_udt_nat_duration_in_another_unit_stays_nat(): assert np.isnat(v.to_coo()[1]).all() r = Vector.from_coo([0, 1], [(nat, 1), (nat, 2)], dtype=rec) assert np.isnat(r.to_coo()[1]["nat_d"]).all() + assert np.isnat(Scalar.from_value(((nat,), 1), nested).value["nat_o"]["nat_d"]) + if pair is not None: + got = Scalar.from_value([nat, np.timedelta64(5000, "ms")], pair).value + assert np.isnat(got[0]) + assert got[1] == np.timedelta64(5, "s") r = Vector(days, size=1) r[0] = (np.datetime64("NaT", unit), 1) assert np.isnat(r[0].new().value["nat_t"]) From 16f679c2c02fc39c4d328da792bb195b82c18750 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 27 Oct 2026 09:05:43 -0500 Subject: [PATCH 36/54] Skip the JIT part of the one-name UDT layouts test without a C compiler test_udt_layouts_under_one_name_get_their_own_jit_c_names set jit_c_control to 'on' and compiled kernels for both layouts, so on a host whose JIT cannot compile (no C compiler, or a broken one) it failed with JitError instead of skipping. CI always installs a compiler, so CI never saw it. It now asks _enable_jit_for_udt, which probes with a real compile once per process, and skips the kernel part when the answer is no. The C name checks and the register_new check need no compiler, so they come first and run everywhere. --- graphblas/tests/test_dtype.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index 57358ec73..46d3e0b16 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -554,6 +554,14 @@ def test_udt_layouts_under_one_name_get_their_own_jit_c_names(): assert small.jit_c_name == "_SameJitName" assert big.jit_c_name != small.jit_c_name assert big.jit_c_name in big.jit_c_definition + # Nor may a typedef registered in C take a name a UDT already has. + with pytest.raises(ValueError, match="unavailable"): + dtypes.ss.register_new("_SameJitName", "typedef struct { int8_t q ; } _SameJitName ;") + # The rest compiles kernels for both layouts, which needs a C compiler the JIT can use. + from graphblas.core.ss.jit_config import _enable_jit_for_udt + + if not _enable_jit_for_udt(): + pytest.skip("the JIT cannot compile here (no usable C compiler)") jit_c_control = gb.ss.config["jit_c_control"] gb.ss.config["jit_c_control"] = "on" try: @@ -569,9 +577,6 @@ def test_udt_layouts_under_one_name_get_their_own_jit_c_names(): assert values == [one.tolist(), one.tolist(), two.tolist()] finally: gb.ss.config["jit_c_control"] = jit_c_control - # Nor may a typedef registered in C take a name a UDT already has. - with pytest.raises(ValueError, match="unavailable"): - dtypes.ss.register_new("_SameJitName", "typedef struct { int8_t q ; } _SameJitName ;") def test_dtype_to_from_string(): From 2a564c22ef6e555105174e385d468e3df5ea1bc1 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 28 Oct 2026 00:56:12 -0500 Subject: [PATCH 37/54] Expand an iso array UDT's value without an extra axis in exports Behavior fix. SuiteSparse exports the iso value of an array UDT with shape (1, *subarray shape), and normalize_values broadcast it to shape + values.shape, keeping that axis of length 1. With the suitesparse backend, for an iso Vector or Matrix of an array UDT such as INT64[4] (as v[:] = [0, 1, 2, 3] makes): - Vector.to_dense: TypeError "unsupported operand type(s) for +: 'int' and 'tuple'" on main and before this commit (it passed its size as an int); now an array of shape (size, 4). - Matrix.to_dense: shape (nrows, ncols, 1, 4); now (nrows, ncols, 4). - Matrix.to_csr, to_csc, to_dcsr and to_dcsc: values of shape (nvals, 1, 4); now (nvals, 4). - Matrix.to_dicts, which uses to_dcsr or to_dcsc: each value as [[0, 1, 2, 3]]; now [0, 1, 2, 3]. Vector.to_coo, Matrix.to_coo, the vanilla backend and records were already right. normalize_values now broadcasts to shape + values.shape[1:], and Vector.to_dense passes (size,). CI found it through #652's test_udt_mixed_jit_c_name_is_its_own (17 jobs, e.g. job 122633175811): SuiteSparse before 10.5 makes a vector built from equal values iso, so that test's constant inputs gave an iso result, and to_dense raised. With 10.5 the vectors are not iso. --- graphblas/core/utils.py | 6 ++---- graphblas/core/vector.py | 2 +- graphblas/tests/test_matrix.py | 13 +++++++++++++ graphblas/tests/test_vector.py | 11 +++++++++++ 4 files changed, 27 insertions(+), 5 deletions(-) diff --git a/graphblas/core/utils.py b/graphblas/core/utils.py index 14f4bbe5d..9f15254ab 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -156,10 +156,8 @@ def normalize_values(self, values, dtype, shape=None, is_iso=False): else: dtype = self.dtype if is_iso: - if dtype.np_type.subdtype is None: - values = np.broadcast_to(values, shape) - else: - values = np.broadcast_to(values, shape + values.shape) + # The iso value of an array UDT is exported as shape (1, *subarray shape) + values = np.broadcast_to(values, shape + values.shape[1:]) return values diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index 0a6087a18..0b980f7f1 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -947,7 +947,7 @@ def to_dense(self, fill_value=None, dtype=None, **opts): ) if backend == "suitesparse": info = self.ss.export("full") - return normalize_values(self, info["values"], dtype, self._size, info["is_iso"]) + return normalize_values(self, info["values"], dtype, (self._size,), info["is_iso"]) return self.to_coo(dtype, indices=False)[1] if dtype is None and not self.dtype._is_udt: diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index fcb2be985..d1405c49f 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -4341,6 +4341,19 @@ def test_wait_chains(A): assert result == 47 +def test_subarray_dtypes_iso(): + # An iso array UDT exports its one value with a leading axis of length 1 + A = Matrix("INT64[4]", nrows=2, ncols=3) + A[:, :] = [0, 1, 2, 3] + if suitesparse: + assert A.ss.is_iso + expected = np.broadcast_to([0, 1, 2, 3], (6, 4)) + np.testing.assert_array_equal(A.to_dense(), expected.reshape(2, 3, 4)) + for method in ["coo", "csr", "csc", "dcsr", "dcsc"]: + np.testing.assert_array_equal(getattr(A, f"to_{method}")()[-1], expected) + assert A.to_dicts() == {i: {j: [0, 1, 2, 3] for j in range(3)} for i in range(2)} + + @pytest.mark.skipif("not supports_udfs") def test_subarray_dtypes(): a = np.arange(3 * 4, dtype=np.int64).reshape(3, 4) diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 87f05dad2..9f4b68747 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -2928,6 +2928,17 @@ def test_ss_sort(v): w, p = v.ss.sort(monoid.lxor) # Weird, but user-defined monoids may not commute, so okay +def test_subarray_dtypes_iso(): + # An iso array UDT exports its one value with a leading axis of length 1 + v = Vector("INT64[4]", size=3) + v[:] = [0, 1, 2, 3] + if suitesparse: + assert v.ss.is_iso + expected = np.broadcast_to([0, 1, 2, 3], (3, 4)) + np.testing.assert_array_equal(v.to_dense(), expected) + np.testing.assert_array_equal(v.to_coo()[1], expected) + + @pytest.mark.skipif("not supports_udfs") def test_subarray_dtypes(): a = np.arange(3 * 4, dtype=np.int64).reshape(3, 4) From 8a109a5d430fe27094820ecc359fb148c1c3fb1d Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 28 Oct 2026 00:57:08 -0500 Subject: [PATCH 38/54] Refuse a zero-size UDT, so a str literal raises TypeError on SuiteSparse 7 Behavior fix. register_anonymous passed a dtype's itemsize to GrB_Type_new without a check. SuiteSparse 8 and later refuse a size of 0, but SuiteSparse 7 takes it as 1. lookup_dtype auto-registers what it does not know, so on SuiteSparse 7 a str became the UDT " Date: Wed, 28 Oct 2026 09:37:14 -0500 Subject: [PATCH 39/54] Note that SuiteSparse:GraphBLAS 10.5.2 fixed select with a value op SuiteSparse:GraphBLAS 10.5.2 fixed the bug _value_select_in_own_type works around (DrTimothyAldenDavis/GraphBLAS#463): select with a VALUE* op of another type than the input's now casts both to the op's type. Its MAINT note said to delete the function once every supported version has the fix, but that would change two results that test_select_value_thunk_of_another_type pins: with the function bypassed on 10.5.2, an INT64 thunk of 2**53 + 1 beside FP32 or FC32 values of 2**53 rounds to 2**53 in the op's type and keeps that element. The other 61 cases of the test match. The note now says so, and gives the versions with the bug as 8.0.2 through 10.5.1. Comments only. --- graphblas/core/operator/select.py | 11 ++++++----- graphblas/tests/test_vector.py | 7 ++++--- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 18e001fcc..2be05fb4c 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -16,7 +16,7 @@ def _value_select_in_own_type(op, dtype, thunk, given): """Return ``op`` and ``thunk`` for ``select`` on built-in ``dtype``, as a value op of ``dtype``. - SuiteSparse:GraphBLAS (8.0 to 10.5, at least) runs ``select`` with a value + SuiteSparse:GraphBLAS 8.0 to 10.5.1 runs ``select`` with a value op typed for another type than the input's by reading the thunk's bytes as the input's type, where the spec casts both to the op's type: on an INT8 vector, ``valueeq`` with ``1.0`` compared with 0, and ``valuelt`` with @@ -31,10 +31,11 @@ def _value_select_in_own_type(op, dtype, thunk, given): it (an int beyond 64 bits is FP64). Other ops and types are returned unchanged. """ - # MAINT 2026-10-07: works around the SuiteSparse:GraphBLAS bug above (seen in - # 8.0.2 through 10.6.0; fix proposed in DrTimothyAldenDavis/GraphBLAS#463). - # Once every supported version has the fix, pass the typed op and thunk - # through and delete this function. + # MAINT 2026-10-28: works around the SuiteSparse:GraphBLAS bug above (seen in + # 8.0.2 through 10.5.1; fixed in 10.5.2 by DrTimothyAldenDavis/GraphBLAS#463). + # The fix casts both to the op's type, which can round the thunk: beside FP32 + # values, an INT64 thunk of 2**53 + 1 is FP64 2**53 and equals 2**53. Once + # every supported version has the fix, keep only what makes the comparison exact. if op.type is dtype or dtype._is_udt or op.parent.name not in _VALUE_OPS: return op, thunk import math diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 9f4b68747..7d5cab524 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -1665,11 +1665,12 @@ def test_python_int_beyond_64_bit_element_compares_exactly(): def test_select_value_thunk_of_another_type(): """Select with a value op compares exactly when the thunk's type differs from the values'. - SuiteSparse:GraphBLAS's select reads such a thunk's bytes as the values' - type: on an INT8 vector, ``valueeq`` with 1.0 compared with 0 and kept + SuiteSparse:GraphBLAS's select before 10.5.2 read such a thunk's bytes as the + values' type: on an INT8 vector, ``valueeq`` with 1.0 compared with 0 and kept nothing, ``valuelt`` with 300 compared with 44, and ``valueeq`` with 356 kept 100. python-graphblas now compares in the values' own type with a - thunk that keeps the same elements. + thunk that keeps the same elements. 10.5.2 casts both to the op's type, + which rounds the 2**53 + 1 thunks below to 2**53 and keeps that element. """ i8 = Vector.from_coo([0, 1, 2, 3], [-5, 1, 2, 100], dtype=dtypes.INT8) u8 = Vector.from_coo([0, 1], [0, 255], dtype=dtypes.UINT8) From 1034473e88e3466cc652c12d6d6f2b2c9be2cf6c Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 28 Oct 2026 23:11:28 -0500 Subject: [PATCH 40/54] Say which array Numba cannot type when registering a UDT Behavior change. Registering a UDT that holds an array of bool, datetime64, timedelta64, str or bytes (as the UDT, or as a field of a record) raised Numba's own error, which does not say what to change: dtypes.register_anonymous(np.dtype((np.bool_, (2,)))) # was AttributeError: 'Boolean' object has no attribute 'bitwidth' # now ValueError: Numba, which python-graphblas uses for UDTs, cannot # type an array of bool, as in ('?', (2,)); use an array of uint8 # instead The ValueError, the type of the other errors register_anonymous raises for a dtype it refuses, is chained from Numba's AttributeError. The np_type of dtypes.ss.register_new gets the same error. Without Numba these dtypes still register, as before. Numba now types the dtype before GrB_Type_new, so a dtype it cannot type no longer leaves an unused GrB_Type behind. lookup_dtype, which registers a numpy dtype it has not seen, still says "Unknown dtype" for these, as for any dtype it cannot register. --- docs/user_guide/udt.rst | 3 +++ graphblas/core/dtypes.py | 42 ++++++++++++++++++++++++++++++++++- graphblas/core/ss/dtypes.py | 2 +- graphblas/tests/test_dtype.py | 21 ++++++++++++++++++ graphblas/tests/test_ssjit.py | 13 +++++++++++ 5 files changed, 79 insertions(+), 2 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 7b731fb44..49dc44ef7 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -74,6 +74,9 @@ Field type rules: struct) are supported as fields, and built-in operators lift through them. In the dict and dataclass forms, write an array field as ``"FP64[3]"``; the dict form also takes an existing UDT, such as ``{"id": int, "pos": point3}``. +- Numba cannot type an array of ``bool``, ``datetime64`` or ``timedelta64``, so + with Numba installed, registering a UDT that holds one raises ``ValueError``. + Use an array of ``uint8`` for ``bool``, and of ``int64`` for the others. - Strings and Python objects are not supported. - Dataclass annotation strings (e.g., ``"int"``) resolve through ``lookup_dtype``. diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 455951485..484dcf354 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -173,6 +173,45 @@ def _set_udt_jit_c_definition(datatype): datatype._jit_c_info = info +# Elements Numba cannot type in an array (its NestedArray needs a bitwidth they lack), and +# what to store instead +_NUMBA_UNTYPED_ARRAY_ELEMENTS = { + "b": ("bool", "uint8"), + "M": ("datetime64", "int64"), + "m": ("timedelta64", "int64"), + "U": ("str", "uint32 (one per character)"), + "S": ("bytes", "uint8 (one per byte)"), +} + + +def _numba_untyped_array_element(dtype): + """The kind of the first array element in dtype that Numba cannot type, or None.""" + if dtype.subdtype is not None: + base = dtype.subdtype[0] + if base.kind in _NUMBA_UNTYPED_ARRAY_ELEMENTS: + return base.kind + return _numba_untyped_array_element(base) + for field_dtype, *_ in (dtype.fields or {}).values(): + if (kind := _numba_untyped_array_element(field_dtype)) is not None: + return kind + return None + + +def _numba_type(dtype): + try: + return numba.typeof(dtype).dtype + except AttributeError as exc: + # Numba's own error is "'Boolean' object has no attribute 'bitwidth'" and the like + kind = _numba_untyped_array_element(dtype) + if kind is None: # pragma: no cover (safety) + raise + element, instead = _NUMBA_UNTYPED_ARRAY_ELEMENTS[kind] + raise ValueError( + f"Numba, which python-graphblas uses for UDTs, cannot type an array of {element}, " + f"as in {dtype}; use an array of {instead} instead" + ) from exc + + def register_new(name, dtype=None): """Register a user-defined type and put it on the ``gb.dtypes`` namespace. @@ -297,6 +336,8 @@ def register_anonymous(dtype, name=None): # SuiteSparse 7 takes a size of 0 as 1 (later versions raise InvalidValue), which # made ``str`` a UDT that holds nothing, so a str literal was not a TypeError raise ValueError(f"dtype must not have a size of 0: {dtype}") + # Before GrB_Type_new, so a dtype Numba cannot type does not leave a GrB_Type behind + numba_type = _numba_type(dtype) if _has_numba else None from ..exceptions import check_status_carg @@ -334,7 +375,6 @@ def register_anonymous(dtype, name=None): # For now, let's use "opaque" unsigned bytes for the c type. if name is None: name = _default_name(dtype) - numba_type = numba.typeof(dtype).dtype if _has_numba else None rv = DataType(name, gb_obj, None, f"uint8_t[{dtype.itemsize}]", numba_type, dtype) _registry[gb_obj] = rv _registry[dtype] = rv diff --git a/graphblas/core/ss/dtypes.py b/graphblas/core/ss/dtypes.py index 029721590..e82dc1224 100644 --- a/graphblas/core/ss/dtypes.py +++ b/graphblas/core/ss/dtypes.py @@ -89,7 +89,7 @@ def register_new(name, jit_c_definition, *, np_type=None): np_type = np.dtype((np.uint8, size)) save_np_type = False if _has_numba: - numba_type = numba.typeof(np_type).dtype + numba_type = core.dtypes._numba_type(np_type) else: numba_type = None diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index b1a5a9655..140d49a4c 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -376,6 +376,27 @@ def test_register_anonymous_zero_size(): dtypes.lookup_dtype(str) +def test_register_anonymous_array_numba_cannot_type(): + # Numba cannot type an array of these, alone or in a record. Its own error was + # "AttributeError: 'Boolean' object has no attribute 'bitwidth'" and the like. + for np_type, element, instead in [ + (np.dtype((np.bool_, (2,))), "bool", "uint8"), + (np.dtype(("M8[s]", (2, 3))), "datetime64", "int64"), + (np.dtype([("_nbt_a", np.int8), ("_nbt_b", "m8[ns]", (2,))]), "timedelta64", "int64"), + (np.dtype(("U3", (2,))), "str", "uint32"), + (np.dtype([("_nbt_c", [("_nbt_d", "S2", (2,))])]), "bytes", "uint8"), + ]: + if not supports_udfs: + # Nothing types it without Numba, and it registers + assert dtypes.register_anonymous(np_type).np_type == np_type + continue + match = f"cannot type an array of {element}.*use an array of {instead}" + with pytest.raises(ValueError, match=match) as exc_info: + dtypes.register_anonymous(np_type) + assert isinstance(exc_info.value.__cause__, AttributeError) + assert np_type not in core.dtypes._registry + + def test_register_anonymous_from_dataclass_deferred_annotations(): """Dataclass fields annotated as strings (PEP 563 / 649) resolve through ``lookup_dtype``.""" from dataclasses import make_dataclass diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index f11d2766a..ad32c5200 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -591,6 +591,19 @@ def test_jit_udt(): assert dtype.np_type == np_type +@pytest.mark.skipif("_IS_SSGB7 or numba is None") +def test_jit_udt_np_type_numba_cannot_type(): + # The np_type gets the error register_anonymous gives for an array Numba cannot type + _require_jit_on() + np_type = np.dtype([("x", np.bool_, (2,))]) + with pytest.raises(ValueError, match="cannot type an array of bool") as exc_info: + dtypes.ss.register_new( + "_jit_bool_pair", "typedef struct { uint8_t x [2] ; } _jit_bool_pair ;", np_type=np_type + ) + assert isinstance(exc_info.value.__cause__, AttributeError) + assert not hasattr(dtypes.ss, "_jit_bool_pair") + + @pytest.mark.skipif("_IS_SSGB7") def test_jit_udt_name_leaves_room_for_the_nul(): """A type name as long as ``GxB_MAX_NAME_LEN`` is rejected before SuiteSparse sees it. From 0a6c4b578ea68f9a1caae2222941699f22bbad34 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 28 Oct 2026 23:52:14 -0500 Subject: [PATCH 41/54] Type an IndexBinaryOp bound to theta when it runs, with theta of its own type Binding theta to an IndexBinaryOp of no type, op(theta), took the op's type from theta alone, and the inputs were cast to it. Now op(theta) returns a BinaryOp whose type is chosen each time it runs, as any other BinaryOp's is. Theta is typed by its value when it is given, as apply types an IndexUnaryOp's thunk: a Python int is INT64, a float FP64 and a bool BOOL, and a numpy scalar or a Scalar keeps its type. A numpy array is typed as its array layout, as other array literals are. The default theta, False, changes no type. An op registered without is_udt runs at the type that holds both inputs and theta. One with is_udt=True is compiled for the types of x, y and theta as they are, and SuiteSparse gets theta's type as the op's theta_type, which GxB_IndexBinaryOp_new takes apart from x's and y's. So theta is never converted. The compiled ops are cached per (x, y, theta) types, and op[X, Y, T] names all three; op[X, Y] and op[T] give theta y's type, as before. A Semiring built from op(theta) is polymorphic in the same way. op[T](theta) and op(theta, dtype=T) still bind at T. IndexBinaryOp has not been released, so nothing is deprecated. Behavior changes (before -> after): - A.ewise_mult(B, op()) on INT64: a BOOL op, so every value became True -> INT64, the values op computes. - u.ewise_mult(w, op(2)) on FP64: computed in INT64 (1.5 + 2 * 2.25 gave 5) -> FP64 (6.0). op(np.int8(2)): computed in INT8, where 300.0 became 127 and wrapped -> FP64. - A Semiring from op(0.5) with INT64, FP32 or INT8 matrices: KeyError -> computed in FP64, with or without is_udt. - The SuiteSparse user guide's argmax (an is_udt multiplier that makes the pair (k, A(i, k)), theta unused) on a matrix of pairs: iop(), iop(False), iop(0) and iop(0.0) were ops on BOOL, INT64 or FP64 values, which failed to compile for the pairs -> all four run, compiled for (pair, BOOL, BOOL), (pair, BOOL, INT64) and (pair, BOOL, FP64). - An is_udt op on INT64 data: op(0.5) or an FP32 Scalar theta made the op FP64 or FP32 and cast the inputs to it -> the inputs stay INT64, theta keeps its type, and Numba types the result (FP64 for x + y + theta). - An is_udt op on int32[3] data: op(np.array([10, 20, 30], np.int32)) raised "Cannot infer a dtype for theta" -> theta is an int32[3] value; an int64 array is an int64[3] one. op(5) gave GrB_DOMAIN_MISMATCH -> theta is an INT64. op((10, 20, 30)) still raises TypeError, now saying to give theta as a Scalar of its UDT; an op without is_udt raises KeyError for an array theta, as for a UDT Scalar. - A JIT op (indexbinary.ss.register_new) is registered for its x, y and theta types. op(theta) took the one registered at theta's type T for x, y and theta, and cast the inputs to T -> it takes the one for the inputs' types and theta's, without casting, as a JIT IndexUnaryOp takes its thunk: with an op registered at FP64, op(10.0) on FP64 data runs as before, op(10.0) on FP32 data ran in FP64 -> KeyError naming (FP32, FP32, FP64), and op(10) raised KeyError when bound -> when it runs. op.types is keyed (x, y, theta), so registering the same x and theta types for another y no longer raises "already defined", and op["FP64", "FP32"] is x and y, no longer x and theta: op["FP64", "FP64", "FP32"] is the op with an FP32 theta. - A bound op pickles with all its types: op["INT64", "FP64"](1) came back as op["INT64"](1), whose y is INT64. - op(theta).type: the bound type -> AttributeError, as for any untyped op; op(theta)[T] is op[T](theta). repr(op(2)): binary.f[INT64] -> indexbinary.f(2). op(theta)(A & B) works, as for other BinaryOps. - A Scalar theta is copied when given, so a later change to it does not reach the op, as when the op bound it at once. find_opclass now reads an opclass attribute on an OpBase subclass, so op(theta) is still a BinaryOp there. The user guide's examples drop the op[float](0.5) that hid the old typing. The comments on IndexUnaryOp's and SelectOp's UDT caches now say they are keyed by (dtype, thunk_type). --- docs/user_guide/operators.rst | 8 +- docs/user_guide/udf.rst | 52 +++++- graphblas/core/operator/base.py | 3 +- graphblas/core/operator/indexbinary.py | 233 ++++++++++++++++++------ graphblas/core/operator/indexunary.py | 4 +- graphblas/core/operator/select.py | 4 +- graphblas/core/operator/semiring.py | 29 +-- graphblas/core/operator/udf.py | 16 +- graphblas/core/ss/indexbinary.py | 21 ++- graphblas/tests/test_indexbinary.py | 241 ++++++++++++++++++++++++- graphblas/tests/test_ssjit.py | 19 +- 11 files changed, 520 insertions(+), 110 deletions(-) diff --git a/docs/user_guide/operators.rst b/docs/user_guide/operators.rst index 59a38c791..fb827afe3 100644 --- a/docs/user_guide/operators.rst +++ b/docs/user_guide/operators.rst @@ -232,7 +232,8 @@ receives both values, the indices of *each* value, and a thunk parameter: There are no built-in IndexBinary operators; all are user-defined (see :doc:`udf`). Binding a ``theta`` value to an IndexBinaryOp produces a regular BinaryOp usable directly in ``ewise_mult`` / ``ewise_add``, or as -the multiplier of a Semiring for ``mxm`` / ``mxv`` / ``vxm``. +the multiplier of a Semiring for ``mxm`` / ``mxv`` / ``vxm``. Like other +BinaryOps, its type is chosen when it runs, from both inputs and ``theta``. Example usage: @@ -243,11 +244,10 @@ Example usage: gb.indexbinary.register_new("discounted_sum", discounted_sum) - binop = gb.indexbinary.discounted_sum[float](0.5) # bind theta + binop = gb.indexbinary.discounted_sum(0.5) # bind theta C << A.ewise_mult(B, binop) - # For mxm, wrap in a Semiring; the additive monoid must accept the - # bound op's return type. + # For mxm, wrap in a Semiring sr = gb.semiring.register_anonymous(gb.monoid.plus, binop) D << A.mxm(B, sr) diff --git a/docs/user_guide/udf.rst b/docs/user_guide/udf.rst index 670cbc6e3..6255c1c56 100644 --- a/docs/user_guide/udf.rst +++ b/docs/user_guide/udf.rst @@ -125,11 +125,43 @@ works in ``ewise_mult``, ``ewise_add``, and the other elementwise paths: indexbinary.register_new("discounted_sum", discounted_sum) - bound = indexbinary.discounted_sum[float](0.5) # theta = 0.5 - A = Matrix.from_coo([0, 1], [0, 1], [1.0, 2.0]) - B = Matrix.from_coo([0, 1], [0, 1], [3.0, 4.0]) + bound = indexbinary.discounted_sum(0.5) # theta = 0.5 + A = Matrix.from_coo([0, 1], [0, 1], [1, 2]) + B = Matrix.from_coo([0, 1], [0, 1], [3, 4]) C = A.ewise_mult(B, bound).new() - # C[0,0] = (1+3)*0.5 = 2; C[1,1] = (2+4)*0.5 = 3 + # C is FP64: C[0,0] = (1+3)*0.5 = 2.0; C[1,1] = (2+4)*0.5 = 3.0 + +Like any other BinaryOp, the bound op takes its type each time it runs: the +type that holds both inputs and ``theta``, as ``apply`` types an IndexUnaryOp +from its input and thunk. One type serves x, y and ``theta``. ``theta`` has +the type of its value: a Python int is INT64 and a float FP64, and a numpy +scalar or a ``Scalar`` keeps its own type. With no ``theta``, ``theta`` is 0 +(``False``), which changes no type. Above, the INT64 matrices and the float +``theta`` run in FP64. To choose the type yourself, name it: +``indexbinary.discounted_sum[float](0.5)``, or +``indexbinary.discounted_sum(0.5, dtype=float)``. The inputs and ``theta`` are +then cast to that type. + +An op registered with ``is_udt=True`` is compiled for the types of x, y and +``theta`` as they are, so none of them is cast: ``op(0.5)`` on INT64 inputs +gets the float 0.5. A numpy array ``theta`` has the array UDT of its dtype +and shape. A tuple has no type of its own, so give a record or array +``theta`` as a ``Scalar`` of its UDT (``Scalar.from_value((1, 2.5), udt)``). +``op[X, Y, T](theta)`` names all three types, and ``op[X, Y](theta)`` gives +``theta`` the type of y. + +This differs from an IndexUnaryOp in ``apply``, which makes a number thunk an +element of the UDT data's type. Beside UDT data, a number ``theta`` stays a +number: with ``x + theta`` on INT32[3] data, ``op(5)`` adds an INT64 to each +element, which gives an int64 array, not an INT32[3] element, so the op +raises ``UdfParseError`` when it is used. To make ``theta`` the data's type, +pass a ``Scalar`` of that type, or a numpy value of the element type +(``op(np.int32(5))``), or name the type: ``op(5, dtype=udt)`` binds +``op[udt](5)``, which makes 5 an element of ``udt``. + +A ``Scalar`` ``theta`` is copied when it is bound, so changing the ``Scalar`` +afterwards does not change the bound op. An ``apply`` expression, by +contrast, reads its ``Scalar`` thunk when it runs. To use a bound IndexBinaryOp as the multiplier in ``mxm`` / ``mxv`` / ``vxm``, wrap it in a Semiring. ``Semiring.register_anonymous`` accepts a @@ -142,12 +174,12 @@ bound IBO directly: sr = semiring.register_anonymous(monoid.plus, bound) C = A.mxm(B, sr).new() -The resulting Semiring is monomorphic in the bound IBO's input/output types -(SuiteSparse builds exactly one ``GrB_Semiring`` for that type pair, rather -than the type-polymorphic table the standard semirings carry). To reuse the -same IBO at a different type, bind theta again under that type and build a -new Semiring. Per SuiteSparse, monoids themselves cannot be built from an -IndexBinaryOp; only the multiplier slot accepts one. +The Semiring is typed each time it runs, as the bound op is. One built from +an op bound at a type (``discounted_sum[float](0.5)``) is monomorphic in that +op's input/output types instead (SuiteSparse builds exactly one +``GrB_Semiring`` for that type pair, rather than the type-polymorphic table +the standard semirings carry). Per SuiteSparse, monoids themselves cannot be +built from an IndexBinaryOp; only the multiplier slot accepts one. IndexBinaryOps require SuiteSparse:GraphBLAS >= 9.4. diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index 536d9aaf3..aea4a6f51 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -603,7 +603,8 @@ def _check_supports_udf(cls, method_name): def find_opclass(gb_op): if isinstance(gb_op, OpBase): - opclass = type(gb_op).__name__ + # A subclass may name the class it is (``ibo(theta)`` is a BinaryOp). + opclass = getattr(gb_op, "opclass", type(gb_op).__name__) elif isinstance(gb_op, TypedOpBase): opclass = gb_op.opclass elif isinstance(gb_op, ParameterizedUdf): diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index c987a64e6..991d8f0f4 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -1,12 +1,16 @@ import inspect +from operator import getitem from types import FunctionType +import numpy as np + from ... import _STANDARD_OPERATOR_NAMES, indexbinary from ...dtypes import BOOL, INT8, UINT64, lookup_dtype from ...exceptions import UdfParseError, check_status_carg from .. import _has_numba, ffi, lib -from ..dtypes import _sample_values +from ..dtypes import _sample_values, unify from .base import OpBase, ParameterizedUdf, TypedOpBase +from .binary import BinaryOp _has_idxbinop = hasattr(lib, "GxB_IndexBinaryOp_new") @@ -24,6 +28,7 @@ def _rebind_indexbinaryop(parent, type_, theta): + # ``type_`` is one type, or ``(x, y, theta)`` types when they differ. # ``theta`` was stored as a raw Python/numpy value. For UDT thunks, # ``TypedBuiltinIndexBinaryOp.__call__`` would route a raw value through # ``Scalar.from_value(theta)`` with no dtype, which then fails inside the @@ -42,26 +47,36 @@ def _theta_to_scalar(theta, dtype): """Wrap a raw (non-Scalar) theta as a GrB Scalar for binding. ``dtype`` is the thunk dtype when known (a UDT thunk or an explicit - ``dtype=``), else ``None`` to infer from the value. A raw user-defined-type - value (array/tuple) can't be inferred, so raise a clear error that names the - value and the working alternatives. + ``dtype=``), else ``None`` to type theta by its own value: a Python + number by its default type, a numpy scalar or ``np.void`` by its dtype, + and a numpy array by its array layout (``INT32[3]`` for an int32 array + of shape ``(3,)``). A tuple, list or dict has no type of its own, so + raise a clear error that names the value and the working alternatives. """ from ..scalar import Scalar + if dtype is None and isinstance(theta, np.ndarray) and theta.ndim > 0: + try: + dtype = lookup_dtype(np.dtype((theta.dtype, theta.shape))) + except (TypeError, ValueError): + pass # a bool array, for one, is no UDT; ``from_value`` says so try: return Scalar.from_value(theta, dtype, is_cscalar=False, name="") # pragma: is_grbscalar except TypeError as exc: if dtype is None: raise TypeError( - f"Cannot infer a dtype for theta from {theta!r}; pass an explicit dtype " - "(e.g., op(value, dtype=...)) or a typed Scalar. Required for user-defined types." + f"Cannot infer a dtype for theta from {theta!r}; give theta as a Scalar of " + "its type (Scalar.from_value(theta, dtype)), or name the op's type " + "(op[dtype](theta), or op(theta, dtype=...))." ) from exc raise class _BoundIndexBinaryOp(TypedOpBase): - """A BinaryOp produced by binding a theta value to an IndexBinaryOp. + """A BinaryOp produced by binding a theta value to a typed IndexBinaryOp. + ``ibo[dtype](theta)`` returns one, and so does ``ibo(theta)`` each time it + runs, for the type it runs in (:class:`_UntypedBoundIndexBinaryOp`). Under the hood this is a single typed ``GrB_BinaryOp`` (constructed via ``GxB_BinaryOp_new_IndexOp``), so SuiteSparse accepts it anywhere a typed ``GrB_BinaryOp`` is expected: ``ewise_mult`` / ``ewise_add`` and, @@ -69,7 +84,7 @@ class _BoundIndexBinaryOp(TypedOpBase): Semiring that then works in ``mxm`` / ``mxv`` / ``vxm``. """ - __slots__ = ("_theta",) + __slots__ = "_theta", "_theta_type" opclass = "BinaryOp" # Each ``iop(theta)`` call allocates a fresh ``GrB_BinaryOp`` via # ``GxB_BinaryOp_new_IndexOp`` and bound IBOs are never cached, so this @@ -99,7 +114,86 @@ class _BoundIndexBinaryOp(TypedOpBase): _custom_dtype = None def __reduce__(self): - return (_rebind_indexbinaryop, (self.parent, self.type, self._theta)) + return ( + _rebind_indexbinaryop, + (self.parent, _ibo_types_key(self), self._theta), + ) + + +class _UntypedBoundIndexBinaryOp(BinaryOp): + """A BinaryOp produced by binding a theta value to an IndexBinaryOp, typed when it runs. + + ``ibo(theta)`` returns one. As with any other BinaryOp, its inputs choose + its type each time it runs. Theta is typed when it is given, as an + IndexUnaryOp's thunk is typed by its value: a Python int is INT64, a float + FP64 and a bool BOOL; a numpy scalar, a numpy array (as an array UDT) or a + Scalar keeps its own type. The op runs at the type that holds both inputs + and theta, and the default theta, False, changes no type. An IndexBinaryOp + registered with ``is_udt=True`` instead runs as ``ibo[X, Y, T]`` for inputs + of types X and Y and theta of type T, compiled for those three types, so + theta is never converted. + + It runs the typed op ``ibo[...](theta)``, built once per pair of input types. + """ + + __slots__ = "_ibo", "_theta" + # For ``find_opclass``, where the class name would not say BinaryOp. + opclass = "BinaryOp" + + def __init__(self, ibo, theta): + # ``is_udt=True`` sends every lookup (``get_typed_op``, a Semiring's) to + # ``_compile_udt`` below, UDT or not. + super().__init__(ibo.name, anonymous=True, is_udt=True) + self._ibo = ibo + self._theta = theta + + def _compile_udt(self, dtype, dtype2): + if dtype2 is None: # pragma: no cover (safety) + dtype2 = dtype + dtypes = (dtype, dtype2) + if dtypes in self._udt_ops: + return self._udt_ops[dtypes] + op = self._ibo._get_typed(dtype, dtype2, self._theta.dtype)(self._theta) + self._udt_types[dtypes] = op.return_type + self._udt_ops[dtypes] = op + return op + + def __getitem__(self, type_): + # A type named here is the op's type, as in ``ibo[type_](theta)``. + return self._ibo[type_](self._theta) + + def __reduce__(self): + return (self._ibo, (self._theta,)) + + def __repr__(self): + return f"{self._ibo!r}({self._theta.value!r})" + + +def _ibo_types_key(op): + """The type, or ``(x, y, theta)`` types when they differ, that index ``op``'s parent.""" + dtype2 = op.type if op._type2 is None else op._type2 + if op.type == dtype2 == op._theta_type: + return op.type + return (op.type, dtype2, op._theta_type) + + +class _TypedIndexBinaryOpMixin: + """The theta type, repr and pickle of a typed IndexBinaryOp, whose theta has its own type.""" + + __slots__ = () + + @property + def thunk_type(self): + return self._theta_type + + def __repr__(self): + dtype2 = self.type if self._type2 is None else self._type2 + if self._theta_type == dtype2: + return super().__repr__() + return f"indexbinary.{self.name}[{self.type.name}, {dtype2.name}, {self._theta_type.name}]" + + def __reduce__(self): + return (getitem, (self.parent, _ibo_types_key(self))) class TypedBuiltinIndexBinaryOp(TypedOpBase): @@ -116,7 +210,8 @@ def __call__(self, theta=None): Parameters ---------- theta : scalar, optional - The theta parameter to bind. Defaults to 0 (False). + The theta parameter to bind. Defaults to 0 (False). Its value is bound + here, so changing a Scalar theta afterwards does not change the BinaryOp. Returns ------- @@ -156,20 +251,20 @@ def __call__(self, theta=None): dtype2=self._type2, ) rv._theta = theta_value + rv._theta_type = self.thunk_type return rv - @property - def thunk_type(self): - return self.type if self._type2 is None else self._type2 - -class TypedUserIndexBinaryOp(TypedOpBase): - __slots__ = () +class TypedUserIndexBinaryOp(_TypedIndexBinaryOpMixin, TypedOpBase): + __slots__ = ("_theta_type",) opclass = "IndexBinaryOp" _owns_gb_obj = True - def __init__(self, parent, name, type_, return_type, gb_obj, dtype2=None): + def __init__(self, parent, name, type_, return_type, gb_obj, dtype2=None, theta_type=None): super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}", dtype2=dtype2) + if theta_type is None: + theta_type = type_ if dtype2 is None else dtype2 + self._theta_type = theta_type @property def orig_func(self): @@ -179,8 +274,6 @@ def orig_func(self): def _numba_func(self): return self.parent._numba_func - thunk_type = TypedBuiltinIndexBinaryOp.thunk_type - def __call__(self, theta=None): return TypedBuiltinIndexBinaryOp.__call__(self, theta) @@ -211,11 +304,14 @@ class IndexBinaryOp(OpBase): where ``ix, jx`` are the row and column indices of ``x``, and ``iy, jy`` are the row and column indices of ``y``. - Binding a theta value (``ibo[dtype](theta)``) produces a BinaryOp usable - directly in ``ewise_mult`` and ``ewise_add``. To use it as a multiplier in - ``mxm`` / ``mxv`` / ``vxm``, wrap it in a Semiring via - :meth:`Semiring.register_anonymous`; per SuiteSparse, the additive monoid - itself cannot be IndexBinaryOp-based. + Binding a theta value (``ibo(theta)``) produces a BinaryOp usable + directly in ``ewise_mult`` and ``ewise_add``. Its type is chosen when it + runs, from both inputs and theta; ``ibo[dtype](theta)`` binds to one type. + With ``is_udt=True``, x, y and theta each keep their own type, and + ``ibo[X, Y, T](theta)`` names all three. + To use it as a multiplier in ``mxm`` / ``mxv`` / ``vxm``, wrap it in a + Semiring via :meth:`Semiring.register_anonymous`; per SuiteSparse, the + additive monoid itself cannot be IndexBinaryOp-based. IndexBinaryOps live in the ``graphblas.indexbinary`` namespace. There are no built-ins; all IndexBinaryOps are user-defined. @@ -361,7 +457,7 @@ def indexbinary_wrapper( return new_type_obj raise UdfParseError("Unable to parse function using Numba") - def _compile_udt(self, dtype, dtype2): + def _compile_udt(self, dtype, dtype2, theta_type=None): if not _has_idxbinop: # KeyError (not RuntimeError) so ``udt in op`` and the resolver # chain in ``__contains__`` / ``[dtype]`` lookups return cleanly. @@ -371,12 +467,15 @@ def _compile_udt(self, dtype, dtype2): ) if dtype2 is None: dtype2 = dtype - dtypes = (dtype, dtype2) + if theta_type is None: + # ``ibo[X, Y]`` and ``ibo[T]``: theta takes y's type, and casts to it + theta_type = dtype2 + dtypes = (dtype, dtype2, theta_type) if dtypes in self._udt_types: return self._udt_ops[dtypes] from ..udt import _refuse_half - _refuse_half(self, dtype, dtype2) + _refuse_half(self, *dtypes) if self._numba_func is None: raise KeyError(f"{self.name} does not work with {dtypes} types") @@ -388,15 +487,15 @@ def _compile_udt(self, dtype, dtype2): dtype2.numba_type, UINT64.numba_type, UINT64.numba_type, - dtype2.numba_type, + theta_type.numba_type, ) _compile_udf_for_udt( - numba_func, sig, op_kind="indexbinary", op_name=self.name, dtypes=(dtype, dtype2) + 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, dtype, dtype2) + ret_type = _resolve_udt_return_type(numba_ret_type, *dtypes) indexbinary_wrapper, wrapper_sig = _get_udt_wrapper_indexbinary( - numba_func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type + numba_func, ret_type, dtype, dtype2, theta_type, numba_ret_type=numba_ret_type ) indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( @@ -410,7 +509,7 @@ def _compile_udt(self, dtype, dtype2): ret_type._carg, dtype._carg, dtype2._carg, - dtype2._carg, + theta_type._carg, ffi_new("char[]", self.name.encode()), ffi.NULL, ), @@ -424,11 +523,38 @@ def _compile_udt(self, dtype, dtype2): ret_type, new_idxbinop[0], dtype2=dtype2, + theta_type=theta_type, ) self._udt_types[dtypes] = ret_type self._udt_ops[dtypes] = op return op + def _get_typed(self, dtype, dtype2, theta_type): + """The typed op for x, y and theta of these types, as ``ibo(theta)`` runs it. + + With ``is_udt``, it is compiled for the three types as they are. Without + it, x, y and theta share one type, the one that holds all three. + """ + if self._is_udt: + return self._compile_udt(dtype, dtype2, theta_type) + if dtype._is_udt or dtype2._is_udt or theta_type._is_udt: + raise KeyError(f"{self.name} does not work with ({dtype}, {dtype2}, {theta_type})") + return self[unify(unify(dtype, dtype2), theta_type)] + + def __getitem__(self, type_): + if type(type_) is tuple and len(type_) == 3: + return self._get_typed(*(lookup_dtype(dtype) for dtype in type_)) + return super().__getitem__(type_) + + def _add(self, op, *, is_jit=False): + if not is_jit: + return super()._add(op) + # A JIT op is registered for its x, y and theta types + dtypes = (op.type, op._type2, op.thunk_type) + self.types[dtypes] = op.return_type + self._udt_types[dtypes] = op.return_type + self._udt_ops[dtypes] = op + @classmethod def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): """Register an IndexBinaryOp without adding it to the ``indexbinary`` namespace. @@ -520,8 +646,8 @@ def __init__(self, name, func=None, *, anonymous=False, is_udt=False, numba_func self._numba_func = numba_func self._is_udt = is_udt if is_udt: - self._udt_types = {} # {(dtype, dtype2): DataType} - self._udt_ops = {} # {(dtype, dtype2): TypedUserIndexBinaryOp} + self._udt_types = {} # {(dtype, dtype2, theta_type): DataType} + self._udt_ops = {} # {(dtype, dtype2, theta_type): TypedUserIndexBinaryOp} def __call__(self, theta=None, dtype=None): """Bind a theta value to produce a BinaryOp. @@ -529,13 +655,18 @@ def __call__(self, theta=None, dtype=None): Parameters ---------- theta : scalar, optional - The theta parameter to bind. Defaults to 0 (False). + The theta parameter to bind. Defaults to 0 (False). Its value is bound + here, so changing a Scalar theta afterwards does not change the BinaryOp. dtype : dtype, optional - The dtype to use. Inferred from ``theta`` when not provided. + The type of the BinaryOp, as in ``op[dtype](theta)``. Without it, the + BinaryOp's type is chosen each time it runs, from the types of both + inputs and theta, as other BinaryOps choose theirs. Theta has the type + of its value: a Python int is INT64 and a float FP64, and a numpy + scalar, a numpy array or a Scalar keeps its own type. Returns ------- - TypedOpBase + BinaryOp A BinaryOp built from this IndexBinaryOp with the given theta. """ @@ -544,22 +675,20 @@ def __call__(self, theta=None, dtype=None): if theta is None: theta = False - if not isinstance(theta, Scalar): - # An explicit dtype supplies a UDT; with no dtype the value is - # inferred (and a raw UDT value raises a clear error). - theta = _theta_to_scalar(theta, _lookup_dtype(dtype) if dtype is not None else None) - elif theta._is_cscalar: - # fmt: off - # Pass dtype explicitly; an array-UDT value would otherwise infer wrong. - val, dt = theta.value, theta.dtype - theta = Scalar.from_value(val, dt, is_cscalar=False, name="") # pragma: is_grbscalar - # fmt: on if dtype is None: - dtype = theta.dtype - else: - dtype = _lookup_dtype(dtype) - typed_op = self[dtype] - return typed_op(theta) + if isinstance(theta, Scalar): + # Bound when it runs, so keep the value theta has now. + theta = theta.dup(name="") + else: + theta = _theta_to_scalar(theta, None) + if not self._is_udt and theta.dtype._is_udt: + raise KeyError(f"{self.name} does not work with {theta.dtype}") + return _UntypedBoundIndexBinaryOp(self, theta) + dtype = _lookup_dtype(dtype) + if not isinstance(theta, Scalar): + # The named dtype types the value, a UDT's included. + theta = _theta_to_scalar(theta, dtype) + return self[dtype](theta) ParameterizedIndexBinaryOp._op_class = IndexBinaryOp diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index 0a7437def..a8eca2460 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -405,8 +405,8 @@ def __init__( self.is_positional = is_positional self._is_udt = is_udt if is_udt: - self._udt_types = {} # {dtype: DataType} - self._udt_ops = {} # {dtype: TypedUserIndexUnaryOp} + self._udt_types = {} # {(dtype, thunk_type): DataType} + self._udt_ops = {} # {(dtype, thunk_type): TypedUserIndexUnaryOp} __call__ = TypedBuiltinIndexUnaryOp.__call__ diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 2be05fb4c..1c1138ae5 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -432,8 +432,8 @@ def __init__( self.is_positional = is_positional self._is_udt = is_udt if is_udt: - self._udt_types = {} # {dtype: DataType} - self._udt_ops = {} # {dtype: TypedUserIndexUnaryOp} + self._udt_types = {} # {(dtype, thunk_type): DataType} + self._udt_ops = {} # {(dtype, thunk_type): TypedUserSelectOp} __call__ = TypedBuiltinSelectOp.__call__ diff --git a/graphblas/core/operator/semiring.py b/graphblas/core/operator/semiring.py index b333259e3..f9d8603db 100644 --- a/graphblas/core/operator/semiring.py +++ b/graphblas/core/operator/semiring.py @@ -224,17 +224,18 @@ def _build(cls, name, monoid, binaryop, *, anonymous=False): raise TypeError(f"monoid must be a Monoid, not {type(monoid)}") # SuiteSparse:GraphBLAS allows the multiplier of a Semiring to be # any GrB_BinaryOp, "including those based on a GxB_IndexBinaryOp" - # (GraphBLAS.h, GrB_Semiring section). A ``_BoundIndexBinaryOp`` is - # already exactly such a typed GrB_BinaryOp (built via - # ``GxB_BinaryOp_new_IndexOp``), so accept it directly and build - # one Semiring for its specific input/output type pair, rather - # than the regular per-typed-op iteration that polymorphic - # ``BinaryOp`` parents go through. + # (GraphBLAS.h, GrB_Semiring section). A ``_BoundIndexBinaryOp`` + # (``ibo[dtype](theta)``) is already exactly such a typed + # GrB_BinaryOp (built via ``GxB_BinaryOp_new_IndexOp``), so accept it + # directly and build one Semiring for its specific input/output type + # pair, rather than the regular per-typed-op iteration that + # polymorphic ``BinaryOp`` parents go through. ``ibo(theta)`` is such + # a parent, a BinaryOp subclass that binds theta per type as it runs. from .indexbinary import _BoundIndexBinaryOp if isinstance(binaryop, _BoundIndexBinaryOp): return cls._build_from_bound_indexbinary(name, monoid, binaryop, anonymous=anonymous) - if type(binaryop) is not BinaryOp: + if not isinstance(binaryop, BinaryOp): raise TypeError( f"binaryop must be a BinaryOp or a bound IndexBinaryOp " f"(``ibo[dtype](theta)``), not {type(binaryop)}" @@ -351,9 +352,10 @@ def register_anonymous(cls, monoid, binaryop, name=None): The monoid of the semiring (like "plus" in the default "plus_times" semiring). binaryop : BinaryOp, ParameterizedBinaryOp, or bound IndexBinaryOp The multiplier of the semiring (like "times" in the default - "plus_times" semiring). A bound IndexBinaryOp (``ibo[dtype](theta)``) - is accepted as well; the resulting semiring is monomorphic in - the bound op's input/output types. + "plus_times" semiring). An IndexBinaryOp bound to theta, + ``ibo(theta)``, is a BinaryOp. One bound at a type, + ``ibo[dtype](theta)``, is accepted as well; the resulting semiring + is monomorphic in the bound op's input/output types. name : str, optional The name of the operator. This *does not* show up as ``gb.semiring.{name}``. @@ -379,9 +381,10 @@ def register_new(cls, name, monoid, binaryop, *, lazy=False): monoid : Monoid or ParameterizedMonoid The monoid of the semiring (like "plus" in the default "plus_times" semiring). binaryop : BinaryOp, ParameterizedBinaryOp, or bound IndexBinaryOp - The multiplier of the semiring. A bound IndexBinaryOp - (``ibo[dtype](theta)``) is also accepted; the resulting - semiring is monomorphic in the bound op's input/output types. + The multiplier of the semiring. ``ibo(theta)`` is a BinaryOp, and + an IndexBinaryOp bound at a type, ``ibo[dtype](theta)``, is also + accepted; the resulting semiring is monomorphic in the bound op's + input/output types. lazy : bool, default False If False (the default), then the function will be automatically compiled for builtin types (unless ``is_udt`` is True). diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index 9eeae225d..ab8144fc2 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -640,17 +640,19 @@ def _get_udt_wrapper( return wrapper, wrapper_sig -def _get_udt_wrapper_indexbinary(numba_func, return_type, dtype, dtype2, *, numba_ret_type=None): +def _get_udt_wrapper_indexbinary( + numba_func, return_type, dtype, dtype2, theta_type, *, numba_ret_type=None +): """Build a Numba cfunc wrapper for IndexBinaryOp UDFs on UDTs. - Signature: ``f(x, ix, jx, y, iy, jy, theta) -> z``. ``dtype2`` is the - shared type of ``y`` and ``theta``. + Signature: ``f(x, ix, jx, y, iy, jy, theta) -> z``, where ``x``, ``y`` + and ``theta`` have types ``dtype``, ``dtype2`` and ``theta_type``. """ nt = numba.types zsetup, zptr_type, zkind, zinfo = _output_handler(return_type, numba_ret_type) xsetup, xderef, xptr_type = _input_operand(dtype, "x") ysetup, yderef, yptr_type = _input_operand(dtype2, "y") - tsetup, tderef, tptr_type = _input_operand(dtype2, "t") + tsetup, tderef, tptr_type = _input_operand(theta_type, "t") wrapper_sig = nt.void( zptr_type, xptr_type, @@ -661,7 +663,7 @@ def _get_udt_wrapper_indexbinary(numba_func, return_type, dtype, dtype2, *, numb UINT64.numba_type, tptr_type, ) - arg_dtypes = [dtype, UINT64, UINT64, dtype2, UINT64, UINT64, dtype2] + arg_dtypes = [dtype, UINT64, UINT64, dtype2, UINT64, UINT64, theta_type] _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes) signature_line = "def wrapper(z_ptr, x_ptr, ix, jx, y_ptr, iy, jy, t_ptr):" @@ -672,7 +674,9 @@ def _get_udt_wrapper_indexbinary(numba_func, return_type, dtype, dtype2, *, numb wrapper = _compile_codegen( text, func_name="wrapper", - source_label=f"", + source_label=( + f"" + ), extra_ns={"numba_func": numba_func}, ) return wrapper, wrapper_sig diff --git a/graphblas/core/ss/indexbinary.py b/graphblas/core/ss/indexbinary.py index 85b3253e0..e34451b0a 100644 --- a/graphblas/core/ss/indexbinary.py +++ b/graphblas/core/ss/indexbinary.py @@ -6,27 +6,28 @@ from ..operator.indexbinary import ( IndexBinaryOp, TypedBuiltinIndexBinaryOp, - TypedUserIndexBinaryOp, _has_idxbinop, + _TypedIndexBinaryOpMixin, ) ffi_new = ffi.new -class TypedJitIndexBinaryOp(TypedOpBase): - __slots__ = "_jit_c_definition" +class TypedJitIndexBinaryOp(_TypedIndexBinaryOpMixin, TypedOpBase): + __slots__ = "_jit_c_definition", "_theta_type" opclass = "IndexBinaryOp" - def __init__(self, parent, name, type_, return_type, gb_obj, jit_c_definition, dtype2=None): + def __init__( + self, parent, name, type_, return_type, gb_obj, jit_c_definition, dtype2, theta_type + ): super().__init__(parent, name, type_, return_type, gb_obj, name, dtype2=dtype2) self._jit_c_definition = jit_c_definition + self._theta_type = theta_type @property def jit_c_definition(self): return self._jit_c_definition - thunk_type = TypedUserIndexBinaryOp.thunk_type - def __call__(self, theta=None): return TypedBuiltinIndexBinaryOp.__call__(self, theta) @@ -101,13 +102,13 @@ def register_new(name, jit_c_definition, x_type, y_type, theta_type, ret_type): if not isinstance(rv, IndexBinaryOp): IndexBinaryOp._remove_nesting(name) if ( - (x_type, theta_type) in rv.types + (x_type, y_type, theta_type) in rv.types or rv._udt_types is not None - and (x_type, theta_type) in rv._udt_types + and (x_type, y_type, theta_type) in rv._udt_types ): raise TypeError( f"IndexBinaryOp gb.indexbinary.{name} already defined for " - f"({x_type}, {theta_type}) input types" + f"({x_type}, {y_type}, {theta_type}) input types" ) else: rv = IndexBinaryOp(name, is_udt=True) @@ -130,7 +131,7 @@ def register_new(name, jit_c_definition, x_type, y_type, theta_type, ret_type): gb_obj[0], ) op = TypedJitIndexBinaryOp( - rv, funcname, x_type, ret_type, gb_obj[0], jit_c_definition, dtype2=theta_type + rv, funcname, x_type, ret_type, gb_obj[0], jit_c_definition, y_type, theta_type ) rv._add(op, is_jit=True) setattr(module, funcname, rv) diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py index 19b65d720..83acac65a 100644 --- a/graphblas/tests/test_indexbinary.py +++ b/graphblas/tests/test_indexbinary.py @@ -83,6 +83,68 @@ def add_theta(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) assert list(C.to_coo()[2]) == [18, 19] +def test_untyped_call_is_typed_when_it_runs(): + """``op(theta)`` takes its type from both inputs and theta each time it runs. + + Theta is typed as an IndexUnaryOp's thunk is: a Python int is INT64 and a + float FP64, a numpy scalar or a Scalar keeps its type, and the default, + False, changes no type. Before, the op had theta's type alone: ``op()`` + was BOOL, so ints became True; ``op(2)`` truncated floats to INT64; and + ``op(np.int8(2))`` computed in INT8, where 300.0 became 127 and wrapped. + """ + + def upper_from_x(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x if ix <= jx else y + + pick = indexbinary.register_anonymous(upper_from_x) + A = Matrix.from_coo([0, 0, 1, 1], [0, 1, 0, 1], [10, 20, 30, 40]) + B = Matrix.from_coo([0, 0, 1, 1], [0, 1, 0, 1], [-1, -2, -3, -4]) + expected = Matrix.from_coo([0, 0, 1, 1], [0, 1, 0, 1], [10, 20, -3, 40]) + for binop in [pick(), pick(0), pick(False)]: + assert A.ewise_mult(B, binop).new().isequal(expected, check_dtype=True) + C = A.ewise_mult(B, pick(0.0)).new() + assert C.isequal(expected.dup("FP64"), check_dtype=True) + + def weighted(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + theta * y + + wsum = indexbinary.register_anonymous(weighted) + u = Vector.from_coo([0, 1], [1.5, 300.0]) + w = Vector.from_coo([0, 1], [2.25, 1.0]) + expected = Vector.from_coo([0, 1], [6.0, 302.0]) + for theta in [2, 2.0, np.int8(2), Scalar.from_value(2, "INT8")]: + binop = wsum(theta) + assert u.ewise_mult(w, binop).new().isequal(expected, check_dtype=True) + assert binop(u & w).new().isequal(expected, check_dtype=True) + # One op runs at each type it meets + i8 = Vector.from_coo([0, 1], [1, 100], dtype="INT8") + binop = wsum(1) + assert u.ewise_mult(w, binop).new().isequal(Vector.from_coo([0, 1], [3.75, 301.0])) + result = i8.ewise_mult(i8, binop).new() + assert result.isequal(Vector.from_coo([0, 1], [2, 200], dtype="INT64"), check_dtype=True) + for binop in [wsum(), wsum(np.int8(1))]: + assert i8.ewise_mult(i8, binop).new().dtype == dtypes.INT8 + # A named type is the op's type, as before + for binop in [ + wsum["INT64"](2), + wsum(2, dtype="INT64"), + wsum(Scalar.from_value(2.0), dtype="INT64"), + wsum(2)["INT64"], + ]: + assert binop.type == dtypes.INT64 + result = u.ewise_mult(w, binop).new() + assert result.isequal(Vector.from_coo([0, 1], [5, 302]), check_dtype=True) + assert repr(wsum(2)) == "indexbinary.weighted(2)" + rec_udt = dtypes.register_anonymous(np.dtype([("theta_a", np.int64)]), "_UntypedThetaRec") + with pytest.raises(KeyError, match="weighted does not work with _UntypedThetaRec"): + wsum(Scalar.from_value((1,), rec_udt)) + # Binding when it runs uses theta's value from when it was given + theta = Scalar.from_value(2.0) + binop = wsum(theta) + theta.value = 100.0 + assert u.ewise_mult(w, binop).new().isequal(expected) + + def test_index_aware(): """The wrapper passes ``(ix, jx, iy, jy)`` through to the user function.""" @@ -326,19 +388,107 @@ def _combine(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) assert good[arr_udt].return_type is arr_udt -def test_bind_raw_udt_theta_without_dtype_errors(): - """A raw UDT theta with no dtype can't be inferred; the error is clear and actionable.""" - indexbinary.register_new("raw_no_dtype_op", _ibo_return_x, is_udt=True) +def test_bind_raw_udt_theta_without_dtype(): + """A numpy array theta is typed as its array layout, as other array literals are. + + With ``is_udt=True``, theta keeps that type, as any theta does, so it is + never converted to y's type. A tuple has no type of its own: give it as a + Scalar of its UDT, or name the op's type. An op without ``is_udt`` does + not take a UDT theta. + """ + + def theta_only(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return theta + + arr_udt = dtypes.register_anonymous(np.dtype((np.int32, (3,))), "_RawThetaArr3i32") + indexbinary.register_new("raw_no_dtype_op", theta_only, is_udt=True) try: op = indexbinary.raw_no_dtype_op + v = Vector.from_coo([0], np.array([[1, 2, 3]], dtype=np.int32), dtype=arr_udt) + for theta in [np.array([10, 20, 30], np.int32), Scalar.from_value((10, 20, 30), arr_udt)]: + result = v.ewise_mult(v, op(theta)).new() + assert result.dtype == arr_udt + assert result[0].new().value.tolist() == [10, 20, 30] + # An int64 or float64 array keeps its type, where it would have become int32 + result = v.ewise_mult(v, op(np.array([10, 20, 2**40]))).new() + assert result.dtype.np_type == np.dtype((np.int64, (3,))) + assert result[0].new().value.tolist() == [10, 20, 2**40] + result = v.ewise_mult(v, op(np.array([10, 20.5, 30]))).new() + assert result.dtype.np_type == np.dtype((np.float64, (3,))) + assert result[0].new().value.tolist() == [10, 20.5, 30] + with pytest.raises(TypeError, match=r"Cannot infer a dtype for theta from \(10, 20, 30\)"): + op((10, 20, 30)) with pytest.raises(TypeError, match="Cannot infer a dtype for theta"): - op(np.array([1, 2, 3], dtype=np.int64)) - with pytest.raises(TypeError, match="Cannot infer a dtype for theta"): - op((1, 2, 3)) + op(np.array([True, False])) + # Its array UDT is shared, under the name it got first (another test's) + with pytest.raises(KeyError, match="_ibo_add_theta does not work with "): + indexbinary.register_anonymous(_ibo_add_theta)(np.array([1, 2, 3])) finally: delattr(indexbinary, "raw_no_dtype_op") +def test_untyped_bind_theta_keeps_its_type_with_is_udt(): + """With ``is_udt=True``, ``op(theta)`` compiles for x, y and theta as they are typed. + + SuiteSparse gives theta a type of its own (``GxB_IndexBinaryOp_new`` takes + ``theta_type``), so theta is not converted: 0.5 stays 0.5 beside INT64 + inputs, and Numba types the result as it would for those three types. + Before, theta took y's type when the op ran. Naming two types, + ``op[X, Y](theta)``, still gives theta y's type and casts it, and + ``op[X, Y, T](theta)`` names all three. + """ + op = indexbinary.register_anonymous(_ibo_add_theta, is_udt=True) + v = Vector.from_coo([0, 1], [1, 2]) + expected = Vector.from_coo([0, 1], [4, 6]) + for theta in [2, np.int8(2), Scalar.from_value(2, "INT16")]: + assert v.ewise_mult(v, op(theta)).new().isequal(expected, check_dtype=True) + for theta in [2.0, np.float64(2), Scalar.from_value(2.0)]: + result = v.ewise_mult(v, op(theta)).new() + assert result.isequal(expected.dup(float), check_dtype=True) + result = v.ewise_mult(v, op(0.5)).new() + assert result.isequal(Vector.from_coo([0, 1], [2.5, 4.5]), check_dtype=True) + i8 = Vector.from_coo([0], [1], dtype="INT8") + result = i8.ewise_mult(i8, op(300)).new() + assert result.isequal(Vector.from_coo([0], [302]), check_dtype=True) + assert {key for key in op._udt_types if key[0] == dtypes.INT64} == { + (dtypes.INT64, dtypes.INT64, dtypes.INT64), + (dtypes.INT64, dtypes.INT64, dtypes.INT8), + (dtypes.INT64, dtypes.INT64, dtypes.INT16), + (dtypes.INT64, dtypes.INT64, dtypes.FP64), + } + # Named types + result = v.ewise_mult(v, op["INT64", "INT64"](0.5)).new() + assert result.isequal(Vector.from_coo([0, 1], [2, 4]), check_dtype=True) + typed = op["INT64", "INT64", "FP64"] + assert typed is op._udt_ops[dtypes.INT64, dtypes.INT64, dtypes.FP64] + assert typed.thunk_type == dtypes.FP64 + assert repr(typed) == "indexbinary._ibo_add_theta[INT64, INT64, FP64]" + assert repr(op["INT64", "INT64"]) == "indexbinary._ibo_add_theta[INT64, INT64]" + result = v.ewise_mult(v, typed(0.5)).new() + assert result.isequal(Vector.from_coo([0, 1], [2.5, 4.5]), check_dtype=True) + # Pickle keeps all three types, for the typed op and for the op it binds + typed2 = pickle.loads(pickle.dumps(typed)) # re-registers the anonymous op + assert repr(typed2) == repr(typed) + assert v.ewise_mult(v, typed2(0.5)).new().isequal(result, check_dtype=True) + bound = v.ewise_mult(v, op(0.5)).op + assert bound.type == bound._type2 == dtypes.INT64 + bound2 = pickle.loads(pickle.dumps(bound)) + assert bound2._theta_type == dtypes.FP64 + assert v.ewise_mult(v, bound2).new().isequal(result, check_dtype=True) + bound2 = pickle.loads(pickle.dumps(op["INT64", "FP64"](1))) # was op["INT64"](1) + assert (bound2.type, bound2._type2, bound2._theta_type) == ( + dtypes.INT64, + dtypes.FP64, + dtypes.FP64, + ) + # Without is_udt, the three types are one, as when the op runs + wsum = indexbinary.register_anonymous(_ibo_add_theta) + assert wsum["INT8", "INT8", "INT64"] is wsum["INT64"] + rec_udt = dtypes.register_anonymous(np.dtype([("theta_b", np.int64)]), "_UntypedThetaRec2") + with pytest.raises(KeyError, match="does not work with .INT64, INT64, _UntypedThetaRec2"): + wsum["INT64", "INT64", rec_udt] + + def test_pickle_bound_record_udt_theta(): """Bound IBO with a record-UDT theta value round-trips through pickle. @@ -508,6 +658,7 @@ def add_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) bound = typed(0) assert bound.opclass == "BinaryOp" + assert find_opclass(op(0))[1] == "BinaryOp" def test_udt_tuple_return(): @@ -715,6 +866,84 @@ def weighted(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) assert C[1, 1].new() == (25, 7) +def _ibo_decay(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + d = ix - jx if ix >= jx else jx - ix + return x * y * theta**d + + +def test_semiring_from_untyped_bound_ibo(): + """A Semiring from ``ibo(theta)`` is typed when it runs, by the inputs and theta. + + Before, ``ibo(0.5)`` was an FP64 op, so its semiring refused INT64 + matrices with KeyError, and ``ibo(1)`` was INT64 and refused FP64 ones. + """ + from graphblas import monoid as monoid_module + from graphblas import semiring as semiring_module + + decay = indexbinary.register_anonymous(_ibo_decay) + A = Matrix.from_coo([0, 0, 1, 2], [0, 2, 1, 0], [1, 1, 1, 1], nrows=3, ncols=3) + v = Vector.from_coo([0, 1, 2], [4, 4, 4]) + half = semiring_module.register_anonymous(monoid_module.plus, decay(0.5)) + one = semiring_module.register_anonymous(monoid_module.plus, decay(1)) + expected = Vector.from_coo([0, 1, 2], [5.0, 4.0, 1.0]) + assert A.mxv(v, half).new().isequal(expected, check_dtype=True) + assert A.dup("FP32").mxv(v.dup("FP32"), half).new().isequal(expected, check_dtype=True) + assert A.mxv(v, one).new().isequal(Vector.from_coo([0, 1, 2], [8, 4, 4]), check_dtype=True) + result = A.dup(float).mxv(v.dup(float), one).new() + assert result.isequal(Vector.from_coo([0, 1, 2], [8.0, 4.0, 4.0]), check_dtype=True) + sr2 = pickle.loads(pickle.dumps(half)) + assert A.mxv(v, sr2).new().isequal(expected, check_dtype=True) + # Theta keeps its type through pickle: FP32 data stays FP32 with an FP32 theta + bound2 = pickle.loads(pickle.dumps(decay(np.float32(0.5)))) + sr2 = semiring_module.register_anonymous(monoid_module.plus, bound2) + result = A.dup("FP32").mxv(v.dup("FP32"), sr2).new() + assert result.isequal(expected.dup("FP32"), check_dtype=True) + + +def test_argmax_semiring_from_untyped_bound_ibo(): + """The SuiteSparse user guide's row argmax, with theta unused, bound with ``iop()``. + + The multiplier turns ``A(i, k)`` into the pair ``(k, A(i, k))`` and a max + monoid on pairs keeps the larger value (the smaller column on a tie). The + guide binds a BOOL theta because theta cannot be NULL. Here the matrix + holds pairs, so the UDF's tuple return takes their type. Before, ``iop()`` + and ``iop(0)`` were ops on BOOL and INT64 values, not the matrix's pairs. + Theta keeps its own type, so each bind compiles the op for it. + """ + from graphblas import binary + from graphblas import monoid as monoid_module + from graphblas import semiring as semiring_module + + pair = dtypes.register_anonymous(np.dtype([("k", np.int64), ("v", np.float64)]), "IboPair") + + def make_pair(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return (np.int64(jx), x["v"]) + + def max_pair(a, b): # pragma: no cover (numba) + if a[1] > b[1] or (a[1] == b[1] and a[0] < b[0]): + return a + return b + + iop = indexbinary.register_anonymous(make_pair, is_udt=True) + max_monoid = monoid_module.register_anonymous( + binary.register_anonymous(max_pair, is_udt=True), (np.iinfo(np.int64).max, -np.inf) + ) + values = np.array([(0, 3.0), (0, 7.0), (0, 7.0), (0, 5.0), (0, -1.0)], dtype=pair.np_type) + A = Matrix.from_coo([0, 0, 0, 1, 1], [0, 1, 2, 0, 2], values, dtype=pair) + ones = Vector.from_scalar(True, A.ncols) + for binop in [iop(), iop(False), iop(0), iop(0.0)]: + argmax = semiring_module.register_anonymous(max_monoid, binop) + c = A.mxv(ones, argmax).new() + assert c.dtype == pair + assert c[0].new() == (1, 7.0) + assert c[1].new() == (0, 5.0) + assert set(iop._udt_types) == { + (pair, dtypes.BOOL, dtypes.BOOL), + (pair, dtypes.BOOL, dtypes.INT64), + (pair, dtypes.BOOL, dtypes.FP64), + } + + def test_dir_and_module(): assert "register_new" in dir(indexbinary) assert "register_anonymous" in dir(indexbinary) diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index ad32c5200..a935fded1 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -799,7 +799,7 @@ def test_jit_indexbinary(v): assert not hasattr(indexbinary, "add_theta") assert indexbinary.ss.add_theta is add_theta assert add_theta.name == "ss.add_theta" - assert add_theta.types == {(dtypes.FP64, dtypes.FP64): dtypes.FP64} + assert add_theta.types == {(dtypes.FP64, dtypes.FP64, dtypes.FP64): dtypes.FP64} assert "FP64" in add_theta assert add_theta["FP64"].return_type == dtypes.FP64 assert add_theta["FP64"].jit_c_definition == cdef @@ -812,6 +812,13 @@ def test_jit_indexbinary(v): # ewise_mult: x+y+theta = 2*val + 10 expected = Vector.from_coo([1, 3, 4, 6], [12.0, 12.0, 14.0, 10.0], dtype="FP64") assert expected.isequal(res) + # Bound with no type, it runs at the inputs' types and theta's, which a JIT op must + # have: like a JIT IndexUnaryOp's thunk, theta is not cast to the op's theta type. + assert expected.isequal(v64.ewise_mult(v64, add_theta(10.0)).new(), check_dtype=True) + with pytest.raises(KeyError, match="add_theta does not work with .FP64, FP64, INT64. types"): + v64.ewise_mult(v64, add_theta(10)) + with pytest.raises(KeyError, match="add_theta does not work with .INT64, INT64, FP64. types"): + v.ewise_mult(v, add_theta(10.0)) # Test duplicate registration fails assert "FP32" not in add_theta with burble(): @@ -827,7 +834,7 @@ def test_jit_indexbinary(v): assert "FP32" in add_theta with pytest.raises( TypeError, - match="IndexBinaryOp gb.indexbinary.ss.add_theta already defined for .FP64, FP64. input", + match="gb.indexbinary.ss.add_theta already defined for .FP64, FP64, FP64. input", ): indexbinary.ss.register_new("add_theta", cdef, "FP64", "FP64", "FP64", "FP64") # Test nested names @@ -844,8 +851,12 @@ def test_jit_indexbinary(v): "add_theta", mixed_cdef, "FP64", "FP64", "FP32", "FP64" ) assert add_theta_mixed is add_theta - assert ("FP64", "FP32") in add_theta - assert ("FP32", "FP64") not in add_theta + assert ("FP64", "FP64", "FP32") in add_theta + assert ("FP64", "FP32") not in add_theta # theta of y's type, FP32 + assert repr(add_theta["FP64", "FP64", "FP32"]) == "indexbinary.add_theta[FP64, FP64, FP32]" + # An FP32 theta picks it. Not computed here: SuiteSparse's JIT knows a user op's kernel + # by its name and x, y, z types, so in this process it would run the FP64 theta's kernel. + assert v64.ewise_mult(v64, add_theta(np.float32(10))).op._theta_type == dtypes.FP32 @pytest.mark.slow From 102f3a0cc15147aa7945f37024643617fe8ab776 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 29 Oct 2026 00:04:06 -0500 Subject: [PATCH 42/54] Refuse a record literal whose array field part does not broadcast to the field Behavior fix. Where a literal for a record UDT is typed (a lifted op such as +, ==, a user's op, an IndexUnaryOp thunk, an ewise_union default, and == on a Scalar), a part for an array field must now broadcast to the field's shape, or ValueError names both shapes. An array UDT's element takes the shape it broadcasts to with a literal, as numpy does, so INT8[3] plus np.array([[1, 300, 3]]) is INT64[1, 3]; that is unchanged. A record's field keeps its shape, so for the record the typing gave up, and the literal was converted into the field as numpy assigns it: the leading axis was dropped and 300 wrapped to 44, silently, even beside a user's op, which refuses a literal that does not fit. Stores (r[0] = ...) still convert as numpy assigns, as they do for an array UDT. With r a Vector of a record with an int8[3] field a and an int64 field b, holding ([1, 2, 3], 0) (before -> after): - r + (np.array([[1, 300, 3]]), 0): ([2, 46, 6], 0) -> ValueError "array([[ 1, 300, 3]]) has shape (1, 3), which does not fit field 'a' of R, of shape (3,): ...". - r.apply(user_op, right=(np.array([[1, 300, 3]]), 0)), with user_op returning y: ([1, 44, 3], 0) -> the same ValueError. - r == (np.array([[1, 300, 3]]), 0): False -> the same ValueError; with np.array([[1, 2, 3]], np.int8) it was True. - r + ([[1, 300, 3]], 0), and the other paths with it: numpy's "setting an array element with a sequence" ValueError -> this ValueError. - r + ([1, 2], 0): numpy's "could not broadcast input array from shape (2,) into shape (3,)" -> this ValueError. - r[0] = (np.array([[1, 300, 3]]), 0): ([1, 44, 3], 0), unchanged. --- docs/user_guide/udt.rst | 6 +++- graphblas/core/literals.py | 42 ++++++++++++++++++++++ graphblas/tests/test_op.py | 71 ++++++++++++++++++++++++++++++++++++-- 3 files changed, 116 insertions(+), 3 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 49dc44ef7..3f3793144 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -167,7 +167,11 @@ dtype, so it never narrows. Below, ``v`` is a Vector of an ``INT8[3]`` UDT. ``{x: FP64, y: FP64}`` * - A built-in Vector, ``Scalar``, numpy scalar or numpy array - Strong: keeps its dtype. A scalar combines with every field or element; - an array pairs as an array UDT of its own shape. + an array pairs as an array UDT of its own shape. A record's array field + keeps its shape, so a literal for it must broadcast to that shape, or + ``ValueError`` names both: ``np.array([[1, 2, 3]])`` does not fit an + ``INT8[3]`` field, though beside an ``INT8[3]`` UDT it gives + ``INT64[1, 3]``. - ``v + np.int64(1)`` is ``INT64[3]`` * - A Python number, or a tuple, list or dict of them laid out like an element diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index d7a24bc48..18dca86db 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -106,6 +106,8 @@ def _literal_operand(dtype, value, op): if is_comparison and not isinstance(value, (np.generic, np.ndarray)): np_type = dtype.np_type parts = _dict_to_record(np_type, value) if isinstance(value, dict) else value + # First: making an int beyond its leaf infinity reshapes the part. + _refuse_misfit_fields(dtype, np_type, parts) if _int_beyond_leaf(np_type, parts) is not None: # As infinity in its leaf, it still equals no element, and it # stays out of an array numpy makes of numpy values beside it. @@ -238,6 +240,7 @@ def _literal_type(dtype, value, *, exact=False): np_type = dtype.np_type if isinstance(value, dict) and np_type.names is not None: value = _dict_to_record(np_type, value) + _refuse_misfit_fields(dtype, np_type, value) if np_type.names is None and np_type.kind in "SU" and isinstance(value, (str, bytes)): if _keeps_value(value, np_type): return dtype @@ -665,6 +668,45 @@ def _holds_numpy_values(value): return any(isinstance(x, (np.number, np.bool_)) for x in items) +def _refuse_misfit_fields(dtype, np_type, value): + """Raise ValueError if a part of literal ``value`` does not broadcast to its array field. + + ``value`` is laid out like an element of record ``np_type`` (UDT ``dtype``, + or a record field in it), as :func:`_weak_literal_udt` reads it, a dict + already a tuple; any other value has no parts to check. An array UDT's + element takes the shape it broadcasts to with a literal, as numpy does, so + ``INT8[3]`` and ``np.array([[1, 300, 3]])`` give ``INT64[1, 3]``. A + record's field keeps its shape, so a part for it must broadcast to that + shape. Converted into the field, as a store converts it, + ``np.array([[1, 300, 3]])`` would lose its leading axis and wrap 300 to 44. + """ + if np_type.names is None or not isinstance(value, (tuple, list)): + return + if len(value) != len(np_type.names): + return # not laid out like an element; typing it reports that + for name, part in zip(np_type.names, value, strict=True): + field = np_type.fields[name][0] + if field.names is not None: + _refuse_misfit_fields(dtype, field, part) + continue + if field.subdtype is None or field.subdtype[0].names is not None: + continue + shape = field.subdtype[1] + try: + given = np.shape(part) + except ValueError: # ragged; converting it reports that + continue + try: + fits = np.broadcast_shapes(given, shape) == shape + except ValueError: + fits = False + if not fits: + raise ValueError( + f"{part!r} has shape {given}, which does not fit field {name!r} of {dtype}, " + f"of shape {shape}: a literal for an array field must broadcast to its shape." + ) + + def _int_beyond_leaf(np_type, value): """Return a Python int in literal ``value`` beyond its integer leaf's range, or ``None``. diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 91c7b3430..80279fc5f 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4883,6 +4883,72 @@ def test_udt_sequence_and_array_literals_type_like_numbers(): assert (s * ones).new().dtype == int8s +def test_udt_record_array_field_refuses_literal_of_another_shape(): + """A literal for a record's array field must broadcast to the field's shape. + + An array UDT's element takes the shape a literal broadcasts it to, as numpy + does: ``INT8[3]`` plus ``np.array([[1, 300, 3]])`` is ``INT64[1, 3]``. A + record's field keeps its shape, so the record refuses such a literal and + names both shapes. Before, its typing gave up and the literal was converted + into the field as numpy assigns it, dropping the leading axis and wrapping + 300 to 44, through ``+`` and through a user's op alike; a nested list + raised numpy's ValueError. A store still converts as numpy assigns. + """ + arr = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_ShapeArrI8") + rec = dtypes.register_anonymous( + np.dtype([("shp_a", np.int8, (3,)), ("shp_b", np.int64)]), "_ShapeRec" + ) + nest = dtypes.register_anonymous( + np.dtype([("shp_n", np.int64), ("shp_in", rec.np_type)]), "_ShapeNest" + ) + v = Vector(arr, size=1) + v[0] = [1, 2, 3] + r = Vector(rec, size=1) + r[0] = ([1, 2, 3], 0) + n = Vector(nest, size=1) + n[0] = (0, ([1, 2, 3], 0)) + exprs = [ + lambda L: r == (L, 0), + lambda L: r.apply(binary.second, right=(L, 0)), + lambda L: r.apply(binary.first, left={"shp_a": L, "shp_b": 0}), + lambda L: r[0].new() == (L, 0), + lambda L: n == (0, (L, 0)), + ] + if supports_udfs: + udf = BinaryOp.register_anonymous(lambda x, y: y, "_shape_right", is_udt=True) + exprs += [ + lambda L: binary.plus(r, (L, 0)), + lambda L: r.apply(udf, right=(L, 0)), + lambda L: r.ewise_union(r, udf, (L, 0), (L, 0)), + ] + for literal in [ + np.array([[1, 300, 3]]), + [[1, 300, 3]], + np.array([[1, 2, 3]], np.int8), + [1, 2], + np.zeros((2, 3), np.int8), + ]: + reason = f"has shape {np.shape(literal)}, which does not fit field 'shp_a'" + for expr in exprs: + with pytest.raises(ValueError, match="of shape \\(3,\\): a literal for an array") as e: + expr(literal) + assert reason in str(e.value) + # A store converts as numpy assigns, on both sides + v[0] = np.array([[1, 300, 3]]) + r[0] = (np.array([[1, 300, 3]]), 0) + assert v[0].new().value.tolist() == [1, 44, 3] + assert r[0].new().value["shp_a"].tolist() == [1, 44, 3] + if supports_udfs: + # The array UDT keeps numpy's arithmetic shape + result = (v + np.array([[1, 300, 3]])).new() + assert result.dtype.np_type == np.dtype((np.array([1]).dtype, (1, 3))) + assert result[0].new().value.tolist() == [[2, 344, 6]] + # and a literal that broadcasts to the field is typed as before + value = binary.plus(r, (np.array([5]), 1)).new()[0].new().value + assert value["shp_a"].tolist() == [6, 49, 8] + assert value["shp_b"] == 1 + + @pytest.mark.skipif("not supports_udfs") def test_udt_ewise_union_defaults_are_typed_not_converted(): """An ``ewise_union`` default beside a UDT is typed as a literal and must fit the operand. @@ -5603,7 +5669,8 @@ def test_udt_literal_rules_edge_cases(): v.apply(udf, right=Fraction(1, 2)) # A literal not laid out like an element (ragged, too long or too short for # a field, not broadcastable), or not a number, has no type of its own, so - # it is converted into the UDT, and numpy says why it cannot be. + # it is converted into the UDT, and numpy says why it cannot be. A record + # says so itself for a part that does not broadcast to its array field. rec = dtypes.register_anonymous( np.dtype([("edge_v", np.float64, (3,)), ("edge_n", np.int8)]), "_EdgeRec" ) @@ -5615,7 +5682,7 @@ def test_udt_literal_rules_edge_cases(): (lambda: v + "abc", "invalid literal"), (lambda: v.apply(udf, right="abc"), "invalid literal"), (lambda: binary.plus(r, (1, 2, 3)), "tuple of length 3"), - (lambda: binary.plus(r, ((1, 2), 3)), "could not broadcast"), + (lambda: binary.plus(r, ((1, 2), 3)), "does not fit field 'edge_v'"), ]: with pytest.raises(ValueError, match=reason): expr() From 530932aa0465d53ba6ebe15444ffee7befbe51cd Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 29 Oct 2026 04:01:12 -0500 Subject: [PATCH 43/54] Cover more literals a UDT refuses and the Numba error of a lifted op Literal typing had arms no test reached: a comparison with a record literal of the wrong length (alone and nested beside an int beyond its leaf), a part holding numpy values that numpy cannot make an array of or that is not numbers, and text a bytes field cannot encode. An array of records now pins the UdfParseError a lifted op gives when Numba cannot compile it, and the NoJITWarning is checked to stay silent for a pair it warned about. Two arms of the monoid identity are excluded from coverage: the binary op compiles first, and it refuses what would reach them. --- graphblas/core/operator/monoid.py | 8 ++++++-- graphblas/tests/test_dtype.py | 5 +++++ graphblas/tests/test_op.py | 19 +++++++++++++++++++ 3 files changed, 30 insertions(+), 2 deletions(-) diff --git a/graphblas/core/operator/monoid.py b/graphblas/core/operator/monoid.py index 0b02d4226..72765e2c2 100644 --- a/graphblas/core/operator/monoid.py +++ b/graphblas/core/operator/monoid.py @@ -69,7 +69,8 @@ def _udt_identity_np(monoid_name, np_type, top_dtype): if np_type.subdtype is not None: base_dtype, shape = np_type.subdtype val = _scalar_identity(monoid_name, base_dtype) - if val is None: + # The binary op compiled first, and min and max compile for arrays of real numbers only + if val is None: # pragma: no cover (safety) raise KeyError( f"monoid.{monoid_name} does not work with {top_dtype}: " f"base dtype {base_dtype} is not numeric" @@ -90,7 +91,10 @@ def _udt_identity_np(monoid_name, np_type, top_dtype): ) vals.append(val) return tuple(vals) - raise KeyError(f"monoid.{monoid_name} does not work with {top_dtype}") + # The binary op compiled first, and lifted ops refuse a UDT neither a record nor an array + raise KeyError( # pragma: no cover (safety) + f"monoid.{monoid_name} does not work with {top_dtype}" + ) ffi_new = ffi.new diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index 140d49a4c..5fb3aad63 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -2,6 +2,7 @@ import pickle import string import sys +import warnings import numpy as np import pytest @@ -469,6 +470,10 @@ def test_udt_with_macro_field_name_falls_back_to_cfunc(): with pytest.warns(NoJITWarning, match="without JIT"): plus_udt = gb.binary.plus[udt] assert plus_udt.jit_c_source is None + # Once per (op, dtype): the same pair again is silent + with warnings.catch_warnings(): + warnings.simplefilter("error") + jit_config._maybe_warn_no_jit(op_name="plus", dtype_name=udt.name) @pytest.mark.skipif("not supports_udfs") diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 80279fc5f..5a86a46ce 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -5671,11 +5671,18 @@ def test_udt_literal_rules_edge_cases(): # a field, not broadcastable), or not a number, has no type of its own, so # it is converted into the UDT, and numpy says why it cannot be. A record # says so itself for a part that does not broadcast to its array field. + # So too for a comparison, which first looks for an int beyond its leaf, a + # part holding numpy values, and text a bytes field cannot hold. rec = dtypes.register_anonymous( np.dtype([("edge_v", np.float64, (3,)), ("edge_n", np.int8)]), "_EdgeRec" ) r = Vector(rec, size=1) r[0] = ([1, 2, 3], 4) + nested = dtypes.register_anonymous( + np.dtype([("edge_in", [("edge_p", np.int8), ("edge_q", np.int8)]), ("edge_m", np.int8)]), + "_EdgeNested", + ) + text = dtypes.register_anonymous(np.dtype([("edge_s", "S3"), ("edge_k", np.int8)]), "_EdgeText") for expr, reason in [ (lambda: binary.plus(v, [1, [2, 3]]), "with a sequence"), (lambda: v + np.array([1, 2]), "could not broadcast"), @@ -5683,6 +5690,12 @@ def test_udt_literal_rules_edge_cases(): (lambda: v.apply(udf, right="abc"), "invalid literal"), (lambda: binary.plus(r, (1, 2, 3)), "tuple of length 3"), (lambda: binary.plus(r, ((1, 2), 3)), "does not fit field 'edge_v'"), + (lambda: r == (1, 2, 3), "tuple of length 3"), + (lambda: Vector(nested, size=1) == ((1, 2, 3), 300), "tuple of length 3"), + (lambda: binary.plus(r, ([np.float64(1), [2, 3]], 4)), "with a sequence"), + (lambda: binary.plus(r, ([np.zeros((3, 2)), np.zeros((3, 1))], 4)), "with a sequence"), + (lambda: binary.plus(r, (np.array(["a", "b", "c"]), 4)), "could not convert string"), + (lambda: Vector(text, size=1) == ("\xe9", 1), "can't encode"), ]: with pytest.raises(ValueError, match=reason): expr() @@ -6121,6 +6134,12 @@ def test_udt_numba_error_summary_is_never_empty(): "Something else went wrong" ) assert summarize(ValueError("")) == "Numba could not compile the function for these input types" + # From a lifted op: Numba cannot add two records, so not an array of records either + arr_rec = dtypes.register_anonymous( + np.dtype(([("sum_p", np.int32), ("sum_q", np.float64)], (2,))), "_SumArrRec" + ) + with pytest.raises(UdfParseError, match=r"plus does not work with \(_SumArrRec, _SumArrRec\)"): + binary.plus[arr_rec] @pytest.mark.skipif("not supports_udfs") From 83643e148973b82f3c6a8adbe8beaaf70b5f8506 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 2 Nov 2026 19:25:03 -0800 Subject: [PATCH 44/54] Give each type signature of a JIT IndexBinaryOp its own C name indexbinary.ss.register_new compiled every type signature of one op under the op's own name. SuiteSparse finds a compiled kernel in memory by its op's name and the types in the kernel's encoding, and an ewise or mxm kernel's encoding has x's, y's and z's types but not theta's. So two signatures that differ only in theta's type ran whichever kernel was compiled first. Now each signature's C function is compiled as name_N_x_M_y_theta, with N and M the lengths of the x and y parts (built-in types by name, UDTs by their C name): the C definition is the user's with the function renamed. A name SuiteSparse would cut at 127 characters is the op name's first 100 and a digest of the whole. The op's jit_c_name and jit_c_source show the C name and the source SuiteSparse compiles; jit_c_definition is still the definition given. JIT IndexUnaryOps and SelectOps are not affected: the thunk's type is in their kernels' encodings. Behavior change (before -> after): with add_theta registered at FP64 for x, y and theta, and then at an FP32 theta with (*z) = (*x) + (*y) + 2 * (double)(*theta), the FP32 signature on FP64 data with theta 10 gave [2, 2, 4, 0] (the FP64 kernel, reading the float as a double) -> [22, 22, 24, 20]. op.jit_c_name: None -> add_theta_4_FP64_4_FP64_FP32. --- graphblas/core/ss/indexbinary.py | 34 +++++++++++++++++++++++++++++--- graphblas/tests/test_ssjit.py | 27 +++++++++++++++++++++---- 2 files changed, 54 insertions(+), 7 deletions(-) diff --git a/graphblas/core/ss/indexbinary.py b/graphblas/core/ss/indexbinary.py index e34451b0a..aee38af3a 100644 --- a/graphblas/core/ss/indexbinary.py +++ b/graphblas/core/ss/indexbinary.py @@ -1,3 +1,6 @@ +import hashlib +import re + from ... import backend from ...dtypes import lookup_dtype from ...exceptions import check_status_carg @@ -13,6 +16,25 @@ ffi_new = ffi.new +def _jit_c_name(funcname, x_type, y_type, theta_type): + """The C name of ``funcname`` registered for these types, unique to the types. + + SuiteSparse finds a compiled kernel in memory by its op's name and the types in the + kernel's encoding, which has x's, y's and z's types but not theta's. So two ops of one + name that differ only in theta's type would run each other's kernel. The name carries + all three types, as ``{funcname}_{len(x)}_{x}_{len(y)}_{y}_{theta}`` (built-in types by + name, UDTs by their C name); the lengths keep a type name with an underscore from + reading as two. SuiteSparse keeps a name of at most 127 characters, so a longer one is + ``funcname``'s start and a digest of the long name. + """ + x, y, theta = (re.sub(r"\W", "_", t.jit_c_name or t.name) for t in (x_type, y_type, theta_type)) + c_name = f"{funcname}_{len(x)}_{x}_{len(y)}_{y}_{theta}" + if len(c_name) >= lib.GxB_MAX_NAME_LEN: + digest = hashlib.sha256(c_name.encode()).hexdigest()[:16] + c_name = f"{funcname[:100]}_{digest}" + return c_name + + class TypedJitIndexBinaryOp(_TypedIndexBinaryOpMixin, TypedOpBase): __slots__ = "_jit_c_definition", "_theta_type" opclass = "IndexBinaryOp" @@ -48,7 +70,9 @@ def register_new(name, jit_c_definition, x_type, y_type, theta_type, ret_type): see the SuiteSparse:GraphBLAS user guide. Only one type signature may be registered at a time, but repeated calls using - the same name with different input types is allowed. + the same name with different input types is allowed. The C function must have the + operator's name (the last part of ``name``). Each type signature compiles it under its + own C name, which the op's ``jit_c_name`` and ``jit_c_source`` show. Parameters ---------- @@ -114,6 +138,9 @@ def register_new(name, jit_c_definition, x_type, y_type, theta_type, ret_type): rv = IndexBinaryOp(name, is_udt=True) from .jit_config import _compiling_c_source + c_name = _jit_c_name(funcname, x_type, y_type, theta_type) + # The function is renamed everywhere it is named, so a recursive call still finds it + c_definition = re.sub(rf"\b{re.escape(funcname)}\b", c_name, jit_c_definition) gb_obj = ffi_new("GxB_IndexBinaryOp*") with _compiling_c_source(f"gb.indexbinary.{name}"): check_status_carg( @@ -124,8 +151,8 @@ def register_new(name, jit_c_definition, x_type, y_type, theta_type, ret_type): x_type._carg, y_type._carg, theta_type._carg, - ffi_new("char[]", funcname.encode()), - ffi_new("char[]", jit_c_definition.encode()), + ffi_new("char[]", c_name.encode()), + ffi_new("char[]", c_definition.encode()), ), "IndexBinaryOp", gb_obj[0], @@ -133,6 +160,7 @@ def register_new(name, jit_c_definition, x_type, y_type, theta_type, ret_type): op = TypedJitIndexBinaryOp( rv, funcname, x_type, ret_type, gb_obj[0], jit_c_definition, y_type, theta_type ) + op._jit_c_info = (c_name, c_definition) rv._add(op, is_jit=True) setattr(module, funcname, rv) return rv diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index a935fded1..70439954c 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -803,6 +803,11 @@ def test_jit_indexbinary(v): assert "FP64" in add_theta assert add_theta["FP64"].return_type == dtypes.FP64 assert add_theta["FP64"].jit_c_definition == cdef + # SuiteSparse compiles it under a C name of its own types + assert add_theta["FP64"].jit_c_name == "add_theta_4_FP64_4_FP64_FP64" + assert add_theta["FP64"].jit_c_source == cdef.replace( + "add_theta", "add_theta_4_FP64_4_FP64_FP64" + ) # Bind theta and use as BinaryOp v64 = v.dup("FP64") binop = add_theta["FP64"](10.0) @@ -845,7 +850,7 @@ def test_jit_indexbinary(v): mixed_cdef = ( "void add_theta (double *z, double *x, GrB_Index ix, GrB_Index jx, " "double *y, GrB_Index iy, GrB_Index jy, float *theta) " - "{ (*z) = (*x) + (*y) + (double)(*theta) ; }" + "{ (*z) = (*x) + (*y) + 2 * (double)(*theta) ; }" ) add_theta_mixed = indexbinary.ss.register_new( "add_theta", mixed_cdef, "FP64", "FP64", "FP32", "FP64" @@ -854,9 +859,23 @@ def test_jit_indexbinary(v): assert ("FP64", "FP64", "FP32") in add_theta assert ("FP64", "FP32") not in add_theta # theta of y's type, FP32 assert repr(add_theta["FP64", "FP64", "FP32"]) == "indexbinary.add_theta[FP64, FP64, FP32]" - # An FP32 theta picks it. Not computed here: SuiteSparse's JIT knows a user op's kernel - # by its name and x, y, z types, so in this process it would run the FP64 theta's kernel. - assert v64.ewise_mult(v64, add_theta(np.float32(10))).op._theta_type == dtypes.FP32 + # An FP32 theta picks it, and its kernel runs. SuiteSparse finds a compiled kernel by its + # op's name, and its encoding has no theta type, so with one C name for both theta types + # this ran the FP64 theta's kernel, which read the float as a double: [2, 2, 4, 0]. + assert add_theta["FP64", "FP64", "FP32"].jit_c_name == "add_theta_4_FP64_4_FP64_FP32" + expected_mixed = Vector.from_coo([1, 3, 4, 6], [22.0, 22.0, 24.0, 20.0]) + res = v64.ewise_mult(v64, add_theta(np.float32(10))).new() + assert expected_mixed.isequal(res, check_dtype=True) + assert expected.isequal(v64.ewise_mult(v64, add_theta(10.0)).new(), check_dtype=True) + # A name too long for SuiteSparse with its types is cut, with a digest of the whole + long_name = "add_theta_" + "long" * 25 + long_op = indexbinary.ss.register_new( + long_name, cdef.replace("add_theta", long_name), "FP64", "FP64", "FP64", "FP64" + ) + c_name = long_op["FP64"].jit_c_name + assert len(c_name) == 117 + assert c_name.startswith(long_name[:100]) + assert expected.isequal(v64.ewise_mult(v64, long_op(10.0)).new(), check_dtype=True) @pytest.mark.slow From cd6e0a85a941daad3f938189fe37f8f5e514aef5 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 2 Nov 2026 19:38:29 -0800 Subject: [PATCH 45/54] Name each op registered from C by its full name and definition The ss.register_new functions of unary, binary, indexunary (and select) gave SuiteSparse the op's last name as its C name, and indexbinary's the last name and types. SuiteSparse finds a compiled kernel in memory by its op's name, and checks an op's definition only when it loads a kernel from disk. So, in one process, an op registered again under its name with another definition (after delattr, as when a notebook cell is run again), or an op of the same last name in another namespace (ss.f and ss.nested.f), ran the kernel compiled first. Now every op registered from C source is compiled under a C name made of its last name, the types for an IndexBinaryOp (whose theta type is not in its kernels' encoding), and a 12-hex digest of these with the full dotted name and the definition: f_3a1c90e2d4b7, or add_theta_4_FP64_4_FP64_4_FP32_... for an IndexBinaryOp. A name SuiteSparse would cut at 127 characters is the last name's first 100 characters and a 16-hex digest. The function is renamed in the definition wherever it is named, except as a struct member (after "." or "->"). The same name and definition give the same C name, so a kernel is compiled once for them. jit_c_name and jit_c_source, None for these ops before (except IndexBinaryOp, from the previous commit), show the C name and the renamed source; jit_c_definition is still the definition given. The helper is ss._utils._jit_c_source. Behavior changes (before -> after), in one process: - binary.ss.register_new("f", x + y) and ("nested.f", x - y), both FP64: nested.f computed x + y -> x - y. - unary.ss.register_new("u", x * x), delattr(unary.ss, "u"), then register_new("u", x + 10): u computed x * x -> x + 10. --- graphblas/core/ss/_utils.py | 32 ++++++++++++++ graphblas/core/ss/binary.py | 14 ++++-- graphblas/core/ss/indexbinary.py | 37 +++++----------- graphblas/core/ss/indexunary.py | 15 +++++-- graphblas/core/ss/select.py | 7 ++- graphblas/core/ss/unary.py | 14 ++++-- graphblas/tests/test_ssjit.py | 74 +++++++++++++++++++++++++++++--- 7 files changed, 150 insertions(+), 43 deletions(-) diff --git a/graphblas/core/ss/_utils.py b/graphblas/core/ss/_utils.py index 6ab0c68f7..85017804f 100644 --- a/graphblas/core/ss/_utils.py +++ b/graphblas/core/ss/_utils.py @@ -1,4 +1,6 @@ +import hashlib import itertools +import re from ...exceptions import _error_code_lookup from .. import ffi, lib @@ -45,3 +47,33 @@ def _resolve_serialized_dtype(data_obj, data_nbytes, obj_kind): f"{obj_kind} deserialize cannot resolve dtype {dtype_name!r} and " f"GxB_Serialized_get_String is unavailable in this SuiteSparse build" ) from orig_error + + +def _jit_c_source(name, funcname, jit_c_definition, types=()): + """Return the C name and definition to give SuiteSparse for an op registered from C. + + SuiteSparse finds a compiled kernel in memory by its op's name (and the types in the + kernel's encoding), and checks an op's definition only when it loads a kernel from disk. + So two ops of one C name with different definitions would run each other's kernel: an op + defined again under its name, or two of one last name in two namespaces (``ss.f`` and + ``ss.nested.f``). The C name is ``funcname``, then ``_{len(t)}_{t}`` for each of + ``types`` (built-in types by name, UDTs by their C name), then a digest of these, the + full dotted ``name`` and the definition, as ``f_3a1c...``. SuiteSparse keeps a name of at + most 127 characters, so a longer one is ``funcname``'s start and a longer digest. One + name and definition give one C name on purpose: registered again, they reuse the kernel + compiled for them. The function is renamed in the definition wherever it is named, so a + recursive call follows, but not a struct member of the same name (after ``.`` or ``->``). + """ + parts = [re.sub(r"\W", "_", t.jit_c_name or t.name) for t in types] + readable = funcname + "".join(f"_{len(part)}_{part}" for part in parts) + digest = hashlib.sha256(f"{readable}\0{name}\0{jit_c_definition}".encode()).hexdigest() + c_name = f"{readable}_{digest[:12]}" + if len(c_name) >= lib.GxB_MAX_NAME_LEN: + c_name = f"{funcname[:100]}_{digest[:16]}" + + def rename(match): + if jit_c_definition[: match.start()].rstrip().endswith((".", "->")): + return match.group() + return c_name + + return c_name, re.sub(rf"\b{re.escape(funcname)}\b", rename, jit_c_definition) diff --git a/graphblas/core/ss/binary.py b/graphblas/core/ss/binary.py index 87c213d32..547d06ccf 100644 --- a/graphblas/core/ss/binary.py +++ b/graphblas/core/ss/binary.py @@ -5,6 +5,7 @@ from ..operator.base import TypedOpBase from ..operator.binary import BinaryOp, TypedUserBinaryOp from . import _IS_SSGB7 +from ._utils import _jit_c_source ffi_new = ffi.new @@ -44,7 +45,12 @@ def register_new(name, jit_c_definition, left_type, right_type, ret_type): see the SuiteSparse:GraphBLAS user guide. Only one type signature may be registered at a time, but repeated calls using - the same name with different input types is allowed. + the same name with different input types is allowed. The C function must have the + operator's name (the last part of ``name``). It is compiled under a C name made from + the full name and the definition, which the op's ``jit_c_name`` and ``jit_c_source`` + show, so another definition, or the same last name in another namespace, gets its own + kernel. Registering one name with one definition again (after ``delattr``) reuses its + C name, and the kernel already compiled for it. Parameters ---------- @@ -108,6 +114,7 @@ def register_new(name, jit_c_definition, left_type, right_type, ret_type): rv = BinaryOp(name, is_udt=True) from .jit_config import _compiling_c_source + c_name, c_definition = _jit_c_source(name, funcname, jit_c_definition) gb_obj = ffi_new("GrB_BinaryOp*") with _compiling_c_source(f"gb.binary.{name}"): check_status_carg( @@ -117,8 +124,8 @@ def register_new(name, jit_c_definition, left_type, right_type, ret_type): ret_type._carg, left_type._carg, right_type._carg, - ffi_new("char[]", funcname.encode()), - ffi_new("char[]", jit_c_definition.encode()), + ffi_new("char[]", c_name.encode()), + ffi_new("char[]", c_definition.encode()), ), "BinaryOp", gb_obj[0], @@ -126,6 +133,7 @@ def register_new(name, jit_c_definition, left_type, right_type, ret_type): op = TypedJitBinaryOp( rv, funcname, left_type, ret_type, gb_obj[0], jit_c_definition, dtype2=right_type ) + op._jit_c_info = (c_name, c_definition) rv._add(op, is_jit=True) setattr(module, funcname, rv) return rv diff --git a/graphblas/core/ss/indexbinary.py b/graphblas/core/ss/indexbinary.py index aee38af3a..5aeed9f11 100644 --- a/graphblas/core/ss/indexbinary.py +++ b/graphblas/core/ss/indexbinary.py @@ -1,6 +1,3 @@ -import hashlib -import re - from ... import backend from ...dtypes import lookup_dtype from ...exceptions import check_status_carg @@ -12,29 +9,11 @@ _has_idxbinop, _TypedIndexBinaryOpMixin, ) +from ._utils import _jit_c_source ffi_new = ffi.new -def _jit_c_name(funcname, x_type, y_type, theta_type): - """The C name of ``funcname`` registered for these types, unique to the types. - - SuiteSparse finds a compiled kernel in memory by its op's name and the types in the - kernel's encoding, which has x's, y's and z's types but not theta's. So two ops of one - name that differ only in theta's type would run each other's kernel. The name carries - all three types, as ``{funcname}_{len(x)}_{x}_{len(y)}_{y}_{theta}`` (built-in types by - name, UDTs by their C name); the lengths keep a type name with an underscore from - reading as two. SuiteSparse keeps a name of at most 127 characters, so a longer one is - ``funcname``'s start and a digest of the long name. - """ - x, y, theta = (re.sub(r"\W", "_", t.jit_c_name or t.name) for t in (x_type, y_type, theta_type)) - c_name = f"{funcname}_{len(x)}_{x}_{len(y)}_{y}_{theta}" - if len(c_name) >= lib.GxB_MAX_NAME_LEN: - digest = hashlib.sha256(c_name.encode()).hexdigest()[:16] - c_name = f"{funcname[:100]}_{digest}" - return c_name - - class TypedJitIndexBinaryOp(_TypedIndexBinaryOpMixin, TypedOpBase): __slots__ = "_jit_c_definition", "_theta_type" opclass = "IndexBinaryOp" @@ -71,8 +50,11 @@ def register_new(name, jit_c_definition, x_type, y_type, theta_type, ret_type): Only one type signature may be registered at a time, but repeated calls using the same name with different input types is allowed. The C function must have the - operator's name (the last part of ``name``). Each type signature compiles it under its - own C name, which the op's ``jit_c_name`` and ``jit_c_source`` show. + operator's name (the last part of ``name``). It is compiled under a C name made from + the full name and the definition, which the op's ``jit_c_name`` and ``jit_c_source`` + show, so another definition, or the same last name in another namespace, gets its own + kernel. Registering one name with one definition again (after ``delattr``) reuses its + C name, and the kernel already compiled for it. Parameters ---------- @@ -138,9 +120,10 @@ def register_new(name, jit_c_definition, x_type, y_type, theta_type, ret_type): rv = IndexBinaryOp(name, is_udt=True) from .jit_config import _compiling_c_source - c_name = _jit_c_name(funcname, x_type, y_type, theta_type) - # The function is renamed everywhere it is named, so a recursive call still finds it - c_definition = re.sub(rf"\b{re.escape(funcname)}\b", c_name, jit_c_definition) + # Theta's type is not in the kernel's encoding, so the C name carries the types + c_name, c_definition = _jit_c_source( + name, funcname, jit_c_definition, (x_type, y_type, theta_type) + ) gb_obj = ffi_new("GxB_IndexBinaryOp*") with _compiling_c_source(f"gb.indexbinary.{name}"): check_status_carg( diff --git a/graphblas/core/ss/indexunary.py b/graphblas/core/ss/indexunary.py index fa394eb4f..a3a462f59 100644 --- a/graphblas/core/ss/indexunary.py +++ b/graphblas/core/ss/indexunary.py @@ -5,6 +5,7 @@ from ..operator.base import TypedOpBase from ..operator.indexunary import IndexUnaryOp, TypedUserIndexUnaryOp from . import _IS_SSGB7 +from ._utils import _jit_c_source ffi_new = ffi.new @@ -39,7 +40,12 @@ def register_new(name, jit_c_definition, input_type, thunk_type, ret_type): see the SuiteSparse:GraphBLAS user guide. Only one type signature may be registered at a time, but repeated calls using - the same name with different input types is allowed. + the same name with different input types is allowed. The C function must have the + operator's name (the last part of ``name``). It is compiled under a C name made from + the full name and the definition, which the op's ``jit_c_name`` and ``jit_c_source`` + show, so another definition, or the same last name in another namespace, gets its own + kernel. Registering one name with one definition again (after ``delattr``) reuses its + C name, and the kernel already compiled for it. This will also create a SelectOp operator under ``gb.select.ss`` if the return type is boolean. @@ -107,6 +113,7 @@ def register_new(name, jit_c_definition, input_type, thunk_type, ret_type): rv = IndexUnaryOp(name, is_udt=True) from .jit_config import _compiling_c_source + c_name, c_definition = _jit_c_source(name, funcname, jit_c_definition) gb_obj = ffi_new("GrB_IndexUnaryOp*") with _compiling_c_source(f"gb.indexunary.{name}"): check_status_carg( @@ -116,8 +123,8 @@ def register_new(name, jit_c_definition, input_type, thunk_type, ret_type): ret_type._carg, input_type._carg, thunk_type._carg, - ffi_new("char[]", funcname.encode()), - ffi_new("char[]", jit_c_definition.encode()), + ffi_new("char[]", c_name.encode()), + ffi_new("char[]", c_definition.encode()), ), "IndexUnaryOp", gb_obj[0], @@ -125,6 +132,7 @@ def register_new(name, jit_c_definition, input_type, thunk_type, ret_type): op = TypedJitIndexUnaryOp( rv, funcname, input_type, ret_type, gb_obj[0], jit_c_definition, dtype2=thunk_type ) + op._jit_c_info = (c_name, c_definition) rv._add(op, is_jit=True) if ret_type == BOOL: from ..operator.select import SelectOp @@ -150,6 +158,7 @@ def register_new(name, jit_c_definition, input_type, thunk_type, ret_type): op2 = TypedJitSelectOp( selectop, funcname, input_type, ret_type, gb_obj[0], jit_c_definition, dtype2=thunk_type ) + op2._jit_c_info = (c_name, c_definition) selectop._add(op2, is_jit=True) setattr(select_module, funcname, selectop) setattr(module, funcname, rv) diff --git a/graphblas/core/ss/select.py b/graphblas/core/ss/select.py index 3ba135eee..48d16b740 100644 --- a/graphblas/core/ss/select.py +++ b/graphblas/core/ss/select.py @@ -38,7 +38,12 @@ def register_new(name, jit_c_definition, input_type, thunk_type): see the SuiteSparse:GraphBLAS user guide. Only one type signature may be registered at a time, but repeated calls using - the same name with different input types is allowed. + the same name with different input types is allowed. The C function must have the + operator's name (the last part of ``name``). It is compiled under a C name made from + the full name and the definition, which the op's ``jit_c_name`` and ``jit_c_source`` + show, so another definition, or the same last name in another namespace, gets its own + kernel. Registering one name with one definition again (after ``delattr``) reuses its + C name, and the kernel already compiled for it. This will also create an IndexUnary operator under ``gb.indexunary.ss`` diff --git a/graphblas/core/ss/unary.py b/graphblas/core/ss/unary.py index e0be45db9..cee7a16a4 100644 --- a/graphblas/core/ss/unary.py +++ b/graphblas/core/ss/unary.py @@ -5,6 +5,7 @@ from ..operator.base import TypedOpBase from ..operator.unary import TypedUserUnaryOp, UnaryOp from . import _IS_SSGB7 +from ._utils import _jit_c_source ffi_new = ffi.new @@ -38,7 +39,12 @@ def register_new(name, jit_c_definition, input_type, ret_type): see the SuiteSparse:GraphBLAS user guide. Only one type signature may be registered at a time, but repeated calls using - the same name with different input types is allowed. + the same name with different input types is allowed. The C function must have the + operator's name (the last part of ``name``). It is compiled under a C name made from + the full name and the definition, which the op's ``jit_c_name`` and ``jit_c_source`` + show, so another definition, or the same last name in another namespace, gets its own + kernel. Registering one name with one definition again (after ``delattr``) reuses its + C name, and the kernel already compiled for it. Parameters ---------- @@ -92,6 +98,7 @@ def register_new(name, jit_c_definition, input_type, ret_type): rv = UnaryOp(name, is_udt=True) from .jit_config import _compiling_c_source + c_name, c_definition = _jit_c_source(name, funcname, jit_c_definition) gb_obj = ffi_new("GrB_UnaryOp*") with _compiling_c_source(f"gb.unary.{name}"): check_status_carg( @@ -100,13 +107,14 @@ def register_new(name, jit_c_definition, input_type, ret_type): NULL, ret_type._carg, input_type._carg, - ffi_new("char[]", funcname.encode()), - ffi_new("char[]", jit_c_definition.encode()), + ffi_new("char[]", c_name.encode()), + ffi_new("char[]", c_definition.encode()), ), "UnaryOp", gb_obj[0], ) op = TypedJitUnaryOp(rv, funcname, input_type, ret_type, gb_obj[0], jit_c_definition) + op._jit_c_info = (c_name, c_definition) rv._add(op, is_jit=True) setattr(module, funcname, rv) return rv diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 70439954c..1a04a315c 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -803,11 +803,11 @@ def test_jit_indexbinary(v): assert "FP64" in add_theta assert add_theta["FP64"].return_type == dtypes.FP64 assert add_theta["FP64"].jit_c_definition == cdef - # SuiteSparse compiles it under a C name of its own types - assert add_theta["FP64"].jit_c_name == "add_theta_4_FP64_4_FP64_FP64" - assert add_theta["FP64"].jit_c_source == cdef.replace( - "add_theta", "add_theta_4_FP64_4_FP64_FP64" - ) + # SuiteSparse compiles it under a C name of its types, name and definition + c_name = add_theta["FP64"].jit_c_name + assert c_name.startswith("add_theta_4_FP64_4_FP64_4_FP64_") + assert len(c_name) == len("add_theta_4_FP64_4_FP64_4_FP64_") + 12 + assert add_theta["FP64"].jit_c_source == cdef.replace("add_theta", c_name) # Bind theta and use as BinaryOp v64 = v.dup("FP64") binop = add_theta["FP64"](10.0) @@ -862,7 +862,9 @@ def test_jit_indexbinary(v): # An FP32 theta picks it, and its kernel runs. SuiteSparse finds a compiled kernel by its # op's name, and its encoding has no theta type, so with one C name for both theta types # this ran the FP64 theta's kernel, which read the float as a double: [2, 2, 4, 0]. - assert add_theta["FP64", "FP64", "FP32"].jit_c_name == "add_theta_4_FP64_4_FP64_FP32" + assert add_theta["FP64", "FP64", "FP32"].jit_c_name.startswith( + "add_theta_4_FP64_4_FP64_4_FP32_" + ) expected_mixed = Vector.from_coo([1, 3, 4, 6], [22.0, 22.0, 24.0, 20.0]) res = v64.ewise_mult(v64, add_theta(np.float32(10))).new() assert expected_mixed.isequal(res, check_dtype=True) @@ -878,6 +880,66 @@ def test_jit_indexbinary(v): assert expected.isequal(v64.ewise_mult(v64, long_op(10.0)).new(), check_dtype=True) +@pytest.mark.slow +def test_jit_c_name_per_name_and_definition(v): + """An op registered from C is compiled under a C name of its full name and definition. + + SuiteSparse finds a compiled kernel in memory by its op's name, and checks a definition + only when it loads a kernel from disk. When the C name was the op's last name, an op + defined again after ``delattr``, or one of the same last name in another namespace, ran + the kernel compiled first. + """ + if _IS_SSGB7: + return + _require_jit_on() + v64 = v.dup("FP64") + # One last name in two namespaces + plus = "void jcn_f (double *z, double *x, double *y) { (*z) = (*x) + (*y) ; }" + f = binary.ss.register_new("jcn_f", plus, "FP64", "FP64", "FP64") + g = binary.ss.register_new("jcn_ns.jcn_f", plus.replace("+", "-"), "FP64", "FP64", "FP64") + try: + expected = Vector.from_coo([1, 3, 4, 6], [2.0, 2.0, 4.0, 0.0]) + assert v64.ewise_mult(v64, f).new().isequal(expected, check_dtype=True) + expected = Vector.from_coo([1, 3, 4, 6], [0.0, 0.0, 0.0, 0.0]) + assert v64.ewise_mult(v64, g).new().isequal(expected, check_dtype=True) + assert f["FP64"].jit_c_name != g["FP64"].jit_c_name + finally: + delattr(binary.ss, "jcn_f") + delattr(binary.ss, "jcn_ns") + # One name defined again, after delattr, with another definition and then the first + square = "void jcn_u (double *z, double *x) { (*z) = (*x) * (*x) ; }" + add10 = square.replace("* (*x)", "+ 10") + names = [] + for cdef, values in [(square, [1, 1, 4, 0]), (add10, [11, 11, 12, 10]), (square, [1, 1, 4, 0])]: + op = unary.ss.register_new("jcn_u", cdef, "FP64", "FP64") + try: + expected = Vector.from_coo([1, 3, 4, 6], values, dtype="FP64") + assert v64.apply(op).new().isequal(expected, check_dtype=True) + names.append(op["FP64"].jit_c_name) + finally: + delattr(unary.ss, "jcn_u") + assert names[0] != names[1] + assert names[0] == names[2] # one definition, one kernel + # A struct member of the op's name is not renamed + rec = dtypes.ss.register_new( + "JcnRec", + "typedef struct { int64_t jcn_m ; double b ; } JcnRec ;", + np_type=np.dtype([("jcn_m", " Date: Mon, 2 Nov 2026 19:55:21 -0800 Subject: [PATCH 46/54] Refuse an empty Scalar theta for an IndexBinaryOp SuiteSparse binds theta without checking that it holds a value: it reads the Scalar's value array. Binding an empty Scalar crashed the interpreter, and op(s) with no type, which binds a copy of s, ran with theta 0. Both now raise EmptyObject, as apply does for an empty IndexUnaryOp thunk, in every bind: typed or not, is_udt or not, a JIT op's too. Behavior change (before -> after), with s = Scalar(int): - op[int](s) and op(s, dtype=int): segmentation fault -> EmptyObject. - op(s): an op that ran with theta 0 -> EmptyObject, at the bind. With s a C scalar (is_cscalar=True): segmentation fault -> EmptyObject. --- graphblas/core/operator/indexbinary.py | 12 +++++++++++- graphblas/tests/test_indexbinary.py | 14 +++++++++++++- graphblas/tests/test_ssjit.py | 5 +++++ 3 files changed, 29 insertions(+), 2 deletions(-) diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index 991d8f0f4..421c94cdf 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -6,7 +6,7 @@ from ... import _STANDARD_OPERATOR_NAMES, indexbinary from ...dtypes import BOOL, INT8, UINT64, lookup_dtype -from ...exceptions import UdfParseError, check_status_carg +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 @@ -72,6 +72,14 @@ def _theta_to_scalar(theta, dtype): raise +def _refuse_empty_theta(theta): + # SuiteSparse binds theta without checking that it holds a value: binding a new + # empty Scalar crashes, and a copy of one binds whatever its value array holds. + # ``apply`` with an IndexUnaryOp raises EmptyObject for an empty thunk. + if theta._is_empty: + raise EmptyObject("theta is an empty Scalar; bind a Scalar that holds a value") + + class _BoundIndexBinaryOp(TypedOpBase): """A BinaryOp produced by binding a theta value to a typed IndexBinaryOp. @@ -235,6 +243,7 @@ def __call__(self, theta=None): val, dt = theta.value, theta.dtype theta = Scalar.from_value(val, dt, is_cscalar=False, name="") # pragma: is_grbscalar # fmt: on + _refuse_empty_theta(theta) new_binop = ffi_new("GrB_BinaryOp*") check_status_carg( lib.GxB_BinaryOp_new_IndexOp(new_binop, self.gb_obj, theta._carg), @@ -677,6 +686,7 @@ def __call__(self, theta=None, dtype=None): theta = False if dtype is None: if isinstance(theta, Scalar): + _refuse_empty_theta(theta) # Bound when it runs, so keep the value theta has now. theta = theta.dup(name="") else: diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py index 83acac65a..304f7230a 100644 --- a/graphblas/tests/test_indexbinary.py +++ b/graphblas/tests/test_indexbinary.py @@ -7,7 +7,7 @@ from graphblas import Matrix, Scalar, Vector, dtypes, indexbinary from graphblas.core import _supports_udfs as supports_udfs from graphblas.core.operator.indexbinary import _has_idxbinop -from graphblas.exceptions import UdfParseError +from graphblas.exceptions import EmptyObject, UdfParseError pytestmark = [ pytest.mark.skipif(not supports_udfs, reason="requires numba"), @@ -388,6 +388,18 @@ def _combine(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) assert good[arr_udt].return_type is arr_udt +def test_empty_scalar_theta(): + # SuiteSparse would read the empty Scalar's value: a crash, or a value it never held + op = indexbinary.register_anonymous(_ibo_add_theta) + opu = indexbinary.register_anonymous(_ibo_add_theta, is_udt=True) + cleared = Scalar.from_value(5) + cleared.clear() + for theta in [Scalar(int), Scalar(int, is_cscalar=True), cleared]: + for bind in [op, opu, op[int], opu[int], lambda t: op(t, dtype=int)]: + with pytest.raises(EmptyObject, match="theta is an empty Scalar"): + bind(theta) + + def test_bind_raw_udt_theta_without_dtype(): """A numpy array theta is typed as its array layout, as other array literals are. diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 1a04a315c..1b8c6d673 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -28,6 +28,7 @@ from graphblas.core.operator.udt_codegen import _has_jit_set from graphblas.core.ss import _IS_SSGB7 from graphblas.core.ss.jit_config import _enable_jit_for_udt +from graphblas.exceptions import EmptyObject from .conftest import autocompute, burble @@ -824,6 +825,10 @@ def test_jit_indexbinary(v): v64.ewise_mult(v64, add_theta(10)) with pytest.raises(KeyError, match="add_theta does not work with .INT64, INT64, FP64. types"): v.ewise_mult(v, add_theta(10.0)) + # An empty theta is refused, as for any IndexBinaryOp (SuiteSparse would read its value) + for bind in [add_theta, add_theta["FP64"]]: + with pytest.raises(EmptyObject, match="theta is an empty Scalar"): + bind(gb.Scalar("FP64")) # Test duplicate registration fails assert "FP32" not in add_theta with burble(): From 8cb81c3c3d91116a51a859f71c6d95146ab5d685 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 07:32:32 -0800 Subject: [PATCH 47/54] Leave a dtypes.ss.register_new name free when the registration fails register_new recorded the name as a JIT C type name as soon as SuiteSparse made the type, before the numpy and Numba types were made. When one of those raised (an np_type with an array of bool, which Numba cannot type), the name stayed reserved, and a correct second call raised "name for dtype is unavailable". Main allowed the second call; the check against C type names came with "Give each UDT layout its own JIT C name". The name is now recorded after the type is registered. The type made by the failed call is never used, so no JIT kernel knows its name. Behavior change (before -> after), with d the typedef of pt_t: - register_new("pt_t", d, np_type=np.dtype((np.bool_, (2,)))) raises ValueError (unchanged), then register_new("pt_t", d, np_type=np.dtype([("x", "f8"), ("y", "f8")])): ValueError "'pt_t' name for dtype is unavailable" -> registers dtypes.ss.pt_t. --- graphblas/core/ss/dtypes.py | 5 ++++- graphblas/tests/test_ssjit.py | 7 +++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/graphblas/core/ss/dtypes.py b/graphblas/core/ss/dtypes.py index e82dc1224..0ccdd8264 100644 --- a/graphblas/core/ss/dtypes.py +++ b/graphblas/core/ss/dtypes.py @@ -61,7 +61,6 @@ def register_new(name, jit_c_definition, *, np_type=None): ffi_new("char[]", jit_c_definition.encode()), ) check_status_carg(status, "Type", gb_obj[0]) - _jit_c_type_defs[name] = jit_c_definition # Let SuiteSparse:GraphBLAS determine the size (we gave 0 as size above) size_ptr = ffi_new("size_t*") @@ -101,4 +100,8 @@ def register_new(name, jit_c_definition, *, np_type=None): if save_np_type or np_type not in core.dtypes._registry: core.dtypes._registry[np_type] = rv setattr(dtypes.ss, name, rv) + # Only now, so that a registration that fails above (Numba cannot type an + # array of bool, say) leaves the name free to register again. The type it + # made is never used, so no JIT kernel knows its name. + _jit_c_type_defs[name] = jit_c_definition return rv diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 1b8c6d673..9e083a2a1 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -603,6 +603,13 @@ def test_jit_udt_np_type_numba_cannot_type(): ) assert isinstance(exc_info.value.__cause__, AttributeError) assert not hasattr(dtypes.ss, "_jit_bool_pair") + # The failure leaves the name free, so it registers with an np_type Numba can type. + dtype = dtypes.ss.register_new( + "_jit_bool_pair", + "typedef struct { uint8_t x [2] ; } _jit_bool_pair ;", + np_type=np.dtype([("x", np.uint8, (2,))]), + ) + assert dtypes.ss._jit_bool_pair is dtype @pytest.mark.skipif("_IS_SSGB7") From b3319636ad995b55dce71b71251fcfcac81d23c5 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 07:37:16 -0800 Subject: [PATCH 48/54] Type a Python date or duration in a record literal as numpy does A record literal's date or duration field takes the field's type only when its value converts into the field and back unchanged; otherwise it keeps its own type, so == compares it whole. A numpy date did this already, but a Python datetime.date, datetime.datetime or datetime.timedelta, or a date given as text, always took the field's type. So it was cut to the field's unit before == compared it: noon compared equal to its day in a field of days, and 5 s and 1 us equal to 5 s in a field of seconds. Main gave the same wrong answers (it gave them for numpy dates too). Now such a value is first made the numpy value it stands for (a datetime in microseconds, a timedelta in microseconds, a date as text in the unit it is written to), and is typed as that numpy value is. Behavior change (before -> after), with v a Vector of a record of an M8[D] and an m8[s] field holding (2020-01-01, 5 s), and s = v[0].new(): - v == (datetime(2020, 1, 1, 12), np.timedelta64(5, "s")), and with "2020-01-01T12:00" for the date: True -> False; the same for s ==. - v == (np.datetime64("2020-01-01"), timedelta(seconds=5, microseconds=1)): True -> False. - v.apply(user_op, right=(datetime(2020, 1, 1, 12), np.timedelta64(5, "s"))): stored the day, silently -> ValueError "does not fit", as for np.datetime64("2020-01-01T12:00"). - v + (datetime(2020, 1, 1, 12), ...): UdfParseError from Numba -> KeyError naming datetime64[D] and datetime64[us]. - A datetime at midnight, a date, and text such as "2020-01-01" still compare equal and fit the field. --- graphblas/core/literals.py | 32 ++++++++++++++++++++++++++---- graphblas/tests/test_op.py | 40 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 68 insertions(+), 4 deletions(-) diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index 18dca86db..97e301416 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -589,12 +589,15 @@ def leaf_dtype(leaf, part, *, in_record=False): # beside it are. It takes the field's own type when it converts into # the field and back unchanged, so ops pair the field with itself; # otherwise its own type, which does not fit: text as the field's - # kind at its own length, a numpy date or duration in its unit. - # Other dates (datetime, str) take the field's type. + # kind at its own length, a date or duration in numpy's unit for it + # (a datetime is microseconds), so == compares it as numpy does. + # Text in a duration field, or text numpy cannot read as a date, + # takes the field's type, and converting it reports what is wrong. if isinstance(part, (str, bytes)) and base.kind in "SU": own = np.dtype(f"{base.kind}{max(len(part), 1)}") - elif isinstance(part, (np.datetime64, np.timedelta64)): - own = part.dtype + elif (date := _as_numpy_date(part, base.kind)) is not None: + part = date + own = date.dtype if own.kind != base.kind: return own elif isinstance(part, (str, bytes, datetime.date, datetime.timedelta)): @@ -655,6 +658,27 @@ def record_leaves(np_type, part): return dtype if literal_type == np_type else lookup_dtype(literal_type) +def _as_numpy_date(value, kind): + """Return ``value`` as numpy's ``datetime64`` or ``timedelta64``, or ``None`` if it is neither. + + A numpy date or duration is itself; a Python ``datetime.date`` or + ``datetime.datetime`` is a date in days or microseconds and a + ``datetime.timedelta`` a duration in microseconds, as numpy makes them; + text is a date in the unit it is written to (``"2020-01-01T12:00"`` is in + minutes) when ``kind``, the leaf's, is ``"M"`` and numpy can read it. + """ + if isinstance(value, (np.datetime64, np.timedelta64)): + return value + try: + if isinstance(value, datetime.timedelta): + return np.timedelta64(value) + if isinstance(value, datetime.date) or (kind == "M" and isinstance(value, (str, bytes))): + return np.datetime64(value) + except ValueError: # text that is not a date; converting it reports that + pass + return None + + def _holds_numpy_values(value): """Whether ``value`` is a numpy array, or a sequence holding a numpy number or bool.""" if isinstance(value, np.ndarray): diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 5a86a46ce..9d18b4afc 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4949,6 +4949,46 @@ def test_udt_record_array_field_refuses_literal_of_another_shape(): assert value["shp_b"] == 1 +def test_udt_record_date_literals_compare_as_numpy(): + """A Python date or duration, or a date as text, in a record literal is typed as numpy types it. + + It takes its field's type only when it converts into the field and back + unchanged, as a numpy date does; otherwise it keeps numpy's type for it (a + ``datetime`` is in microseconds), so ``==`` compares it whole: noon is not + the day, and 5 s and 1 us is not 5 s. Before, it always took the field's + type, so it was cut to the field's unit and compared equal. + """ + import datetime + + rec = dtypes.register_anonymous(np.dtype([("dl_d", "M8[D]"), ("dl_t", "m8[s]")]), "_DateLit") + day, secs = np.datetime64("2020-01-01"), np.timedelta64(5, "s") + v = Vector(rec, size=1) + v[0] = (day, secs) + s = v[0].new() + # numpy's dates have no time zone, so neither do these. + midnight = datetime.datetime(2020, 1, 1) # noqa: DTZ001 + noon = datetime.datetime(2020, 1, 1, 12) # noqa: DTZ001 + for literal, equal in [ + ((noon, secs), False), + ((day, datetime.timedelta(seconds=5, microseconds=1)), False), + (("2020-01-01T12:00", secs), False), + ((datetime.date(2020, 1, 1), datetime.timedelta(seconds=5)), True), + ((midnight, secs), True), + (("2020-01-01", secs), True), + ]: + assert (s == literal) is equal + if supports_udfs: + assert (v == literal).new()[0].new().value == equal + assert (v != literal).new()[0].new().value != equal + if supports_udfs: + # A user's op refuses what its field would cut short, as for a numpy date. + second = BinaryOp.register_anonymous(lambda x, y: y, "_date_lit_second", is_udt=True) + with pytest.raises(ValueError, match="does not fit _DateLit"): + v.apply(second, right=(noon, secs)) + result = v.apply(second, right=(midnight, secs)).new() + assert result[0].new().value.tolist() == (day, secs) + + @pytest.mark.skipif("not supports_udfs") def test_udt_ewise_union_defaults_are_typed_not_converted(): """An ``ewise_union`` default beside a UDT is typed as a literal and must fit the operand. From cf64e896efd02b50ae3e91b1fccbe688a095f432 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 07:43:35 -0800 Subject: [PATCH 49/54] Compare a UDT of dates or durations with one in another unit as numpy does A UDT that is a bare datetime64 or timedelta64 (not a record or an array) converted a date or duration literal into its own unit before == compared it, so noon compared equal to its day in a UDT of days. Main raised TypeError here for a Vector, and gave True for a Scalar; "Compare UDTs that are neither records nor arrays in eq and ne" made the Vector give True too. A numpy duration in another unit was a KeyError instead, even when it was the same duration. Now such a literal is typed as in a record's date field: the UDT itself when it converts into the UDT and back unchanged, else numpy's type for it (a datetime is in microseconds, a timedelta too, a date as text in the unit it is written to). Under eq and ne, a UDT of dates or durations pairs with one in another unit, as text of one kind pairs at another width, and they compare as numpy compares them. Behavior change (before -> after), with day a Vector of an M8[D] UDT holding 2020-01-01 and five one of an m8[s] UDT holding 5 s: - day == np.datetime64("2020-01-01T12:00"), datetime(2020, 1, 1, 12) or "2020-01-01T12:00": True -> False; != the opposite; the same for day[0].new() ==. - five == np.timedelta64(5000, "ms"): KeyError -> True; 5001 ms: KeyError -> False. five == timedelta(seconds=5, microseconds=1): True -> False. - binary.eq(five & a Vector of m8[ms] holding 5000 ms), and five.isequal of it: KeyError -> True. - day.apply(user_op, right=datetime(2020, 1, 1, 12)): stored the day, silently -> ValueError "does not fit". - five.apply(user_op, right=np.timedelta64(5000, "ms")): converted with a DeprecationWarning -> converts with none, as in a record's field. - day == timedelta(5): numpy's ValueError -> binary.eq's KeyError. --- graphblas/core/literals.py | 9 +++++- graphblas/core/operator/binary.py | 2 +- graphblas/core/udt.py | 11 ++++---- graphblas/tests/test_op.py | 46 +++++++++++++++++++++++++------ 4 files changed, 52 insertions(+), 16 deletions(-) diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index 97e301416..bb0ecc724 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -233,7 +233,10 @@ def _literal_type(dtype, value, *, exact=False): (:func:`_array_literal_udt`); a sequence holding numpy values is the array numpy makes of it. Text beside a UDT of text is typed as in a record's text field: the UDT itself when it converts there and back - unchanged, else its own type at its own length, which does not fit. + unchanged, else its own type at its own length, which does not fit. A + date or duration beside a UDT of dates or durations is typed so too, its + own type being numpy's for it (:func:`_as_numpy_date`), so ``==`` + compares it as numpy does. Anything else has no type here, and ``Scalar.from_value`` reports it. """ value = _as_python_number(value) @@ -245,6 +248,10 @@ def _literal_type(dtype, value, *, exact=False): if _keeps_value(value, np_type): return dtype return lookup_dtype(np.dtype(f"{np_type.kind}{max(len(value), 1)}")) + if np_type.kind in "mM" and (date := _as_numpy_date(value, np_type.kind)) is not None: + if date.dtype.kind == np_type.kind and _keeps_value(date, np_type): + return dtype + return lookup_dtype(date.dtype) if type(value) in _LITERAL_KINDS or isinstance(value, (tuple, list)): literal_type = _weak_literal_udt(dtype, value, exact=exact) if literal_type is not None or type(value) in _LITERAL_KINDS: diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 4386587aa..e17889d2a 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -87,7 +87,7 @@ def _make_udt_comparison(dtype, dtype2, *, is_eq): # Bytes, a datetime, or a complex number where the backend has no # complex types: an element is one value, compared as numpy compares # it, as a record field of that type is. ``_check_udt_pair`` has made - # sure both are the one UDT, or text of one kind. + # sure both are the one UDT, text of one kind, or dates or durations. wrapper = _compile_codegen( f"def wrapper(z_ptr, x_ptr, y_ptr):\n z_ptr[0] = x_ptr[0] {cmp_op} y_ptr[0]\n", func_name="wrapper", diff --git a/graphblas/core/udt.py b/graphblas/core/udt.py index 90dbcea47..41f4ceb38 100644 --- a/graphblas/core/udt.py +++ b/graphblas/core/udt.py @@ -113,25 +113,26 @@ def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): ``None`` (a built-in scalar side broadcasts), no further check fires. A UDT that is neither a record nor an array (bytes, say) holds one value, which pairs only with its own type, on either side; for ``eq`` and ``ne`` - only, text pairs with text of the same kind at another width, as numpy - and a record's text fields compare it. Without this gate the generated + only, text pairs with text of the same kind at another width, and a date + or duration with one in another unit, as numpy and a record's fields + compare them. Without this gate the generated wrapper code raises a cryptic Numba ``TypingError`` on the first field access, or, for array fields that do not broadcast, fails inside the cfunc without telling the caller. Element dtypes may differ; :func:`_udt_op_types` promotes them. """ - texts = ( + compared_as_numpy = ( op_name in {"eq", "ne"} and info_x is None and info_y is None and dtype._is_udt and dtype2._is_udt and dtype.np_type.kind == dtype2.np_type.kind - and dtype.np_type.kind in "SU" + and dtype.np_type.kind in "SUmM" ) if ( dtype != dtype2 - and not texts + and not compared_as_numpy and (info_x is None and dtype._is_udt or info_y is None and dtype2._is_udt) ): raise KeyError( diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 9d18b4afc..3a572d71b 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4679,13 +4679,6 @@ def test_udt_literal_converts_only_when_exact(): expr() with pytest.warns(DeprecationWarning, match="does not fit _ExactText"): assert t.apply(second, right=text_literal).new()[0].new().value.tolist() == (b"cd", 2) - with pytest.warns(DeprecationWarning, match="does not fit _ExactBareSeconds"): - assert d.apply(same, right=np.timedelta64(5000, "ms")).new()[0].new().value - # NaT stays NaT in any unit, and equals nothing. - with pytest.warns(DeprecationWarning, match="does not fit _ExactBareSeconds"): - assert np.isnat(d.apply(second, right=np.timedelta64("NaT", "ms")).new()[0].new().value) - with pytest.warns(DeprecationWarning, match="does not fit _ExactBareSeconds"): - assert not d.apply(same, right=np.timedelta64("NaT", "ms")).new()[0].new().value # Text for a text leaf converts, bytes or str. for literal in [b"cd", "cd"]: assert t.apply(second, right=(literal, 1)).new()[0].new().value.tolist() == (b"cd", 1) @@ -4699,7 +4692,10 @@ def test_udt_literal_converts_only_when_exact(): assert fine["ex_d"] == np.datetime64(5, "s") nat = when.apply(second, right=(np.datetime64("NaT", "ms"), 1)).new()[0].new().value assert np.isnat(nat["ex_d"]) - # A duration in a coarser unit casts safely, so it fits by type. + # A UDT of durations alone too; NaT stays NaT in any unit, and equals nothing. + assert d.apply(same, right=np.timedelta64(5000, "ms")).new()[0].new().value + assert np.isnat(d.apply(second, right=np.timedelta64("NaT", "ms")).new()[0].new().value) + assert not d.apply(same, right=np.timedelta64("NaT", "ms")).new()[0].new().value assert np.isnat(d.apply(second, right=np.timedelta64("NaT", "m")).new()[0].new().value) assert (t == (b"ab", 1)).new()[0].new().value assert not (t == (b"ab", 1.5)).new()[0].new().value @@ -6065,7 +6061,8 @@ def test_udt_eq_compares_a_udt_that_is_neither_record_nor_array(): Such a UDT holds one value per element, which ``binary.eq`` and ``binary.ne`` compare as a record field of that type compares: bytes by value, and NaT equal to nothing. Only two elements of the one type - compare; any other pair raises binary.eq's KeyError. + compare, or text of one kind at two widths, or dates or durations in two + units; any other pair raises binary.eq's KeyError. """ text = dtypes.register_anonymous(np.dtype("S4"), "_EqOpaqueBytes") v = Vector(text, size=3) @@ -6140,6 +6137,37 @@ def test_udt_eq_compares_a_udt_that_is_neither_record_nor_array(): d[0] = np.datetime64("NaT", "s") d[1] = np.datetime64(0, "s") assert (d == d).new().to_coo()[1].tolist() == [False, True] + # Under eq and ne, a date or duration pairs with one in another unit too, as + # numpy compares them, so a literal in a finer unit is not cut to the element's. + import datetime + + days = dtypes.register_anonymous(np.dtype("M8[D]"), "_EqOpaqueDays") + day = Vector(days, size=1) + day[0] = np.datetime64("2020-01-01") + seconds = dtypes.register_anonymous(np.dtype("m8[s]"), "_EqOpaqueSeconds") + five = Vector(seconds, size=1) + five[0] = np.timedelta64(5, "s") + for x, literal, equal in [ + (day, np.datetime64("2020-01-01T12:00"), False), + (day, datetime.datetime(2020, 1, 1, 12), False), # noqa: DTZ001 (numpy's are naive) + (day, "2020-01-01T12:00", False), + (day, np.datetime64("2020-01-01T00"), True), + (day, datetime.date(2020, 1, 1), True), + (day, "2020-01-01", True), + (five, np.timedelta64(5001, "ms"), False), + (five, datetime.timedelta(seconds=5, microseconds=1), False), + (five, np.timedelta64(5000, "ms"), True), + (five, datetime.timedelta(seconds=5), True), + ]: + assert (x == literal).new()[0].new().value == equal + assert (x != literal).new()[0].new().value != equal + assert (x[0].new() == literal) is equal + in_ms = Vector(dtypes.register_anonymous(np.dtype("m8[ms]"), "_EqOpaqueMillis"), size=1) + in_ms[0] = np.timedelta64(5000, "ms") + assert binary.eq(five & in_ms).new()[0].new().value + assert five.isequal(in_ms) + with pytest.raises(KeyError, match="pairs only with its own type"): + day == np.timedelta64(5, "s") # noqa: B015 # Nor with a record or an array UDT, for a Scalar as for a Vector. pair = dtypes.register_anonymous(np.dtype((np.int8, (5,))), "_EqOpaqueArr5") rec = dtypes.register_anonymous(np.dtype([("eqo_a", np.int8)]), "_EqOpaqueRec") From 2385b588a8119b3f34d4babbac7d5ad907f9b860 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 07:48:58 -0800 Subject: [PATCH 50/54] Type a literal beside a typed op in the op's own input type A literal operand beside a UDT became an element of that UDT even when the op was typed and takes that operand in another type: op[A, INT64] takes its thunk, or its second input, as an INT64. Apply and select pass a UDT thunk or operand to SuiteSparse as a pointer to its bytes, so the op read the bytes of the A element as an int64: silently wrong answers. Main did the same. A Scalar of the op's type gave the right answer. Now a typed op's own input type on the literal's side decides: its first input type for a left operand, else its second input or thunk type. GraphBLAS converts a value to a built-in input type, so the literal is typed as beside a built-in typed op (its own type: 2.5 for an INT64 thunk is converted to 2, as for an op typed on built-in types); a UDT input type takes only a literal that fits it. _literal_operand takes left=True from apply's left operand and from ewise_union's left default. Behavior change (before -> after), with v a Vector of a record A of two float64 fields holding (1.0, 2.0) and (3.0, 4.0), and user ops typed [A, INT64] (or [INT64, A]) that compare or add field a with the other: - v.apply(iu[A, INT64], 2): [False, False] -> [False, True], as with Scalar.from_value(2); the same for thunks 2.5 and np.int8(2). - v.select(se[A, INT64], 2), on a Vector or Matrix: empty -> index 1. - v.apply(f[A, INT64], right=2) and v.apply(g[INT64, A], left=2): about 4.6e18 for each element -> [3.0, 5.0]. - With B a record of two int64 fields: v.apply(h[A, B], right=(2, 3)): about 4.6e18 -> [3.0, 5.0]; v.apply(iu_rec[A, B], (2, 0)): [False, False] -> [False, True], and with (2.5, 0): [False, False] -> ValueError "does not fit". - Scalar(A).ewise_mult(2, f[A, INT64]): GraphBLAS's DomainMismatch -> 3.0. --- graphblas/core/literals.py | 11 +++++++++- graphblas/core/matrix.py | 2 +- graphblas/core/scalar.py | 4 +++- graphblas/core/vector.py | 2 +- graphblas/tests/test_op.py | 41 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 56 insertions(+), 4 deletions(-) diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index bb0ecc724..389b07332 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -38,13 +38,20 @@ ) -def _literal_operand(dtype, value, op): +def _literal_operand(dtype, value, op, *, left=False): """Return literal ``value`` and the dtype to make it, as ``op``'s operand beside ``dtype``. A dtype of ``None`` means the value's own dtype. This is where python-graphblas decides how a literal operand is typed; beside a built-in dtype, see :func:`_weak_builtin_literal_dtype`. + A typed op (``op[T]``) takes each operand in its own input type, so the + literal is typed beside that instead of beside ``dtype``: the op's first + input type when the literal is its left operand (``left``), else its + second input or thunk type. GraphBLAS converts a value to a built-in input + type; a UDT input type it cannot, so the literal becomes an element of + that UDT. + Beside a UDT, a literal becomes an element of that UDT, as it must for a user's op, when its type allows it (:func:`_check_literal_fits`). With a built-in op that lifts to UDTs, a literal is an operand in its own right, @@ -76,6 +83,8 @@ def _literal_operand(dtype, value, op): """ if getattr(op, "is_positional", False): return value, None + if isinstance(op, TypedOpBase): + dtype = op.type if left or op._type2 is None else op._type2 if not dtype._is_udt: literal_dtype = _weak_builtin_literal_dtype(dtype, value, op) if ( diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 37707f3b2..166f263a0 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -2459,7 +2459,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = None elif right is None: if type(left) is not Scalar: - left, dtype = _literal_operand(self.dtype, left, op) + left, dtype = _literal_operand(self.dtype, left, op, left=True) try: left = Scalar.from_value(left, dtype, is_cscalar=None, name="") except TypeError: diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index c44a97c40..bb7339ce8 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -1202,7 +1202,9 @@ def _union_default(obj, default, input_dtype, op, keyword_name): if dtype is not None: _check_literal_fits(dtype, default) else: - default, dtype = _literal_operand(input_dtype, default, op) + default, dtype = _literal_operand( + input_dtype, default, op, left=keyword_name == "left_default" + ) try: return Scalar.from_value( default, diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index 0b980f7f1..f4b72fca4 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -1441,7 +1441,7 @@ def apply(self, op, right=None, *, left=None): expr_repr = None elif right is None: if type(left) is not Scalar: - left, dtype = _literal_operand(self.dtype, left, op) + left, dtype = _literal_operand(self.dtype, left, op, left=True) try: left = Scalar.from_value(left, dtype, is_cscalar=None, name="") except TypeError: diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 3a572d71b..64cd1b604 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4985,6 +4985,47 @@ def test_udt_record_date_literals_compare_as_numpy(): assert result[0].new().value.tolist() == (day, secs) +@pytest.mark.skipif("not supports_udfs") +def test_udt_typed_op_takes_a_literal_in_its_own_input_type(): + """A literal beside a typed op is typed by the op's input type on its side, not by the UDT's. + + ``op[A, INT64]`` takes its second input or thunk as INT64, so a literal + ``2`` there is an INT64, as a Scalar would be, and ``2.5`` is converted to + it by GraphBLAS, as for a typed op on built-in types. Before, the literal + became an element of A, whose bytes the op read as an int64. + """ + A = dtypes.register_anonymous(np.dtype([("tl_a", "f8"), ("tl_b", "f8")]), "_TypedLitA") + B = dtypes.register_anonymous(np.dtype([("tl_a", "i8"), ("tl_b", "i8")]), "_TypedLitB") + v = Vector(A, size=2) + v[0] = (1.0, 2.0) + v[1] = (3.0, 4.0) + M = Matrix(A, nrows=2, ncols=1) + M[0, 0] = (1.0, 2.0) + M[1, 0] = (3.0, 4.0) + iu = IndexUnaryOp.register_anonymous( + lambda x, i, j, th: x.tl_a > th, "_typed_lit_iu", is_udt=True + ) + iu_rec = IndexUnaryOp.register_anonymous( + lambda x, i, j, th: x.tl_a > th.tl_a, "_typed_lit_iu_rec", is_udt=True + ) + se = SelectOp.register_anonymous(lambda x, i, j, th: x.tl_a > th, "_typed_lit_se", is_udt=True) + right = BinaryOp.register_anonymous(lambda x, y: x.tl_a + y, "_typed_lit_right", is_udt=True) + left = BinaryOp.register_anonymous(lambda x, y: x + y.tl_a, "_typed_lit_left", is_udt=True) + for thunk in [2, 2.5, np.int8(2), gb.Scalar.from_value(2)]: + assert v.apply(iu[A, INT64], thunk).new().to_coo()[1].tolist() == [False, True] + assert v.select(se[A, INT64], thunk).new().to_coo()[0].tolist() == [1] + assert M.select(se[A, INT64], thunk).new().to_coo()[0].tolist() == [1] + assert v[0].new().apply(iu[A, INT64], thunk).new().value is False + assert v.apply(iu_rec[A, B], (2, 0)).new().to_coo()[1].tolist() == [False, True] + assert v.apply(right[A, INT64], right=2).new().to_coo()[1].tolist() == [3, 5] + assert v.apply(left[INT64, A], left=2).new().to_coo()[1].tolist() == [3, 5] + assert M.apply(left[INT64, A], left=2).new().to_coo()[2].tolist() == [3, 5] + assert v[1].new().ewise_mult(2, right[A, INT64]).new().value == 5 + # A UDT input type still takes only a literal that fits it. + with pytest.raises(ValueError, match="does not fit _TypedLitB"): + v.apply(iu_rec[A, B], (2.5, 0)) + + @pytest.mark.skipif("not supports_udfs") def test_udt_ewise_union_defaults_are_typed_not_converted(): """An ``ewise_union`` default beside a UDT is typed as a literal and must fit the operand. From 7afa685b8dddb8bf9b4a6dd223752824368f8194 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 07:51:56 -0800 Subject: [PATCH 51/54] Accumulate into a UDT Scalar as into a built-in Scalar The accumulator is typed on the object's dtype, and for a built-in Scalar GraphBLAS converts the value to that type before it accumulates: an INT8 Scalar of -1 with s += 0.5 (Scalar.__iadd__ is s(binary.plus) << 0.5) adds 0 and stays -1, as v(accum=binary.plus) << 0.5 does for a Vector of either kind. A UDT Scalar used the untyped op instead (_scalar_accum_op), so it computed -1 + 0.5 in FP64 and truncated the sum to 0. That was meant to match a Vector's v += 0.5, but that is v << v + 0.5, not an accumulate. Main gave the built-in answer. Now a UDT Scalar accumulates with the typed op, and the value is converted to the Scalar's type first, as storing it would convert it: a literal by Scalar.from_value, a Scalar or an element of another UDT by the store cast. A built-in Scalar into a UDT Scalar still raises DomainMismatch, as a store of it does, into a Vector too. Behavior change (before -> after), with s an INT8[3] UDT Scalar of [-1, -1, -1]: - s += 0.5 and s(accum=binary.plus) << 0.5: [0, 0, 0] -> [-1, -1, -1] (an INT8 Scalar gives -1; a Vector of the UDT gives [-1, -1, -1]). - s(accum=binary.plus) << (0.5, 1.5, 2.5): [0, 0, 1] -> [-1, 0, 1], as for a Vector of the UDT. - s(accum=binary.plus) << v[1], with v a Vector of an FP64[3] UDT holding [0.5, 1.5, 2.5]: [0, 0, 1] -> [-1, 0, 1], as with the same value in a Scalar of that UDT, which was cast first already. --- docs/user_guide/udt.rst | 5 +++- graphblas/core/base.py | 10 +++++--- graphblas/core/udt.py | 29 ++++++++++++++++------ graphblas/tests/test_op.py | 51 ++++++++++++++++++++++++++++++++++---- 4 files changed, 77 insertions(+), 18 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 3f3793144..2a0b86555 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -256,7 +256,10 @@ data. (``w << expr.new()``); to assign, convert the value with ``.dup(dtype=...)`` first. * - ...from a ``Scalar`` or one element - - Computed in its own type, then cast. + - Computed in its own type, then cast. With an accumulator, the value is + cast before it accumulates, as GraphBLAS converts a value of a built-in + dtype, and so is a literal: an ``INT8[3]`` Scalar's ``s += 0.5`` adds 0, + as an ``INT8`` Scalar's does. * - Between a UDT and a built-in dtype, or layouts that do not correspond - ``DomainMismatch``. diff --git a/graphblas/core/base.py b/graphblas/core/base.py index d4e76e159..e8940ad9b 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -10,7 +10,7 @@ from .literals import _literal_store_dtype from .mask import Mask from .operator import UNKNOWN_OPCLASS, binary_from_string, find_opclass, get_typed_op -from .udt import _fused_store_expr, _scalar_accum_op, _store_cast_op +from .udt import _fused_store_expr, _scalar_accum_operand, _store_cast_op from .utils import _Pointer, libget, output_type _recorder = ContextVar("recorder") @@ -343,7 +343,8 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * if expr._is_scalar and self._is_scalar: # Extract element (s << v[1]) if accum is not None: - self(**opts) << self.ewise_add(expr, _scalar_accum_op(self, accum)) + expr = _scalar_accum_operand(self, expr) + self(**opts) << self.ewise_add(expr, accum) return expr.parent._extract_element( expr.resolved_indexes, @@ -370,7 +371,7 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * return if self._is_scalar: if accum is not None: - self(**opts) << self.ewise_add(expr, _scalar_accum_op(self, accum)) + self(**opts) << self.ewise_add(expr, accum) return if opts: # Ignore opts for now @@ -392,7 +393,8 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * # s << (v @ v) expr = expr._to_expr() elif accum is not None: - self(**opts) << self.ewise_add(expr, _scalar_accum_op(self, accum)) + expr = _scalar_accum_operand(self, expr) + self(**opts) << self.ewise_add(expr, accum) return else: if opts: diff --git a/graphblas/core/udt.py b/graphblas/core/udt.py index 41f4ceb38..94e7b3136 100644 --- a/graphblas/core/udt.py +++ b/graphblas/core/udt.py @@ -497,15 +497,28 @@ def _fused_store_expr(expr, output_dtype): return rv -def _scalar_accum_op(scalar, accum): - """Return the op to accumulate into ``scalar`` with, as ``scalar.ewise_add`` runs it. - - ``__call__`` types ``accum`` on the object's dtype. For a UDT Scalar, that - would convert a literal into the UDT (``s += 0.5``), where a Vector - computes ``s + 0.5`` in its own type and casts as it stores; the untyped - op does the same here. +def _scalar_accum_operand(scalar, value): + """Return ``value`` to accumulate into ``scalar``, converted to its type if that is a UDT. + + ``__call__`` types the accumulator on the object's dtype, as GraphBLAS + needs, so for a built-in Scalar GraphBLAS converts the value to that type + before it accumulates: an INT8 Scalar's ``s += 0.5`` adds 0, as a Vector's + ``v(accum=binary.plus) << 0.5`` does. GraphBLAS cannot convert a value to a + UDT, so for a UDT Scalar it is converted here, as storing it would convert + it: a literal as ``Scalar.from_value`` converts it, and a Scalar or an + element by :func:`_cast_for_store`, which raises for a type that does not + cast. ``value`` is a literal, a Scalar, or an element to extract. """ - return accum.parent if scalar.dtype._is_udt else accum + if not scalar.dtype._is_udt: + return value + from .scalar import Scalar + + if output_type(value) is Scalar: + return _cast_for_store(value, scalar.dtype) + try: + return Scalar.from_value(value, scalar.dtype, is_cscalar=False, name="") + except TypeError: # not a scalar; ewise_add reports it + return value def _cast_for_store(value, dtype): diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 64cd1b604..9c866ed21 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -5787,9 +5787,8 @@ def test_udt_store_casts_reach_every_store_path(): ``agg.first`` into a Scalar of another UDT handed GraphBLAS the output as it was, which read the int8 bytes as float64 (24 bytes from a 3-byte - element). Indexed assignment of an expression skipped the cast, a UDT - Scalar's ``s += 0.5`` converted 0.5 into the UDT where a Vector casts the - sum, and ``dup`` failed for a dtype given as a numpy dtype or a string. + element). Indexed assignment of an expression skipped the cast, and + ``dup`` failed for a dtype given as a numpy dtype or a string. """ i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_PathI8") f8s = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_PathF8") @@ -5822,8 +5821,6 @@ def test_udt_store_casts_reach_every_store_path(): x[[0, 1]] << w[[0, 1]].new().dup(dtype=i8s) assert x.to_coo()[1].tolist() == [[0, 0, 0], [1, 1, 1], [4, 5, 6]] s = gb.Scalar.from_value([1, 2, 3], dtype=i8s) - s += 0.5 - np.testing.assert_array_equal(s.value, [1, 2, 3]) s += 1 np.testing.assert_array_equal(s.value, [2, 3, 4]) result = s.dup(dtype=np.dtype((np.float64, (3,)))) @@ -5833,6 +5830,50 @@ def test_udt_store_casts_reach_every_store_path(): s.dup(dtype="FP64") +@pytest.mark.skipif("not supports_udfs") +def test_udt_scalar_accumulates_its_value_converted_to_its_type(): + """A UDT Scalar's accumulator adds a value converted to the Scalar's type, as a built-in's does. + + The accumulator is typed on the Scalar's dtype, and GraphBLAS converts the + value to that type first: an INT8 Scalar of -1 ``+= 0.5`` stays -1, and so + does a Vector's ``v(accum=binary.plus) << 0.5``. A UDT Scalar used the + untyped op instead, so it computed -1 + 0.5 and truncated the sum to 0. + """ + i8s = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_AccumI8") + f8s = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_AccumF8") + halves = Vector(f8s, size=2) + halves[1] = [0.5, 1.5, 2.5] + + def added(value): + s = gb.Scalar.from_value([-1, -1, -1], dtype=i8s) + s(accum=binary.plus) << value + return s.value.tolist() + + def added_to_int8(value): + s = gb.Scalar.from_value(-1, dtype=INT8) + s(accum=binary.plus) << value + return s.value + + assert added(0.5) == [added_to_int8(0.5)] * 3 == [-1, -1, -1] + assert added(2.5) == [added_to_int8(2.5)] * 3 == [1, 1, 1] + assert added((0.5, 1.5, 2.5)) == [-1, 0, 1] + assert added(halves[1]) == [-1, 0, 1] + assert added(gb.Scalar.from_value([0.5, 1.5, 2.5], dtype=f8s)) == [-1, 0, 1] + v = Vector(i8s, size=1) + v[0] = [-1, -1, -1] + v(accum=binary.plus) << (0.5, 1.5, 2.5) + assert v[0].new().value.tolist() == [-1, 0, 1] + s = gb.Scalar.from_value([-1, -1, -1], dtype=i8s) + s += 0.5 + assert s.value.tolist() == [-1, -1, -1] + s.clear() + s(accum=binary.plus) << 0.5 + assert s.value.tolist() == [0, 0, 0] + # A built-in Scalar casts only into a built-in type, as for a store. + with pytest.raises(DomainMismatch, match="a UDT casts only to another UDT"): + added(gb.Scalar.from_value(0.5)) + + @pytest.mark.skipif("not supports_udfs") def test_udt_literal_fits_in_select_bool_arrays_and_overflow(): """Three edges of the literal rules found by review. From 658b522f65c19d2a09aae1f98bdbcbf0cbbc209b Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 07:54:20 -0800 Subject: [PATCH 52/54] Pair an array literal with an array UDT when their shapes broadcast udt.rst says a numpy array beside a UDT pairs as an array UDT of its own shape, and array UDTs pair when their shapes broadcast. A numpy array paired only when it broadcast to the UDT's shape, or had that shape apart from leading 1s. Any other array, even one whose shape broadcasts with the UDT's, was converted into the UDT, where numpy raised "could not broadcast", while the same values as a UDT operand worked. Main raised too (a KeyError for the UDT pair). Now such an array is typed in its own shape and pairs as an array UDT operand does; the result has the shape numpy's would. An array that broadcasts to the UDT's shape is typed in the UDT's shape, as before, which a literal that must become an element needs. Behavior change (before -> after), with v a Vector of an INT8[3] UDT holding [1, 2, 3] and u one of an INT8[1] UDT holding [1]: - v + np.array([[1, 2, 3], [4, 5, 6]]), and apply(binary.plus) with it on either side: ValueError "could not broadcast input array from shape (2,3) into shape (3,)" -> INT64[2, 3] [[2, 4, 6], [5, 7, 9]], as with the same values in an INT64[2, 3] UDT. - v == np.array([[1, 2, 3], [1, 2, 3]]), on a Vector or a Scalar: the same ValueError -> True. - u + np.array([1, 2, 3]), and u + (1, 2, 3): ValueError -> INT64[3] [2, 3, 4]. - v + [[1, 2, 3], [4, 5, 6]]: numpy's "setting an array element with a sequence" ValueError -> INT64[2, 3], as for the array. - v.ewise_union(v, binary.plus, a 2-by-3 array, 0): numpy's ValueError -> ValueError "does not fit". --- graphblas/core/literals.py | 18 ++++++++++++------ graphblas/tests/test_op.py | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index 389b07332..810bde22c 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -805,11 +805,13 @@ def _array_literal_udt(dtype, value): An array has a dtype, so it is strong, as numpy 2 treats it: ``int8_udt + np.array([0.5, 1, 2])`` has float64 elements and ``+ np.array([1, 2, 3])`` int64 ones. Its type is the array layout of its - dtype, in its own shape when that pairs with the UDT's and in the UDT's - shape when it broadcasts to that (``np.array([1])`` beside ``INT8[3]``). - A 0-d structured array or ``np.void`` beside a record is its record dtype. - Anything else, such as a numeric array beside a record, is converted into - the UDT. + dtype: in the UDT's shape when it broadcasts to that (``np.array([1])`` + beside ``INT8[3]``), which a literal that must become an element needs, + else in its own shape when the two broadcast together, so that an op's + result takes the shape numpy's would (``INT8[3]`` and a 2-by-3 array give + ``[2, 3]``). A 0-d structured array or ``np.void`` beside a record is its + record dtype. Anything else, such as a numeric array beside a record, is + converted into the UDT. """ np_type = dtype.np_type if value.dtype.names is not None: @@ -825,5 +827,9 @@ def _array_literal_udt(dtype, value): try: np.broadcast_to(value, shape) except ValueError: - return None + try: + np.broadcast_shapes(value.shape, shape) + except ValueError: + return None + shape = value.shape return lookup_dtype(np.dtype((value.dtype, shape))) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 9c866ed21..d51993c2e 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -4945,6 +4945,39 @@ def test_udt_record_array_field_refuses_literal_of_another_shape(): assert value["shp_b"] == 1 +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_literal_broadcasts_with_an_array_udt(): + """An array literal pairs with an array UDT when their shapes broadcast, as a UDT operand does. + + The result has the shape numpy's would: ``INT8[3]`` and a 2-by-3 array + give ``[2, 3]``, as ``INT8[3]`` and an ``INT64[2, 3]`` UDT do. Before, an + array paired only when it broadcast to the UDT's shape; any other was + converted into the UDT and raised numpy's "could not broadcast" error. + """ + i8x3 = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_BcastI8x3") + i8x1 = dtypes.register_anonymous(np.dtype((np.int8, (1,))), "_BcastI8x1") + int64 = np.array([1]).dtype # int32 on Windows with numpy 1 + v = Vector(i8x3, size=1) + v[0] = [1, 2, 3] + u = Vector(i8x1, size=1) + u[0] = [1] + rows = np.array([[1, 2, 3], [4, 5, 6]]) + for expr in [lambda: v + rows, lambda: v.apply(binary.plus, left=rows)]: + result = expr().new() + assert result.dtype.np_type == np.dtype((int64, (2, 3))) + assert result[0].new().value.tolist() == [[2, 4, 6], [5, 7, 9]] + result = (u + np.array([1, 2, 3])).new() + assert result.dtype.np_type == np.dtype((int64, (3,))) + assert result[0].new().value.tolist() == [2, 3, 4] + assert (v + np.array([[1], [2]])).new()[0].new().value.tolist() == [[2, 3, 4], [3, 4, 5]] + assert (v == np.array([[1, 2, 3], [1, 2, 3]])).new()[0].new().value + assert not (v == rows).new()[0].new().value + assert v[0].new() == np.array([[1, 2, 3], [1, 2, 3]]) + # A literal that must become an element must still fit it. + with pytest.raises(ValueError, match="does not fit _BcastI8x3"): + v.ewise_union(v, binary.plus, rows, 0) + + def test_udt_record_date_literals_compare_as_numpy(): """A Python date or duration, or a date as text, in a record literal is typed as numpy types it. From 339fd930418031e3894ccae66002e5ef143eb04e Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 07:56:32 -0800 Subject: [PATCH 53/54] Give IndexBinaryOp is_positional, so dtype in op and op[dtype] work OpBase.__getitem__ reads is_positional for a type the op does not have, and IndexBinaryOp had none, so asking an IndexBinaryOp (registered without is_udt) for a type it does not take raised AttributeError. Main did the same; IndexBinaryOp has not been released. It is now False, as on a bound IndexBinaryOp and every other op class. Behavior change (before -> after), with op an IndexBinaryOp of x + y + theta and R a record UDT: - R in op: AttributeError -> False, as for an IndexUnaryOp. - op[R] and op(Scalar.from_value((1.0, 2), R), dtype=R): AttributeError -> KeyError "... does not work with R". --- graphblas/core/operator/indexbinary.py | 3 +++ graphblas/tests/test_indexbinary.py | 21 +++++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index 421c94cdf..2bf5110ab 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -330,6 +330,9 @@ class IndexBinaryOp(OpBase): _module = indexbinary _modname = "indexbinary" _custom_dtype = None + # Read by ``OpBase.__getitem__`` for a type the op does not have; no + # IndexBinaryOp is positional. + is_positional = False _typed_class = TypedBuiltinIndexBinaryOp _typed_user_class = TypedUserIndexBinaryOp # No built-in IndexBinaryOps; no parse config needed diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py index 304f7230a..704f55e3e 100644 --- a/graphblas/tests/test_indexbinary.py +++ b/graphblas/tests/test_indexbinary.py @@ -612,6 +612,27 @@ def test_bad_type(): indexbinary.register_anonymous(42) +def test_type_the_op_does_not_take(): + """``R in op`` is False and ``op[R]`` raises KeyError, as for an IndexUnaryOp. + + Both raised AttributeError: IndexBinaryOp had no ``is_positional``, which + ``op[R]`` reads for a type the op does not have. + """ + + def add_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + op = indexbinary.register_anonymous(add_vals) + rec = dtypes.register_anonymous( + np.dtype([("ibt_a", np.float64), ("ibt_b", np.int64)]), "_IbtRec" + ) + assert dtypes.INT64 in op + assert rec not in op + for expr in [lambda: op[rec], lambda: op(Scalar.from_value((1.0, 2), rec), dtype=rec)]: + with pytest.raises(KeyError, match="does not work with _IbtRec"): + expr() + + def test_with_mask(): def add_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) return x + y + theta From 82f33123323043cfa67be6551efcd4d24c8fd9da Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 07:58:03 -0800 Subject: [PATCH 54/54] Refuse a bound IndexBinaryOp in reduce with the usual TypeError reduce turns a BinaryOp that is part of a Monoid into that Monoid, and refuses any other BinaryOp with a TypeError naming it. It reads the op's monoid for that, and the error message its parent's; neither a bound IndexBinaryOp nor IndexBinaryOp had the attribute, so reduce raised AttributeError, for ibo[T](theta) on main and for the new untyped ibo(theta) too. Both now have monoid = None, as a BinaryOp of no Monoid has. IndexBinaryOp has not been released. Behavior change (before -> after), with v an INT64 Vector and ibo an IndexBinaryOp of x + y + theta: - v.reduce(ibo(1)) and v.reduce(ibo[int](1)): AttributeError "'_BoundIndexBinaryOp' object has no attribute 'monoid'" -> TypeError "Bad type for argument `op` in Vector.reduce(...) ... is not known to be part of a Monoid", as for binary.minus. - A.reduce_rowwise(ibo(1)): AttributeError -> NotImplementedException, as for binary.minus. --- graphblas/core/operator/indexbinary.py | 5 +++++ graphblas/tests/test_indexbinary.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index 2bf5110ab..4539387a9 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -113,6 +113,9 @@ class _BoundIndexBinaryOp(TypedOpBase): _semiring_commutes_to = None commutes_to = None is_commutative = False + # No Monoid is made of it, so reduce refuses it as it refuses a BinaryOp + # that is not part of one. + monoid = None # ``is_positional`` and ``_custom_dtype`` are BinaryOp-specific dispatch # hints: when a polymorphic BinaryOp is queried for a type it doesn't # support, the resolver in ``utils.get_typed_op`` falls back through @@ -333,6 +336,8 @@ class IndexBinaryOp(OpBase): # Read by ``OpBase.__getitem__`` for a type the op does not have; no # IndexBinaryOp is positional. is_positional = False + # Read by the error for a bound op where a Monoid is expected (reduce). + monoid = None _typed_class = TypedBuiltinIndexBinaryOp _typed_user_class = TypedUserIndexBinaryOp # No built-in IndexBinaryOps; no parse config needed diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py index 704f55e3e..333c35fcc 100644 --- a/graphblas/tests/test_indexbinary.py +++ b/graphblas/tests/test_indexbinary.py @@ -662,6 +662,23 @@ def add_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) assert C.to_coo()[2][0] == 118 # 100 + (3+5+10) +def test_reduce_refuses_a_bound_op(): + """A bound IndexBinaryOp gets the TypeError reduce gives a BinaryOp that is in no Monoid. + + It raised AttributeError: neither the bound op nor its IndexBinaryOp had + ``monoid``, which reduce and its error message read. + """ + + def add_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + op = indexbinary.register_anonymous(add_vals) + v = Vector.from_coo([0, 1], [1, 2]) + for binop in [op(1), op[int](1)]: + with pytest.raises(TypeError, match="is not known to be part of a Monoid"): + v.reduce(binop) + + def test_ewise_with_bound_binop(): """Confirm bound IndexBinaryOp works in all ewise operations."""