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: | 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/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/types.rst b/docs/user_guide/types.rst index aaeba1c6b..36d93e8a4 100644 --- a/docs/user_guide/types.rst +++ b/docs/user_guide/types.rst @@ -30,6 +30,33 @@ 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, 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`` +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/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/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index ae2cef9e2..2a0b86555 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``. @@ -134,12 +137,152 @@ 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. 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. +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. 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 + - 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``, 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`` + - 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. + - ``v == 300`` is ``False``; an ``FP32[3]`` UDT ``== 0.1`` compares in + 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 + 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, 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 +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. 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``. + +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 15f66bc2f..e8940ad9b 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_operand, _store_cast_op from .utils import _Pointer, libget, output_type _recorder = ContextVar("recorder") @@ -341,6 +343,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: + expr = _scalar_accum_operand(self, expr) self(**opts) << self.ewise_add(expr, accum) return expr.parent._extract_element( @@ -362,6 +365,10 @@ 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) @@ -386,6 +393,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: + expr = _scalar_accum_operand(self, expr) self(**opts) << self.ewise_add(expr, accum) return else: @@ -399,16 +407,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() @@ -418,7 +430,7 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * 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: @@ -446,6 +458,21 @@ 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: + 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) + 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/dtypes.py b/graphblas/core/dtypes.py index 93b3d4d8a..484dcf354 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 @@ -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,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_codegen 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 @@ -163,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. @@ -283,6 +332,12 @@ def register_anonymous(dtype, name=None): return rv if dtype.hasobject: raise ValueError("dtype must not allow Python objects") + if dtype.itemsize == 0: + # 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 @@ -320,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 @@ -566,14 +620,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 (``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. """ 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)) @@ -649,6 +704,43 @@ def _view_if_same_layout(values, np_type): return values +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. 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 + 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): + """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 _dtype_to_string(dtype): """Convert a dtype to a string that can be safely evaluated to recreate the dtype. diff --git a/graphblas/core/infixmethods.py b/graphblas/core/infixmethods.py index 67248223a..53f9d76f0 100644 --- a/graphblas/core/infixmethods.py +++ b/graphblas/core/infixmethods.py @@ -3,6 +3,7 @@ from .infix import MatrixInfixExpr, ScalarInfixExpr, VectorInfixExpr from .matrix import Matrix, MatrixExpression, MatrixIndexExpr, TransposedMatrix 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 @@ -14,6 +15,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/literals.py b/graphblas/core/literals.py new file mode 100644 index 000000000..810bde22c --- /dev/null +++ b/graphblas/core/literals.py @@ -0,0 +1,835 @@ +"""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 datetime +import math +import numbers +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 _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 +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, *, 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, + 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 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 ( + 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" + 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. + value = _beyond_as_infinity(np_type, parts) + literal_type = _literal_type(dtype, value, exact=exact) + if literal_type is not None: + 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(" 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. + + 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 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: + 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: + 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/core/matrix.py b/graphblas/core/matrix.py index 4625e4258..166f263a0 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -10,6 +10,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, @@ -19,6 +20,7 @@ get_typed_op, op_from_string, ) +from .operator.select import _value_select_in_own_type from .scalar import ( _COMPLETE, _MATERIALIZE, @@ -27,6 +29,15 @@ ScalarIndexExpr, _as_scalar, _scalar_index, + _union_default, +) +from .udt import ( + _cast_for_store, + _element_shapes_differ, + _ewise_add_cast_error, + _ewise_add_needs_cast, + _fused_store_expr, + _store_cast_op, ) from .utils import ( _CArray, @@ -100,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") @@ -107,6 +127,14 @@ 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. + 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. # 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. @@ -139,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 @@ -158,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): @@ -384,6 +412,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 @@ -394,6 +427,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: @@ -707,7 +742,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 @@ -1499,7 +1538,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, @@ -1921,14 +1961,22 @@ 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})." + ) + 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: 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})." + 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( method_name, None, @@ -2099,42 +2147,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: - 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: - 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: - 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: - 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") @@ -2419,6 +2433,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 +2459,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 + left, dtype = _literal_operand(self.dtype, left, op, left=True) try: left = Scalar.from_value(left, dtype, is_cscalar=None, name="") except TypeError: @@ -2476,7 +2495,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 + right, dtype = _literal_operand(self.dtype, right, op) try: right = Scalar.from_value(right, dtype, is_cscalar=None, name="") except TypeError: @@ -2596,8 +2615,10 @@ 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: - 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. + thunk, dtype = _literal_operand(self.dtype, thunk, op) try: thunk = Scalar.from_value(thunk, dtype, is_cscalar=None, name="") except TypeError: @@ -2610,6 +2631,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: @@ -2949,7 +2971,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="") @@ -3026,6 +3048,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 @@ -3088,7 +3117,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: @@ -3099,6 +3128,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]) @@ -3115,6 +3148,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 @@ -3311,7 +3346,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/base.py b/graphblas/core/operator/base.py index e4c41b95e..aea4a6f51 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -1,20 +1,15 @@ -import re from functools import lru_cache from operator import getitem from types import BuiltinFunctionType, ModuleType -import numpy as np - from ... import _STANDARD_OPERATOR_NAMES, backend, op -from ...dtypes import BOOL, INT8, UINT64, _supports_complex, lookup_dtype -from ...exceptions import UdfParseError, check_status_carg +from ...dtypes import BOOL, UINT64, _supports_complex, lookup_dtype +from ...exceptions import UdfParseError from .. import _has_numba, _supports_udfs, ffi, lib -from ..dtypes import _sample_values from ..expr import InfixExprBase from ..utils import output_type if _has_numba: - import numba from numba import NumbaError else: NumbaError = TypeError @@ -98,20 +93,6 @@ def _hasop(module, name): ) -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 - - class OpPath: def __init__(self, parent, name): self._parent = parent @@ -182,632 +163,13 @@ def _call_op(op, left, right=None, thunk=None, **kwargs): ) -if _has_numba: - - 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``. - """ - wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) - 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): - from .udt_utils import _iter_record_leaves - - n = len(numba_ret_type.types) - - def _leaves(d): - """Yield ``(python_path, c_path, leaf_dtype)`` for 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. - from .udt_utils import _iter_record_leaves - - 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" - f"{body_setup}" - f" _result = {call_expr}\n" - f"{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": - 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", @@ -1241,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/binary.py b/graphblas/core/operator/binary.py index 2cdab6eff..e17889d2a 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, @@ -61,85 +61,100 @@ 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_utils 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 "!=" + 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, 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", + source_label=f"", + ) + 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 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: 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) - ) - y_ptr_type = ( - nt.CPointer(numba.from_dtype(dtype2.np_type)) - if y_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.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) - 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):"] @@ -389,7 +404,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, @@ -617,21 +632,36 @@ 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 + # 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 +691,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( @@ -677,7 +713,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..4539387a9 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -1,19 +1,23 @@ 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 ...exceptions import EmptyObject, 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") if _has_numba: import numba - from .base import ( + from .udf import ( _bool_to_int8, _compile_udf_for_udt, _get_udt_wrapper_indexbinary, @@ -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,44 @@ 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 +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 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 +92,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 @@ -90,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 @@ -99,7 +125,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 +221,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 ------- @@ -140,6 +246,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), @@ -156,20 +263,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 +286,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 +316,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. @@ -225,6 +333,11 @@ 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 + # 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 @@ -361,7 +474,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,9 +484,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, *dtypes) if self._numba_func is None: raise KeyError(f"{self.name} does not work with {dtypes} types") @@ -385,15 +504,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")( @@ -407,7 +526,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, ), @@ -421,11 +540,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. @@ -517,8 +663,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. @@ -526,13 +672,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. """ @@ -541,22 +692,21 @@ 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): + _refuse_empty_theta(theta) + # 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 5820e66c5..a8eca2460 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, @@ -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") @@ -402,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/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/core/operator/select.py b/graphblas/core/operator/select.py index 3cc506d84..1c1138ae5 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -7,7 +7,138 @@ 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"} + + +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.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 + ``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; 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-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 + + 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 "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() + 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 + 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 + else: + # 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 + 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"}: + # 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 type(value) is float and 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): @@ -129,6 +260,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") @@ -298,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 new file mode 100644 index 000000000..ab8144fc2 --- /dev/null +++ b/graphblas/core/operator/udf.py @@ -0,0 +1,682 @@ +"""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 + 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) + + 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): # 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. + 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: # 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 + + +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): + # _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) + ] + 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, 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``, 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(theta_type, "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, 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):" + 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_codegen.py b/graphblas/core/operator/udt_codegen.py new file mode 100644 index 000000000..81d78f73f --- /dev/null +++ b/graphblas/core/operator/udt_codegen.py @@ -0,0 +1,1536 @@ +"""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 +import itertools +import linecache +from functools import lru_cache, reduce +from operator import mul + +import numpy as np + +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 + + +_codegen_counter = itertools.count(1) + + +def _compile_codegen(src, *, func_name, source_label, extra_ns=None): + """Compile a generated Python source string and return the named function. + + Centralizes the small amount of ``exec``-based code generation that the + UDT operator wrappers need (Numba has to see a literal function body, + not a closure, to type-check shape-specialized arithmetic). Three things + this helper does that a bare ``exec`` doesn't: + + 1. ``ast.parse`` runs first so a codegen typo raises a clear + ``RuntimeError`` with the offending source attached, at the call + site, rather than as a cryptic ``SyntaxError`` from ``exec`` or a + confusing ``TypingError`` from Numba's first compile. + 2. ``compile(..., filename, "exec")`` uses a human-readable synthetic + filename like ``" #7"``, and the + generated source is registered with ``linecache`` so any later + traceback shows real lines instead of ``:??``. + 3. The execution namespace is constructed here, so the names visible + to the generated code are auditable in one place. Extra entries + (e.g., the user's compiled ``numba_func`` for wrapper bodies) come + in via ``extra_ns``. + + Returns ``namespace[func_name]``. + """ + try: + ast.parse(src) + except SyntaxError as exc: + # Codegen bug, not user input; surface the source so a future + # maintainer can see exactly what was generated. + raise RuntimeError( + f"Generated code for {source_label!r} is not valid Python " + f"(parse error: {exc}). Source:\n{src}" + ) from exc + # Counter suffix so two codegens with the same label still get distinct + # cache keys (e.g., the same op compiled for two different UDTs that + # happen to print the same shape summary). + filename = f"{source_label} #{next(_codegen_counter)}" + linecache.cache[filename] = ( + len(src), + None, + src.splitlines(keepends=True), + filename, + ) + code = compile(src, filename, "exec") + # 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. 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) + return namespace[func_name] + + +# 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"} + +NP_TO_C_TYPES = { + np.dtype(np.bool_): "_Bool", + np.dtype(np.int8): "int8_t", + np.dtype(np.int16): "int16_t", + np.dtype(np.int32): "int32_t", + np.dtype(np.int64): "int64_t", + np.dtype(np.uint8): "uint8_t", + np.dtype(np.uint16): "uint16_t", + np.dtype(np.uint32): "uint32_t", + np.dtype(np.uint64): "uint64_t", + np.dtype(np.float32): "float", + np.dtype(np.float64): "double", + # Complex types are layout-compatible with numpy's; SS's GraphBLAS.h + # typedefs ``GxB_FC32_t`` / ``GxB_FC64_t`` to C99 ``float _Complex`` / + # ``double _Complex`` (or MSVC's ``_Fcomplex`` / ``_Dcomplex``). The + # 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; ``udt._element_types`` rejects them before any codegen. + np.dtype(np.complex64): "GxB_FC32_t", + np.dtype(np.complex128): "GxB_FC64_t", +} + + +# ``_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. +_has_jit_set = backend == "suitesparse" and hasattr(lib, "GxB_JIT_C_NAME") + +# Names that can't be used as UDT type names or record field names in the +# JIT C source. Covers two failure modes from the JIT include chain +# (``GraphBLAS.h`` transitively pulls in ````, ````, +# ````, ````, ````, ````): +# +# 1. **Macro names**: the preprocessor expands them inside any declarator +# or expression they appear in, so a field declared ``double M_PI ;`` +# becomes ``double 3.14159... ;`` and won't compile. The C-standard +# separate namespace for struct members doesn't help here. +# 2. **Typedef names** at the *outer* type-name position: emitting +# ``typedef struct { ... } FILE ;`` collides with the existing +# ``FILE`` typedef in scope (most compilers reject as a redefinition). +# Struct member names of the same spelling are fine (members live in +# their own namespace), but our gate runs the same check at both +# positions for simplicity. +# +# Either way the JIT compile fails, SuiteSparse swallows the failure +# silently, and the op runs through the slower Numba cfunc path. Block +# these eagerly so the user sees the "without JIT" warning instead. +_C_RESERVED = frozenset( + { + # C keywords + "auto", + "break", + "case", + "char", + "const", + "continue", + "default", + "do", + "double", + "else", + "enum", + "extern", + "float", + "for", + "goto", + "if", + "inline", + "int", + "long", + "register", + "restrict", + "return", + "short", + "signed", + "sizeof", + "static", + "struct", + "switch", + "typedef", + "union", + "unsigned", + "void", + "volatile", + "while", + "_Alignas", + "_Alignof", + "_Atomic", + "_Bool", + "_Complex", + "_Generic", + "_Imaginary", + "_Noreturn", + "_Static_assert", + "_Thread_local", + # C++ keywords that some compilers also reserve; cheap insurance + "bool", + "class", + "new", + "delete", + "template", + "namespace", + "this", + "true", + "false", + "nullptr", + # / : macros + the FILE typedef + "NULL", + "EOF", + "offsetof", + "stdin", + "stdout", + "stderr", + "FILE", + # / : typedefs at risk in the outer position + "size_t", + "ptrdiff_t", + "wchar_t", + "intptr_t", + "uintptr_t", + "intmax_t", + "uintmax_t", + "int8_t", + "int16_t", + "int32_t", + "int64_t", + "uint8_t", + "uint16_t", + "uint32_t", + "uint64_t", + # numeric macros + "INFINITY", + "NAN", + "HUGE_VAL", + "HUGE_VALF", + "HUGE_VALL", + "M_E", + "M_LOG2E", + "M_LOG10E", + "M_LN2", + "M_LN10", + "M_PI", + "M_PI_2", + "M_PI_4", + "M_1_PI", + "M_2_PI", + "M_2_SQRTPI", + "M_SQRT2", + "M_SQRT1_2", + # : ``complex`` and ``I`` in particular expand inside + # any field declarator. + "complex", + "imaginary", + "I", + "_Complex_I", + "_Imaginary_I", + "CMPLX", + "CMPLXF", + "CMPLXL", + # + "errno", + } +) + + +def _is_valid_c_identifier(name): + """True if ``name`` is a valid C identifier and not a reserved word.""" + return isinstance(name, str) and name.isidentifier() and name not in _C_RESERVED + + +# Process-global counter for synthesizing C names when the user-supplied +# Python name isn't a valid C identifier (or is a C reserved word, or the +# UDT was registered without ``name=`` at all). The Python-side ``DataType.name`` +# is left alone; this only affects ``GxB_JIT_C_NAME``. ``itertools.count`` is +# atomic enough for our purposes; collisions across processes are harmless +# because the SS JIT cache files are keyed by ``(c_name, content_hash)`` per +# 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. + + When ``python_name`` is already a valid C identifier (and not reserved), + return it unchanged so introspection and JIT cache filenames remain + readable. Otherwise mint a fresh ``_gbudt_NNN`` name. This lets UDTs + that were anonymous in the Python sense (no ``name=`` passed, name with + dots / special chars, name that collides with a C keyword) still take + the JIT path; users never see the synthetic name unless they read + ``DataType.jit_c_name``. + """ + if _is_valid_c_identifier(python_name): + return python_name + return f"_gbudt_{next(_synthetic_udt_counter)}" + + +def _udt_c_typedef(python_name, np_type): + """Build the JIT C typedef for a record or array UDT, or return ``None``. + + Returns ``(type_name, typedef)`` on success, where ``type_name`` is the + C identifier we register with SuiteSparse (may differ from + ``python_name`` if that wasn't a valid C identifier; see + ``_pick_c_type_name``). Returns ``None`` when the UDT still can't be + expressed in C after synthesizing the top-level name (a field name + collides with a C reserved word, or a leaf field type isn't in + ``NP_TO_C_TYPES``, e.g. object or datetime). + + Nested record UDTs (structured field whose dtype is itself a record) + are supported: each inner struct is emitted with a synthesized + ``_gbnest_NNN`` C name before the outer typedef in the same + ``GxB_JIT_C_DEFINITION`` string, so SuiteSparse's JIT C file ends up + with the full chain of typedefs in scope. Array UDTs are flattened to + a single-dimension C array so that the generated per-element operator + code can use flat indices like ``x->v[i]``, matching the + (row-major contiguous) numpy memory layout. + """ + # If the numpy layout disagrees with what a C compiler would produce, + # the JIT kernel would read fields at the wrong offsets. Skip rather + # than register a broken typedef; the cfunc path is still correct. + if np_type.names is not None and not _is_c_compatible_layout(np_type): + return None + type_name = _pick_c_type_name(python_name) + if np_type.names is not None: + # Inner typedefs (for nested structured fields) are prepended to the + # final definition string so the outer struct's field types are in + # scope when SuiteSparse compiles the kernel. + inner_typedefs = [] + fields = [] + for field_name in np_type.names: + # Field names still have to be valid C; synthesizing per-field + # names would mean tracking a name map for the codegen and would + # break the human-readable JIT source (``z->_f0 = ...``). Not + # worth it for the corner case where a user named a field + # ``"class"``; fall back to cfunc there. + if not _is_valid_c_identifier(field_name): + return None + field_dtype = np_type.fields[field_name][0] + if field_dtype.names is not None: + # Nested struct field. Recurse to build the inner typedef. + # The inner C name is synthesized to avoid collisions with + # user-supplied type names elsewhere in the process. + inner_name = f"_gbnest_{next(_synthetic_udt_counter)}" + inner_info = _udt_c_typedef(inner_name, field_dtype) + if inner_info is None: + return None + inner_resolved_name, inner_typedef = inner_info + inner_typedefs.append(inner_typedef) + fields.append(f"{inner_resolved_name} {field_name}") + continue + c_type = NP_TO_C_TYPES.get(field_dtype) + if c_type is None: + return None + fields.append(f"{c_type} {field_name}") + outer_typedef = f"typedef struct {{ {' ; '.join(fields)} ; }} {type_name} ;" + if inner_typedefs: + typedef = " ".join([*inner_typedefs, outer_typedef]) + else: + typedef = outer_typedef + return type_name, typedef + if np_type.subdtype is not None: + base_dtype, shape = np_type.subdtype + base_c_type = NP_TO_C_TYPES.get(base_dtype) + if base_c_type is None: + return None + size = reduce(mul, shape) + typedef = f"typedef struct {{ {base_c_type} v [{size}] ; }} {type_name} ;" + return type_name, typedef + return None + + +def _c_aligned_version(np_type): + """Return an ``align=True`` version of ``np_type`` with every nested + record also aligned recursively. + + ``np.dtype([..., (name, inner_struct)], align=True)`` respects the + *inner_struct*'s own declared alignment, so a packed inner struct stays + packed at its natural-1-byte alignment and the outer's ``coord`` field + lands at the outer's next-after-flag offset (4 for ``int32``), not at + the 8-byte boundary a C compiler would pick. To compare against what C + would emit, we need to rebuild the inner as ``align=True`` first, then + use that rebuilt type as the outer's field type. + """ + fields = [] + for name in np_type.names: + f = np_type.fields[name][0] + if f.names is not None: + f = _c_aligned_version(f) + fields.append((name, f)) + return np.dtype(fields, align=True) + + +def _is_c_compatible_layout(np_type): + """Return True iff ``np_type``'s layout matches what a C compiler will + pick for the corresponding ``typedef struct``. + + A user-supplied ``np.dtype([(name, dtype), ...])`` is *packed* by default + (no padding between fields). A C struct *aligns* fields to their natural + boundary. For ``[(int32, float64)]`` numpy uses offsets ``0, 4`` and + itemsize 12; C uses offsets ``0, 8`` and itemsize 16. The JIT-compiled + kernel reads fields at C offsets but the numpy buffer holds them at the + packed offsets, so the JIT would read garbage. Detect and refuse JIT in + that case; the user can either re-register with ``align=True``, use the + dict / dataclass form (which already does), or accept the cfunc path. + + Array UDTs are always C-compatible (single contiguous run of one type). + """ + if np_type.names is None: + return True + aligned = _c_aligned_version(np_type) + if aligned.itemsize != np_type.itemsize: + return False + for name in np_type.names: + if np_type.fields[name][1] != aligned.fields[name][1]: + return False + if not _is_c_compatible_layout(np_type.fields[name][0]): + return False + 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, 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. 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 + 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. + - ``*`` 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 + ``x / 0`` is an infinity, as there. + """ + if py_op in _FUNC_BINARY_OPS: + 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"_complex_truediv({x_expr}, {y_expr})" + return f"{x_expr} {py_op} {y_expr}" + + @numba.njit(error_model="numpy") + 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) + 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): # 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 + 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``. + + ``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 or py_op == "np.abs": + return f"{py_op}({operand})" + return f"{py_op}{operand}" + + 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. 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"{name}({expr})" + + def _make_record_func( + 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. + + ``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 (``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). + + 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. + 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, 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: + # 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})" + lines.append(f" return {ret}") + src = f"def _op({sig}):\n" + "\n".join(lines) + "\n" + op_func = _compile_codegen( + 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( + shape, + base_numba_type, + arity, + py_op, + *, + x_type, + y_type=None, + x_shape=None, + y_shape=None, + 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. 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. ``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) + # 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: + # 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 = "" + 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" + 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] = {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)) + 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 + + # MAINT 2026-05-21: this function, ``compile_udt_unary_wrapper`` below, + # and ``_make_jit_c_definition`` all branch on ``kind == "record"`` vs the + # 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, out=None): + """Compile a built-in element-wise binary op for a UDT. + + Handles two cases: + + - 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 ``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`). + + 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 .udf import _get_udt_wrapper + + info_x = _get_udt_info(dtype) + info_y = _get_udt_info(dtype2) + 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])." + ) + _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 + udt_dtype = dtype2 if x_is_scalar else dtype + + # 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 .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 + # fields, with paths like ``"['a']"``. + 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)) + 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, + 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 + # 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: + # Each side has its own element type, since the bases may differ. + 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(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.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 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 .udf 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): + """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``. + + ``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). + """ + 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)) + ) + 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, 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. + """ + from .udf import _get_udt_wrapper, _resolve_udt_return_type + + info = _get_udt_info(dtype) + if info is None: + raise KeyError( + f"unary.{op_name} does not work with {dtype}. " + f"Element-wise UDT ops require a record dtype (named fields) " + f"or an array dtype (e.g., FP64[3])." + ) + + kind, _detail = info + if kind == "record": + 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) + sig = (dtype.numba_type,) + _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(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 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 (``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; + ``udt._store_cast_op`` has already accepted the pair of types. + """ + from .udf 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_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 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 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_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}") # pragma: no cover (safety) + 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 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. + ``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. + + 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. 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) + 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_int_floordiv_expr(lhs, rhs, field_dtype): + """Return a C expression for integer floor division. + + Three ways C ``/`` differs from what the cfunc computes: + + - 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.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): + """Return a C expression for a unary op: e.g., ``-(x->a)``. + + ``field_dtype`` is only consulted for ``abs``: + + - Float fields use ``fabs`` / ``fabsf`` so ``abs(-0.0)`` returns + ``+0.0`` (matching Python). The naive ternary preserves the sign bit. + - Complex fields use ``cabs`` / ``cabsf``: the magnitude (a real). It's + assigned to a complex field via implicit ``double -> _Complex`` + conversion (``imag = 0``), matching Numba's behavior when Python + ``abs`` on a complex value is written back to a complex record field. + - Integer fields keep the ternary; ``abs`` of unsigned is a no-op and + ``abs`` of signed int wraps on INT_MIN, both matching the cfunc. + """ + if py_op == "abs": + 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})" + raise ValueError(f"Unknown unary C op: {py_op}") # pragma: no cover (safety) + + +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). + An op not defined on the field types (``min`` on a complex field) never + 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`` + # 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 + # use an undefined struct name, and SS would silently fall back to the + # Numba cfunc instead of JIT-compiling. ``jit_c_name`` is also where the + # synthetic ``_gbudt_NNN`` name (for UDTs whose Python name isn't a valid + # C identifier) is recorded; using it keeps op codegen consistent. + pinned_name = dtype.jit_c_name + typedef_info = _udt_c_typedef(pinned_name or dtype.name, dtype.np_type) + if typedef_info is None: + return None + type_name, _typedef = typedef_info + c_name = f"{op_name}_{type_name}" + + if arity == 2: + params = f"{type_name} *z, const {type_name} *x, const {type_name} *y" + else: + params = f"{type_name} *z, const {type_name} *x" + + if np_type.names is not None: + # Use leaf C paths so nested records emit ``z->outer.inner_a = ...`` + # alongside the inner+outer typedefs. Non-nested records degenerate + # to plain ``z->name``. + leaves = list(_iter_record_leaves(np_type)) + if arity == 2: + # Pass the leaf dtype to the binary statement builder so + # type-sensitive ops (min, max, and the divisions) emit correct C. + assigns = " ".join( + _c_assign_binary(py_op, f"z->{c}", f"x->{c}", f"y->{c}", leaf_dtype) + for _py, c, leaf_dtype in leaves + ) + else: + assigns = " ".join( + f"z->{c} = {_c_expr_unary(py_op, f'x->{c}', leaf_dtype)} ;" + for _py, c, leaf_dtype in leaves + ) + 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: + # 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: + 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} }}" + + +def _make_jit_c_comparison_definition(op_name, dtype, *, is_eq): + """Generate a JIT C function definition for ``binary.eq[udt]`` / + ``binary.ne[udt]``. + + Returns ``(c_name, c_defn)``, or ``None`` when the dtype can't be + 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 loop + over their elements. + + Array-valued sub-fields inside a record (e.g. + ``[("weights", (float64, 3))]``) aren't supported: ``_udt_c_typedef`` + rejects records with array sub-fields (``NP_TO_C_TYPES`` has no entry + for sub-array dtypes), so this function returns ``None`` before + iterating leaves in that case. Adding support would mean extending + ``_udt_c_typedef`` to emit ``double name[N]`` and ``_make_jit_c_definition`` + to unroll array sub-fields in arithmetic codegen too; both legs need + to land together or the eq/ne path is alone in supporting it. + + IEEE NaN propagation comes for free: C ``a == b`` is false when either + side is NaN, so two records both carrying NaN compare unequal under + ``eq`` (and equal under ``ne``), matching the cfunc path's leaf-wise + semantic comparison. + """ + np_type = dtype.np_type + pinned_name = dtype.jit_c_name + # ``_udt_c_typedef`` returns None when the numpy layout doesn't match a + # C compiler's layout for the same struct. + typedef_info = _udt_c_typedef(pinned_name or dtype.name, np_type) + if typedef_info is None: + return None + type_name, _typedef = typedef_info + c_name = f"{op_name}_{type_name}" + op = "==" if is_eq else "!=" + join = " && " if is_eq else " || " + if np_type.subdtype is not None: + # 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) + 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 ;" + ) + 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)} ;" + params = f"_Bool *z, const {type_name} *x, const {type_name} *y" + return c_name, f"void {c_name} ({params}) {{ {body} }}" + + +def _maybe_warn_jit_skipped(jit_info, op_name, dtype_name): + """Emit a ``NoJITWarning`` when JIT setup was skipped despite SS supporting JIT. + + Called by op-class ``_compile_udt`` paths after ``set_jit_c_*_on_op`` + returns ``None``. The ``_has_jit_set`` gate keeps the warning silent on + SS < 8 where JIT is fundamentally absent. + """ + if jit_info is None and _has_jit_set: + from ..ss.jit_config import _maybe_warn_no_jit + + _maybe_warn_no_jit(op_name=op_name, dtype_name=dtype_name) + + +def _set_jit_c_strings(gb_obj, c_name, c_defn, set_string_func): + """Pin ``GxB_JIT_C_NAME`` + ``GxB_JIT_C_DEFINITION`` on a fresh op handle. + + Both strings are one-shot on SuiteSparse: a second set returns + ``GrB_ALREADY_SET`` silently. The return is not checked here because + callers always pass a freshly-allocated handle. + + Arming an op with C source is the point at which this process has + committed to wanting a JIT compiler, so it is also where SuiteSparse's + non-compiling default gets raised. + """ + from ..ss.jit_config import _enable_jit_for_udt + + _enable_jit_for_udt() + set_string_func(gb_obj, ffi.new("char[]", c_name.encode()), lib.GxB_JIT_C_NAME) + set_string_func(gb_obj, ffi.new("char[]", c_defn.encode()), lib.GxB_JIT_C_DEFINITION) + + +def set_jit_c_comparison_on_op(gb_obj, op_name, dtype, set_string_func, *, is_eq): + """Generate and set JIT C name+definition for ``eq``/``ne`` on a UDT. + + 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 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`. + """ + result = _make_jit_c_comparison_definition(op_name, dtype, is_eq=is_eq) + if result is None: + return None + _set_jit_c_strings(gb_obj, *result, set_string_func) + return result + + +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 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. + """ + result = _make_jit_c_definition(op_name, py_op, dtype, arity) + if result is None: + return None + _set_jit_c_strings(gb_obj, *result, set_string_func) + return result diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py deleted file mode 100644 index ede25b467..000000000 --- a/graphblas/core/operator/udt_utils.py +++ /dev/null @@ -1,1003 +0,0 @@ -"""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. -""" - -import ast -import itertools -import linecache -from functools import reduce -from operator import mul - -import numpy as np - -from ... import backend -from .. import _has_numba, ffi, lib - -if _has_numba: - import numba - - -_codegen_counter = itertools.count(1) - - -def _compile_codegen(src, *, func_name, source_label, extra_ns=None): - """Compile a generated Python source string and return the named function. - - Centralizes the small amount of ``exec``-based code generation that the - UDT operator wrappers need (Numba has to see a literal function body, - not a closure, to type-check shape-specialized arithmetic). Three things - this helper does that a bare ``exec`` doesn't: - - 1. ``ast.parse`` runs first so a codegen typo raises a clear - ``RuntimeError`` with the offending source attached, at the call - site, rather than as a cryptic ``SyntaxError`` from ``exec`` or a - confusing ``TypingError`` from Numba's first compile. - 2. ``compile(..., filename, "exec")`` uses a human-readable synthetic - filename like ``" #7"``, and the - generated source is registered with ``linecache`` so any later - traceback shows real lines instead of ``:??``. - 3. The execution namespace is constructed here, so the names visible - to the generated code are auditable in one place. Extra entries - (e.g., the user's compiled ``numba_func`` for wrapper bodies) come - in via ``extra_ns``. - - Returns ``namespace[func_name]``. - """ - try: - ast.parse(src) - except SyntaxError as exc: - # Codegen bug, not user input; surface the source so a future - # maintainer can see exactly what was generated. - raise RuntimeError( - f"Generated code for {source_label!r} is not valid Python " - f"(parse error: {exc}). Source:\n{src}" - ) from exc - # Counter suffix so two codegens with the same label still get distinct - # cache keys (e.g., the same op compiled for two different UDTs that - # happen to print the same shape summary). - filename = f"{source_label} #{next(_codegen_counter)}" - linecache.cache[filename] = ( - len(src), - None, - src.splitlines(keepends=True), - filename, - ) - code = compile(src, filename, "exec") - namespace = {"min": min, "max": max, "abs": abs} - if _has_numba: - namespace["numba"] = numba - if extra_ns: - namespace.update(extra_ns) - exec(code, namespace) - 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"} - -NP_TO_C_TYPES = { - np.dtype(np.bool_): "_Bool", - np.dtype(np.int8): "int8_t", - np.dtype(np.int16): "int16_t", - np.dtype(np.int32): "int32_t", - np.dtype(np.int64): "int64_t", - np.dtype(np.uint8): "uint8_t", - np.dtype(np.uint16): "uint16_t", - np.dtype(np.uint32): "uint32_t", - np.dtype(np.uint64): "uint64_t", - np.dtype(np.float32): "float", - np.dtype(np.float64): "double", - # Complex types are layout-compatible with numpy's; SS's GraphBLAS.h - # typedefs ``GxB_FC32_t`` / ``GxB_FC64_t`` to C99 ``float _Complex`` / - # ``double _Complex`` (or MSVC's ``_Fcomplex`` / ``_Dcomplex``). The - # 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. - 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_NOT_FOR_COMPLEX = frozenset({"min", "max", "floordiv"}) - -# C operator equivalents for JIT code generation -_C_INFIX_OPS = {"+": "+", "-": "-", "*": "*", "/": "/", "//": "/"} - -# Vanilla strips GxB callables but keeps GxB constants, so the bare -# ``hasattr`` would lie; gate on the backend too. -_has_jit_set = backend == "suitesparse" and hasattr(lib, "GxB_JIT_C_NAME") - -# Names that can't be used as UDT type names or record field names in the -# JIT C source. Covers two failure modes from the JIT include chain -# (``GraphBLAS.h`` transitively pulls in ````, ````, -# ````, ````, ````, ````): -# -# 1. **Macro names**: the preprocessor expands them inside any declarator -# or expression they appear in, so a field declared ``double M_PI ;`` -# becomes ``double 3.14159... ;`` and won't compile. The C-standard -# separate namespace for struct members doesn't help here. -# 2. **Typedef names** at the *outer* type-name position: emitting -# ``typedef struct { ... } FILE ;`` collides with the existing -# ``FILE`` typedef in scope (most compilers reject as a redefinition). -# Struct member names of the same spelling are fine (members live in -# their own namespace), but our gate runs the same check at both -# positions for simplicity. -# -# Either way the JIT compile fails, SuiteSparse swallows the failure -# silently, and the op runs through the slower Numba cfunc path. Block -# these eagerly so the user sees the "without JIT" warning instead. -_C_RESERVED = frozenset( - { - # C keywords - "auto", - "break", - "case", - "char", - "const", - "continue", - "default", - "do", - "double", - "else", - "enum", - "extern", - "float", - "for", - "goto", - "if", - "inline", - "int", - "long", - "register", - "restrict", - "return", - "short", - "signed", - "sizeof", - "static", - "struct", - "switch", - "typedef", - "union", - "unsigned", - "void", - "volatile", - "while", - "_Alignas", - "_Alignof", - "_Atomic", - "_Bool", - "_Complex", - "_Generic", - "_Imaginary", - "_Noreturn", - "_Static_assert", - "_Thread_local", - # C++ keywords that some compilers also reserve; cheap insurance - "bool", - "class", - "new", - "delete", - "template", - "namespace", - "this", - "true", - "false", - "nullptr", - # / : macros + the FILE typedef - "NULL", - "EOF", - "offsetof", - "stdin", - "stdout", - "stderr", - "FILE", - # / : typedefs at risk in the outer position - "size_t", - "ptrdiff_t", - "wchar_t", - "intptr_t", - "uintptr_t", - "intmax_t", - "uintmax_t", - "int8_t", - "int16_t", - "int32_t", - "int64_t", - "uint8_t", - "uint16_t", - "uint32_t", - "uint64_t", - # numeric macros - "INFINITY", - "NAN", - "HUGE_VAL", - "HUGE_VALF", - "HUGE_VALL", - "M_E", - "M_LOG2E", - "M_LOG10E", - "M_LN2", - "M_LN10", - "M_PI", - "M_PI_2", - "M_PI_4", - "M_1_PI", - "M_2_PI", - "M_2_SQRTPI", - "M_SQRT2", - "M_SQRT1_2", - # : ``complex`` and ``I`` in particular expand inside - # any field declarator. - "complex", - "imaginary", - "I", - "_Complex_I", - "_Imaginary_I", - "CMPLX", - "CMPLXF", - "CMPLXL", - # - "errno", - } -) - - -def _is_valid_c_identifier(name): - """True if ``name`` is a valid C identifier and not a reserved word.""" - return isinstance(name, str) and name.isidentifier() and name not in _C_RESERVED - - -# Process-global counter for synthesizing C names when the user-supplied -# Python name isn't a valid C identifier (or is a C reserved word, or the -# UDT was registered without ``name=`` at all). The Python-side ``DataType.name`` -# is left alone; this only affects ``GxB_JIT_C_NAME``. ``itertools.count`` is -# atomic enough for our purposes; collisions across processes are harmless -# because the SS JIT cache files are keyed by ``(c_name, content_hash)`` per -# process state. -_synthetic_udt_counter = itertools.count(1) - - -def _pick_c_type_name(python_name): - """Return a valid C identifier for use as ``GxB_JIT_C_NAME`` on a UDT. - - When ``python_name`` is already a valid C identifier (and not reserved), - return it unchanged so introspection and JIT cache filenames remain - readable. Otherwise mint a fresh ``_gbudt_NNN`` name. This lets UDTs - that were anonymous in the Python sense (no ``name=`` passed, name with - dots / special chars, name that collides with a C keyword) still take - the JIT path; users never see the synthetic name unless they read - ``DataType.jit_c_name``. - """ - if _is_valid_c_identifier(python_name): - return 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. - """ - 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" 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 == "array" and detail_x != detail_y: - 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}." - ) - - -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``. - - Returns ``(type_name, typedef)`` on success, where ``type_name`` is the - C identifier we register with SuiteSparse (may differ from - ``python_name`` if that wasn't a valid C identifier; see - ``_pick_c_type_name``). Returns ``None`` when the UDT still can't be - expressed in C after synthesizing the top-level name (a field name - collides with a C reserved word, or a leaf field type isn't in - ``NP_TO_C_TYPES``, e.g. object or datetime). - - Nested record UDTs (structured field whose dtype is itself a record) - are supported: each inner struct is emitted with a synthesized - ``_gbnest_NNN`` C name before the outer typedef in the same - ``GxB_JIT_C_DEFINITION`` string, so SuiteSparse's JIT C file ends up - with the full chain of typedefs in scope. Array UDTs are flattened to - a single-dimension C array so that the generated per-element operator - code can use flat indices like ``x->v[i]``, matching the - (row-major contiguous) numpy memory layout. - """ - # If the numpy layout disagrees with what a C compiler would produce, - # the JIT kernel would read fields at the wrong offsets. Skip rather - # than register a broken typedef; the cfunc path is still correct. - if np_type.names is not None and not _is_c_compatible_layout(np_type): - return None - type_name = _pick_c_type_name(python_name) - if np_type.names is not None: - # Inner typedefs (for nested structured fields) are prepended to the - # final definition string so the outer struct's field types are in - # scope when SuiteSparse compiles the kernel. - inner_typedefs = [] - fields = [] - for field_name in np_type.names: - # Field names still have to be valid C; synthesizing per-field - # names would mean tracking a name map for the codegen and would - # break the human-readable JIT source (``z->_f0 = ...``). Not - # worth it for the corner case where a user named a field - # ``"class"``; fall back to cfunc there. - if not _is_valid_c_identifier(field_name): - return None - field_dtype = np_type.fields[field_name][0] - if field_dtype.names is not None: - # Nested struct field. Recurse to build the inner typedef. - # The inner C name is synthesized to avoid collisions with - # user-supplied type names elsewhere in the process. - inner_name = f"_gbnest_{next(_synthetic_udt_counter)}" - inner_info = _udt_c_typedef(inner_name, field_dtype) - if inner_info is None: - return None - inner_resolved_name, inner_typedef = inner_info - inner_typedefs.append(inner_typedef) - fields.append(f"{inner_resolved_name} {field_name}") - continue - c_type = NP_TO_C_TYPES.get(field_dtype) - if c_type is None: - return None - fields.append(f"{c_type} {field_name}") - outer_typedef = f"typedef struct {{ {' ; '.join(fields)} ; }} {type_name} ;" - if inner_typedefs: - typedef = " ".join([*inner_typedefs, outer_typedef]) - else: - typedef = outer_typedef - return type_name, typedef - if np_type.subdtype is not None: - base_dtype, shape = np_type.subdtype - base_c_type = NP_TO_C_TYPES.get(base_dtype) - if base_c_type is None: - return None - size = reduce(mul, shape) - typedef = f"typedef struct {{ {base_c_type} v [{size}] ; }} {type_name} ;" - return type_name, typedef - return None - - -def _c_aligned_version(np_type): - """Return an ``align=True`` version of ``np_type`` with every nested - record also aligned recursively. - - ``np.dtype([..., (name, inner_struct)], align=True)`` respects the - *inner_struct*'s own declared alignment, so a packed inner struct stays - packed at its natural-1-byte alignment and the outer's ``coord`` field - lands at the outer's next-after-flag offset (4 for ``int32``), not at - the 8-byte boundary a C compiler would pick. To compare against what C - 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] - if f.names is not None: - f = _c_aligned_version(f) - fields.append((name, f)) - return np.dtype(fields, align=True) - - -def _is_c_compatible_layout(np_type): - """Return True iff ``np_type``'s layout matches what a C compiler will - pick for the corresponding ``typedef struct``. - - A user-supplied ``np.dtype([(name, dtype), ...])`` is *packed* by default - (no padding between fields). A C struct *aligns* fields to their natural - boundary. For ``[(int32, float64)]`` numpy uses offsets ``0, 4`` and - itemsize 12; C uses offsets ``0, 8`` and itemsize 16. The JIT-compiled - kernel reads fields at C offsets but the numpy buffer holds them at the - packed offsets, so the JIT would read garbage. Detect and refuse JIT in - that case; the user can either re-register with ``align=True``, use the - dict / dataclass form (which already does), or accept the cfunc path. - - Array UDTs are always C-compatible (single contiguous run of one type). - """ - if np_type.names is None: - return True - aligned = _c_aligned_version(np_type) - if aligned.itemsize != np_type.itemsize: - return False - for name in np_type.names: - if np_type.fields[name][1] != aligned.fields[name][1]: - return False - if not _is_c_compatible_layout(np_type.fields[name][0]): - return False - 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.""" - if py_op in _FUNC_BINARY_OPS: - return f"{py_op}({x_expr}, {y_expr})" - return f"{x_expr} {py_op} {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: - 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): - """Build a Numba njit function for a record UDT. - - ``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 - 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). - - 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. - """ - 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}" - parts.append(_expr_binary(py_op, x_expr, y_expr)) - sig = "x, y" - else: - parts = [_expr_unary(py_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" - op_func = _compile_codegen( - src, - func_name="_op", - source_label=f"", - ) - return numba.njit(op_func, error_model="numpy") - - def _make_array_wrapper( - size, - base_numba_type, - arity, - py_op, - *, - x_scalar_type=None, - y_scalar_type=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). - - Returns (wrapper_func, wrapper_sig). - """ - nt = numba.types - 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) - ) - 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) - else: - assigns = "\n".join( - f" z[{i}] = {_expr_unary(py_op, f'x[{i}]')}" for i in range(size) - ) - 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" - op_func = _compile_codegen( - src, - func_name="_op", - source_label=f"", - ) - return op_func, sig - - # MAINT 2026-05-21: this function, ``compile_udt_unary_wrapper`` below, - # and ``_make_jit_c_definition`` all branch on ``kind == "record"`` vs the - # 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): - """Compile a built-in element-wise binary op for a UDT. - - Handles two cases: - - - Both sides are the same UDT: 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). - - 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 - - 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: - 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." - ) - - 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 - - if kind == "record": - 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)] - func = _make_record_func( - leaf_paths, - 2, - py_op, - x_is_scalar=x_is_scalar, - y_is_scalar=y_is_scalar, - ) - 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) - 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 - wrapper, wrapper_sig = _make_array_wrapper( - size, - numba.from_dtype(base_dtype), - 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, - ) - return wrapper, wrapper_sig, ret_type - - def compile_udt_unary_wrapper(op_name, py_op, dtype): - """Compile a built-in element-wise unary op for a UDT. - - 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 - - info = _get_udt_info(dtype) - if info is None: - raise KeyError( - f"unary.{op_name} does not work with {dtype}. " - f"Element-wise UDT ops require a record dtype (named fields) " - f"or an array dtype (e.g., FP64[3])." - ) - - kind, detail = info - if kind == "record": - from .base 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) - sig = (dtype.numba_type,) - _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) - wrapper, wrapper_sig = _get_udt_wrapper( - func, ret_type, dtype, numba_ret_type=numba_ret_type - ) - else: - base_dtype, size = detail - ret_type = dtype - wrapper, wrapper_sig = _make_array_wrapper(size, numba.from_dtype(base_dtype), 1, py_op) - return wrapper, wrapper_sig, ret_type - - -# JIT C code generators below. - - -def _c_expr_binary(py_op, lhs, rhs, field_dtype=None): - """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. - """ - 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 == "//": - return _c_floordiv_expr(lhs, rhs, field_dtype) - c_op = _C_INFIX_OPS.get(py_op, py_op) - return f"({lhs}) {c_op} ({rhs})" - - -def _c_floordiv_expr(lhs, rhs, field_dtype): - """Return a C expression for Python-semantics 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. - - 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. - """ - 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))))" - - -def _c_expr_unary(py_op, operand, field_dtype=None): - """Return a C expression for a unary op: e.g., ``-(x->a)``. - - ``field_dtype`` is only consulted for ``abs``: - - - Float fields use ``fabs`` / ``fabsf`` so ``abs(-0.0)`` returns - ``+0.0`` (matching Python). The naive ternary preserves the sign bit. - - Complex fields use ``cabs`` / ``cabsf``: the magnitude (a real). It's - assigned to a complex field via implicit ``double -> _Complex`` - conversion (``imag = 0``), matching Numba's behavior when Python - ``abs`` on a complex value is written back to a complex record field. - - Integer fields keep the ternary; ``abs`` of unsigned is a no-op and - ``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})" - return f"(({operand}) < 0 ? -({operand}) : ({operand}))" - if py_op == "-": - return f"-({operand})" - raise ValueError(f"Unknown unary C op: {py_op}") - - -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). - """ - 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 - # use an undefined struct name, and SS would silently fall back to the - # Numba cfunc instead of JIT-compiling. ``jit_c_name`` is also where the - # synthetic ``_gbudt_NNN`` name (for UDTs whose Python name isn't a valid - # C identifier) is recorded; using it keeps op codegen consistent. - pinned_name = dtype.jit_c_name - typedef_info = _udt_c_typedef(pinned_name or dtype.name, dtype.np_type) - if typedef_info is None: - return None - type_name, _typedef = typedef_info - c_name = f"{op_name}_{type_name}" - - if arity == 2: - params = f"{type_name} *z, const {type_name} *x, const {type_name} *y" - else: - params = f"{type_name} *z, const {type_name} *x" - - if np_type.names is not None: - # Use leaf C paths so nested records emit ``z->outer.inner_a = ...`` - # alongside the inner+outer typedefs. Non-nested records degenerate - # 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. - assigns = " ".join( - f"z->{c} = {_c_expr_binary(py_op, f'x->{c}', f'y->{c}', leaf_dtype)} ;" - for _py, c, leaf_dtype in leaves - ) - else: - assigns = " ".join( - f"z->{c} = {_c_expr_unary(py_op, f'x->{c}', leaf_dtype)} ;" - for _py, c, leaf_dtype in leaves - ) - else: # array UDT, flattened to v[size] - base_dtype, shape = np_type.subdtype - size = reduce(mul, shape) - 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) - ) - else: - assigns = " ".join( - f"z->v[{i}] = {_c_expr_unary(py_op, f'x->v[{i}]', base_dtype)} ;" - for i in range(size) - ) - return c_name, f"void {c_name} ({params}) {{ {assigns} }}" - - -def _make_jit_c_comparison_definition(op_name, dtype, *, is_eq): - """Generate a JIT C function definition for ``binary.eq[udt]`` / - ``binary.ne[udt]``. - - Returns ``(c_name, c_defn)``, or ``None`` when the dtype can't be - 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. - - Array-valued sub-fields inside a record (e.g. - ``[("weights", (float64, 3))]``) aren't supported: ``_udt_c_typedef`` - rejects records with array sub-fields (``NP_TO_C_TYPES`` has no entry - for sub-array dtypes), so this function returns ``None`` before - iterating leaves in that case. Adding support would mean extending - ``_udt_c_typedef`` to emit ``double name[N]`` and ``_make_jit_c_definition`` - to unroll array sub-fields in arithmetic codegen too; both legs need - to land together or the eq/ne path is alone in supporting it. - - IEEE NaN propagation comes for free: C ``a == b`` is false when either - side is NaN, so two records both carrying NaN compare unequal under - ``eq`` (and equal under ``ne``), matching the cfunc path's leaf-wise - semantic comparison. - """ - np_type = dtype.np_type - pinned_name = dtype.jit_c_name - # ``_udt_c_typedef`` returns None when the numpy layout doesn't match a - # C compiler's layout for the same struct. - typedef_info = _udt_c_typedef(pinned_name or dtype.name, np_type) - if typedef_info is None: - return None - type_name, _typedef = typedef_info - c_name = f"{op_name}_{type_name}" - op = "==" if is_eq else "!=" - join = " && " if is_eq else " || " - if np_type.subdtype is not None: - # Array UDT: unroll all elements. - _base, shape = np_type.subdtype - size = reduce(mul, shape) - terms = [f"((x->v[{i}]) {op} (y->v[{i}]))" for i in range(size)] - elif np_type.names is not None: - terms = [f"((x->{c}) {op} (y->{c}))" for _py, c, _d in _iter_record_leaves(np_type)] - 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} ; }}" - - -def _maybe_warn_jit_skipped(jit_info, op_name, dtype_name): - """Emit a ``NoJITWarning`` when JIT setup was skipped despite SS supporting JIT. - - Called by op-class ``_compile_udt`` paths after ``set_jit_c_*_on_op`` - returns ``None``. The ``_has_jit_set`` gate keeps the warning silent on - SS < 8 where JIT is fundamentally absent. - """ - if jit_info is None and _has_jit_set: - from ..ss.jit_config import _maybe_warn_no_jit - - _maybe_warn_no_jit(op_name=op_name, dtype_name=dtype_name) - - -def _set_jit_c_strings(gb_obj, c_name, c_defn, set_string_func): - """Pin ``GxB_JIT_C_NAME`` + ``GxB_JIT_C_DEFINITION`` on a fresh op handle. - - Both strings are one-shot on SuiteSparse: a second set returns - ``GrB_ALREADY_SET`` silently. The return is not checked here because - callers always pass a freshly-allocated handle. - - Arming an op with C source is the point at which this process has - committed to wanting a JIT compiler, so it is also where SuiteSparse's - non-compiling default gets raised. - """ - from ..ss.jit_config import _enable_jit_for_udt - - _enable_jit_for_udt() - set_string_func(gb_obj, ffi.new("char[]", c_name.encode()), lib.GxB_JIT_C_NAME) - set_string_func(gb_obj, ffi.new("char[]", c_defn.encode()), lib.GxB_JIT_C_DEFINITION) - - -def set_jit_c_comparison_on_op(gb_obj, op_name, dtype, set_string_func, *, is_eq): - """Generate and set JIT C name+definition for ``eq``/``ne`` on a UDT. - - 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``). - """ - 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 - _set_jit_c_strings(gb_obj, *result, set_string_func) - return result - - -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``). - """ - if not _has_jit_set: - return None - result = _make_jit_c_definition(op_name, py_op, dtype, arity) - if result is None: - return None - _set_jit_c_strings(gb_obj, *result, set_string_func) - return result diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index fc9c92661..c0bdf0434 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,10 +120,11 @@ 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, compile_udt_unary_wrapper, set_jit_c_on_op, ) @@ -255,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: @@ -494,3 +498,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``. + + ``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) + 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..37da5cc5f 100644 --- a/graphblas/core/operator/utils.py +++ b/graphblas/core/operator/utils.py @@ -68,6 +68,23 @@ 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): + # A Monoid takes one type, but used element-wise it is its + # 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 try: return op[ @@ -219,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) @@ -274,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 ce191f374..bb7339ce8 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -1,5 +1,4 @@ import itertools -from functools import lru_cache import numpy as np @@ -8,9 +7,28 @@ 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 .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, + _get_udt_info, + _has_half_leaf, + _is_lifted_op, + _literal_shape_differs, + _scalar_ewise_add_converted, + _udt_eq_without_numba, +) from .utils import _Pointer, output_type, wrapdoc if _supports_udfs: @@ -133,15 +151,19 @@ 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. + Without Numba, which ``binary.eq`` needs for UDTs, a UDT's values are + compared in numpy, which broadcasts them the same way. """ - 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. @@ -198,6 +220,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 @@ -206,8 +233,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( @@ -221,6 +255,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: @@ -353,12 +389,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 @@ -369,19 +400,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]) - arr[:] = _view_if_same_layout(val, arr.dtype) - 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 @@ -390,7 +409,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]) @@ -468,7 +487,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 (``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: # 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 +691,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 + other, dtype = _literal_operand(self.dtype, other, op) try: other = Scalar.from_value(other, dtype, is_cscalar=False, name="") except TypeError: @@ -682,6 +705,8 @@ 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): + self, other, op = _scalar_ewise_add_converted(self, other, op) return ScalarExpression( method_name, f"GrB_Vector_eWiseAdd_{op.opclass}", @@ -739,7 +764,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 + other, dtype = _literal_operand(self.dtype, other, op) try: other = Scalar.from_value(other, dtype, is_cscalar=False, name="") except TypeError: @@ -816,8 +841,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, dtype, is_cscalar=False, name="") + other = Scalar.from_value(other, literal_dtype, is_cscalar=False, name="") except TypeError: other = self._expect_type( other, @@ -828,46 +854,17 @@ 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") - left_dtype = temp_op.type - dtype = left_dtype if left_dtype._is_udt else None - if type(left_default) is not Scalar: - 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: - 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: - 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: - 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") @@ -1140,6 +1137,47 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=False, name=None, **opts): # End auto-generated code: Scalar +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: + 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 _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() + 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 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, 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) + + 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="") @@ -1148,76 +1186,41 @@ def _as_scalar(scalar, dtype=None, *, is_cscalar): return scalar -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) - - -@lru_cache -def _padding_ranges(np_type): - """The ``(start, stop)`` byte ranges of ``np_type`` that are alignment padding. +def _union_default(obj, default, input_dtype, op, keyword_name): + """Return ``default`` for ``obj.ewise_union`` as a Scalar, where the op takes ``input_dtype``. - Walks records, nested records, and array dtypes, and returns the gaps - between the field data. + GraphBLAS casts a default to the op's input type, and a UDT casts only to + itself, so where that is a UDT the default becomes an element of it and must + fit (``literals._check_literal_fits``). Otherwise a literal default is typed + as any literal operand is (``literals._literal_operand``), and a Scalar keeps + its type. ``keyword_name`` names the argument when ``default`` is not a scalar. """ - 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) + 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, left=keyword_name == "left_default" + ) + 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/_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/dtypes.py b/graphblas/core/ss/dtypes.py index 12f74fbce..0ccdd8264 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_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. + 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 @@ -84,7 +88,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 @@ -96,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/core/ss/indexbinary.py b/graphblas/core/ss/indexbinary.py index 85b3253e0..5aeed9f11 100644 --- a/graphblas/core/ss/indexbinary.py +++ b/graphblas/core/ss/indexbinary.py @@ -6,27 +6,29 @@ from ..operator.indexbinary import ( IndexBinaryOp, TypedBuiltinIndexBinaryOp, - TypedUserIndexBinaryOp, _has_idxbinop, + _TypedIndexBinaryOpMixin, ) +from ._utils import _jit_c_source 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) @@ -47,7 +49,12 @@ 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``). 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 ---------- @@ -101,18 +108,22 @@ 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) from .jit_config import _compiling_c_source + # 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( @@ -123,15 +134,16 @@ 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], ) 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 ) + 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/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/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/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/core/udt.py b/graphblas/core/udt.py new file mode 100644 index 000000000..94e7b3136 --- /dev/null +++ b/graphblas/core/udt.py @@ -0,0 +1,911 @@ +"""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 _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 +# 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`` (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, 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. + """ + 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 "SUmM" + ) + if ( + dtype != dtype2 + 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( + 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 + 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)) + + +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. + + +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`), 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 + + +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_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. + """ + 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): + """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) + _check_leaves_compare("eq", scalar.dtype, other.dtype) + 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 = _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 + # 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_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.""" + 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/utils.py b/graphblas/core/utils.py index 72f7219bd..9f15254ab 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 _nats_in_leaf_units, _view_if_same_layout _NP2 = np.__version__.startswith("2.") @@ -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: # 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 + frame = frame.f_back + stacklevel += 1 + warnings.warn(message, category, stacklevel=stacklevel) + + def output_type(val): try: return _output_types[type(val)] @@ -97,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), + _nats_in_leaf_units(_view_if_same_layout(array, np_type), np_type), np_type, copy=copy or _NP2 and None, order=order, @@ -132,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 5d8c8fef9..f4b72fca4 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, @@ -18,6 +19,7 @@ get_typed_op, op_from_string, ) +from .operator.select import _value_select_in_own_type from .scalar import ( _COMPLETE, _MATERIALIZE, @@ -26,6 +28,13 @@ ScalarIndexExpr, _as_scalar, _scalar_index, + _union_default, +) +from .udt import ( + _cast_for_store, + _element_shapes_differ, + _ewise_add_cast_error, + _ewise_add_needs_cast, ) from .utils import ( _CArray, @@ -351,6 +360,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 @@ -359,6 +373,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: @@ -593,7 +609,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 @@ -927,14 +947,15 @@ 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: # 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, @@ -1026,15 +1047,23 @@ 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 + 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_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( - "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})." + f"Vector sizes must match in {method_name}: {self._size} and {other._size}" ) + raise _ewise_add_cast_error(op, self.dtype, other.dtype) + + if other.ndim == 2: return MatrixExpression( method_name, None, @@ -1204,42 +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: - 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: - 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: - 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: - 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") @@ -1420,6 +1415,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 +1441,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 + left, dtype = _literal_operand(self.dtype, left, op, left=True) try: left = Scalar.from_value(left, dtype, is_cscalar=None, name="") except TypeError: @@ -1477,7 +1477,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 + right, dtype = _literal_operand(self.dtype, right, op) try: right = Scalar.from_value(right, dtype, is_cscalar=None, name="") except TypeError: @@ -1593,8 +1593,10 @@ 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: - 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. + thunk, dtype = _literal_operand(self.dtype, thunk, op) try: thunk = Scalar.from_value(thunk, dtype, is_cscalar=None, name="") except TypeError: @@ -1607,6 +1609,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: @@ -1843,6 +1846,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) @@ -1874,7 +1886,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: @@ -1885,6 +1897,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]) @@ -1901,6 +1917,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 @@ -1935,7 +1953,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_dtype.py b/graphblas/tests/test_dtype.py index a9550ff51..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 @@ -10,7 +11,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 @@ -368,6 +369,35 @@ class _Empty: dtypes.register_anonymous(_Empty) +def test_register_anonymous_zero_size(): + for dtype in [str, bytes, np.dtype("U"), np.dtype([]), np.dtype((np.int64, (0,)))]: + with pytest.raises(ValueError, match="size of 0"): + dtypes.register_anonymous(dtype) + with pytest.raises(ValueError, match="Unknown dtype"): + 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 @@ -440,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") @@ -539,6 +573,46 @@ 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 + # 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: + 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: + gb.ss.config["jit_c_control"] = jit_c_control + + def test_dtype_to_from_string(): types = [dtypes.BOOL, dtypes.FP64] for c in string.ascii_letters: diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py index 19b65d720..333c35fcc 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"), @@ -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,119 @@ 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_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. + + 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. @@ -450,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 @@ -479,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.""" @@ -508,6 +708,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 +916,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_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_matrix.py b/graphblas/tests/test_matrix.py index ea8a29f26..d1405c49f 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) @@ -4340,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_op.py b/graphblas/tests/test_op.py index 36152dfe0..d51993c2e 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 DomainMismatch, UdfParseError +from graphblas.exceptions import DimensionMismatch, DomainMismatch, InvalidValue, UdfParseError from .conftest import shouldhave @@ -1263,9 +1264,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 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): return x # pragma: no cover (numba) @@ -1274,7 +1278,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 @@ -1351,10 +1359,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) @@ -1426,6 +1435,470 @@ 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( # pragma: no cover (only on failure) + "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") +@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( + "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): @@ -1579,6 +2052,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 @@ -1596,6 +2080,24 @@ 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] + # 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] + @pytest.fixture(scope="module") def record_udt(): @@ -1636,7 +2138,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 +2147,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") @@ -1760,7 +2271,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") @@ -1776,13 +2287,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))) @@ -1798,7 +2310,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") @@ -1824,20 +2336,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 "v[0]" in defn - assert "v[4]" in defn + assert "plus_Vec37Jit" in defn + assert "z->v[i] = (x->v[i]) + (y->v[i])" 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") @@ -1953,7 +2465,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") @@ -2332,7 +2844,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) @@ -2381,16 +2893,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( @@ -2576,7 +3088,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") @@ -2585,16 +3097,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") @@ -3073,6 +3585,77 @@ def test_udt_eq_ne_scalar_broadcast_nan_propagates(): assert eq_nan.isequal(Vector.from_coo([0, 1, 2], [False, True, False])) +@pytest.mark.skipif("not supports_udfs") +def test_udt_eq_ne_type_a_literal_as_builtin_comparisons_do(): + """``eq`` and ``ne`` beside a UDT type a Python number as a built-in comparison does. + + 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 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") + v = Vector(f32, size=1) + v[0] = np.full(7, 0.1) + assert (v == 0.1).new()[0].new().value + assert not (v != 0.1).new()[0].new().value + assert (v == (0.1,) * 7).new()[0].new().value + assert not (v == np.float64(0.1)).new()[0].new().value + builtin = Vector.from_coo([0], [0.1], dtype=dtypes.FP32) + assert (builtin == 0.1).new()[0].new().value + w = Vector(i8, size=1) + w[0] = np.arange(7) + for number in [300, -(2**63) - 1, 2**63, 2**64]: + assert not (w == number).new()[0].new().value + 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] + 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__] + with np.errstate(over="ignore"): + v[0] = [1e300, 1] # no warning, as numpy + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_eq_ne_scalar_broadcast_array_1d(): @@ -3167,6 +3750,2612 @@ def test_udt_eq_ne_rejects_incompatible_pairs(): binary.eq(v_uv & v_arr).new() +@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. +# 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") +@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. + """ + 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)], + ), + "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, 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", b_dtype)]), name) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.parametrize("op", ["plus", "eq"]) +@pytest.mark.parametrize( + ("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, 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. + """ + 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") +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. + + 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. + """ + 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( + ("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_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 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 + 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. + """ + try: + T = dtypes.register_anonymous(np.dtype((np.float64, (64, 64))), "_LoopBig") + except InvalidValue: + if sys.platform != "win32": + 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") + 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 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") +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 + # truediv gives float elements whatever the int, as numpy divides in float64. + result = (w / 300).new() + assert result.dtype.np_type.base == np.float64 + np.testing.assert_array_equal(result[0].new().value, (base + 1) / 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 its type fits: ``0.5`` became an element of + zeros. + """ + udt = dtypes.register_anonymous(np.dtype((np.int64, (2, 3))), "_NamedLikeLifted") + + def scaled_plus(x, y): # pragma: no cover (numba) + return x * 10 + y + + def difference(x, y): # pragma: no cover (numba) + 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, 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 + # 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])) + + +@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_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() + 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]) + 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] + 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), + 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)]: + 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.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 + # ...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 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.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 be of that type. + + 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``. 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 datetime + import re + import warnings + + 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" + ) + 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") + 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) + 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) + 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)])[()] + 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) + # 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 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 + 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 [ + (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: 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: 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"), + ( + 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"), + # 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: 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"), + ]: + 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. 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: 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() + with pytest.warns(DeprecationWarning, match="does not fit _ExactText"): + assert t.apply(second, right=text_literal).new()[0].new().value.tolist() == (b"cd", 2) + # 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"]) + # 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 + 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"): + v.apply(udf, right=(1, 2, 300)) + # Defaults that fit still work, and fill in for the missing side. + 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)]: + 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 + + +# _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.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)))), + ]: + 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)) + # 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 + # 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() + # 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)) + # 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 + 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)]) + ) + 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])) + 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 + + +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_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. + + 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_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. + + 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(): + """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), 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 = [ + ("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 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] + # 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] + # 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) + # 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. + 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] + # 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) + 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 + # 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 + 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() + + +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. + + 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] + 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) + 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(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) + 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 _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 combine, either way round. + + The result is the left operand's type, since it holds every leaf. Converting + 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") + 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)]: + 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 = 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") +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)]: + # 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.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, "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,))) + + +@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, 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. 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 + + 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") + 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")]: + 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] + # 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. 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"), + (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)), "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() + 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 + + +@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, 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) + # 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) + 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() + 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]] + s = gb.Scalar.from_value([1, 2, 3], dtype=i8s) + 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_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. + + 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_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)) + 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 + 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. + + ``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. + + 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_codegen 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_on_leaves_of_other_kinds_is_a_keyerror(): + """``eq`` refuses leaves that hold different kinds of value with a KeyError, as ``plus`` does. + + 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( + 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(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") +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, 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) + 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" + # 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) + 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") + 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") + 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: 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, + ]: + 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. + + 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" + # 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") +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_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] + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_aggregators(): @@ -3536,7 +6725,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. @@ -3549,7 +6737,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_scalar.py b/graphblas/tests/test_scalar.py index b8c00ffec..686e4861a 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -279,6 +279,56 @@ 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)])) + 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) + 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() + 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"]) + # 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() @@ -367,6 +417,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( @@ -628,6 +725,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): @@ -660,6 +760,69 @@ 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) + text = dtypes.register_anonymous(np.dtype("S6")) + t = Scalar.from_value(b"ab", text) + 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)), + (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( + 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), + (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)), + ] + + 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_ssjit.py b/graphblas/tests/test_ssjit.py index eeeeac206..9e083a2a1 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -25,9 +25,10 @@ ) 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 +from graphblas.exceptions import EmptyObject from .conftest import autocompute, burble @@ -591,6 +592,26 @@ 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") + # 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") 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. @@ -786,10 +807,15 @@ 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 + # 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) @@ -799,6 +825,17 @@ 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)) + # 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(): @@ -814,7 +851,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 @@ -825,14 +862,94 @@ 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" ) 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, 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.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) + 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 +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", "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 @@ -1105,41 +1221,181 @@ 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 + 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") + _require_jit_on() + + 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. + + 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. 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") + _require_jit_on() + + 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") @@ -1475,6 +1731,44 @@ 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)])], + ) + 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(): diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 52c382742..7d5cab524 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -3,6 +3,7 @@ import pickle import sys import types +import warnings import weakref import numpy as np @@ -1493,6 +1494,292 @@ def test_ss_nbytes(v): assert v.ss.nbytes > 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 + else: + # 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 + A = Matrix.from_coo([0], [0], [100], dtype=dtypes.INT8) + assert (A * 2).dtype == dtypes.INT8 + # 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: + 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 + # 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 + 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 + 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] + # 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 + + +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'. + + 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. 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) + 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, 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]), + (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]), + (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, 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 + # 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) + 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 @@ -2097,11 +2384,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) @@ -2641,6 +2929,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) 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