diff --git a/README.md b/README.md index 258e764bb..1f7ee3bbb 100644 --- a/README.md +++ b/README.md @@ -180,9 +180,3 @@ m = gb.io.to_scipy_sparse_matrix(m, format='csr') A = gb.io.from_networkx(g) g = gb.io.to_networkx(A) ``` - -## Attribution -This library borrows some great ideas from [pygraphblas](https://github.com/michelp/pygraphblas), -especially around parsing operator names from SuiteSparse and the concept of a Scalar which the backend -implementation doesn't need to know about. - diff --git a/docs/index.rst b/docs/index.rst index f260ea6d8..d4bd71cc8 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -219,14 +219,6 @@ Import/Export connectors to the Python ecosystem A = gb.io.from_networkx(g) g = gb.io.to_networkx(A) -Attribution ------------ - -This library borrows some great ideas from `pygraphblas `_, -especially around parsing operator names from SuiteSparse and the concept of a Scalar which the backend -implementation doesn't need to know about. - - Indices and tables ================== diff --git a/grblas/_agg.py b/grblas/_agg.py index cfc87107a..d45052be4 100644 --- a/grblas/_agg.py +++ b/grblas/_agg.py @@ -4,8 +4,8 @@ import numpy as np from . import agg, binary, monoid, semiring, unary -from .dtypes import lookup_dtype, unify -from .operator import _normalize_type +from .dtypes import INT64, lookup_dtype +from .operator import get_typed_op from .scalar import Scalar @@ -14,8 +14,8 @@ def _get_types(ops, initdtype): if initdtype is None: prev = dict(ops[0].types) else: - initdtype = lookup_dtype(initdtype) - prev = {key: unify(lookup_dtype(val), initdtype).name for key, val in ops[0].types.items()} + op = ops[0] + prev = {key: get_typed_op(op, key, initdtype).return_type for key in op.types} for op in ops[1:]: cur = {} types = op.types @@ -43,11 +43,12 @@ def __init__( composite=None, custom=None, types=None, + any_dtype=None, ): self.name = name self._initval_orig = initval self._initval = False if initval is None else initval - self._initdtype = lookup_dtype(type(self._initval)) + self._initdtype = lookup_dtype(type(self._initval), self._initval) self._monoid = monoid self._semiring = semiring self._semiring2 = semiring2 @@ -68,6 +69,7 @@ def __init__( self._types_orig = types self._types = None self._typed_ops = {} + self._any_dtype = any_dtype @property def types(self): @@ -82,16 +84,16 @@ def types(self): return self._types def __getitem__(self, dtype): - dtype = _normalize_type(dtype) - if dtype not in self.types: + dtype = lookup_dtype(dtype) + if not self._any_dtype and dtype not in self.types: raise KeyError(f"{self.name} does not work with {dtype}") if dtype not in self._typed_ops: self._typed_ops[dtype] = TypedAggregator(self, dtype) return self._typed_ops[dtype] def __contains__(self, dtype): - dtype = _normalize_type(dtype) - return dtype in self.types + dtype = lookup_dtype(dtype) + return self._any_dtype or dtype in self.types def __repr__(self): return f"agg.{self.name}" @@ -107,7 +109,12 @@ def __init__(self, agg, dtype): self.name = agg.name self.parent = agg self.type = dtype - self.return_type = agg.types[dtype] + if dtype in agg.types: + self.return_type = agg.types[dtype] + elif agg._any_dtype is True: + self.return_type = dtype + else: + self.return_type = agg._any_dtype def __repr__(self): return f"agg.{self.name}[{self.type}]" @@ -160,8 +167,7 @@ def _new(self, updater, expr, *, in_composite=False): if agg._custom is not None: return agg._custom(self, updater, expr, in_composite=in_composite) - dtype = unify(lookup_dtype(self.type), lookup_dtype(agg._initdtype)) - semiring = agg._semiring[dtype] + semiring = get_typed_op(agg._semiring, self.type, agg._initdtype) if expr.cfunc_name == "GrB_Matrix_reduce_Aggregator": # Matrix -> Vector A = expr.args[0] @@ -242,13 +248,15 @@ def __reduce__(self): agg.any = Aggregator("any", monoid=monoid.lor) agg.min = Aggregator("min", monoid=monoid.min) agg.max = Aggregator("max", monoid=monoid.max) -agg.any_value = Aggregator("any_value", monoid=monoid.any) +agg.any_value = Aggregator("any_value", monoid=monoid.any, any_dtype=True) agg.bitwise_all = Aggregator("bitwise_all", monoid=monoid.band) agg.bitwise_any = Aggregator("bitwise_any", monoid=monoid.bor) # Other monoids: bxnor bxor eq lxnor lxor # Semiring-only -agg.count = Aggregator("count", semiring=semiring.plus_pair, semiring2=semiring.plus_first) +agg.count = Aggregator( + "count", semiring=semiring.plus_pair, semiring2=semiring.plus_first, any_dtype=INT64 +) agg.count_nonzero = Aggregator( "count_nonzero", semiring=semiring.plus_isne, semiring2=semiring.plus_first ) @@ -264,7 +272,9 @@ def __reduce__(self): semiring=semiring.plus_pow, semiring2=semiring.plus_first, ) -agg.exists = Aggregator("exists", semiring=semiring.any_pair, semiring2=semiring.any_pair) +agg.exists = Aggregator( + "exists", semiring=semiring.any_pair, semiring2=semiring.any_pair, any_dtype=INT64 +) # Semiring and finalize agg.hypot = Aggregator( @@ -564,11 +574,13 @@ def _first_last(agg, updater, expr, *, in_composite, semiring_): "first", custom=partial(_first_last, semiring_=semiring.min_secondi), types=[binary.first], + any_dtype=True, ) agg.last = Aggregator( "last", custom=partial(_first_last, semiring_=semiring.max_secondi), types=[binary.second], + any_dtype=True, ) @@ -601,9 +613,11 @@ def _first_last_index(agg, updater, expr, *, in_composite, semiring): "first_index", custom=partial(_first_last_index, semiring=semiring.min_secondi), types=[semiring.min_secondi], + any_dtype=INT64, ) agg.last_index = Aggregator( "last_index", custom=partial(_first_last_index, semiring=semiring.max_secondi), types=[semiring.min_secondi], + any_dtype=INT64, ) diff --git a/grblas/_ss/matrix.py b/grblas/_ss/matrix.py index 7c958ee2e..ad5db5bce 100644 --- a/grblas/_ss/matrix.py +++ b/grblas/_ss/matrix.py @@ -786,7 +786,7 @@ def _export(self, format=None, *, sort=False, give_ownership=False, raw=False, m parent = self._parent else: parent = self._parent.dup(name=f"M_{method}") - dtype = np.dtype(parent.dtype.np_type) + dtype = parent.dtype.np_type index_dtype = np.dtype(np.uint64) nrows = parent._nrows @@ -847,25 +847,27 @@ def _export(self, format=None, *, sort=False, give_ownership=False, raw=False, m parent.gb_obj = ffi.NULL else: parent.clear() - return rv elif format == "coor": - info = self._export( + rv = self._export( "csr", sort=sort, give_ownership=give_ownership, raw=False, method=method ) - info["rows"] = indptr_to_indices(info.pop("indptr")) - info["cols"] = info.pop("col_indices") - info["sorted_rows"] = True - info["format"] = "coor" - return info + rv["rows"] = indptr_to_indices(rv.pop("indptr")) + rv["cols"] = rv.pop("col_indices") + rv["sorted_rows"] = True + rv["format"] = "coor" elif format == "cooc": - info = self._export( + rv = self._export( "csc", sort=sort, give_ownership=give_ownership, raw=False, method=method ) - info["cols"] = indptr_to_indices(info.pop("indptr")) - info["rows"] = info.pop("row_indices") - info["sorted_cols"] = True - info["format"] = "cooc" - return info + rv["cols"] = indptr_to_indices(rv.pop("indptr")) + rv["rows"] = rv.pop("row_indices") + rv["sorted_cols"] = True + rv["format"] = "cooc" + else: + raise ValueError(f"Invalid format: {format}") + if parent.dtype._is_udt: + rv["dtype"] = parent.dtype + return rv if method == "export": mhandle = ffi_new("GrB_Matrix*", parent._carg) @@ -1175,6 +1177,8 @@ def _export(self, format=None, *, sort=False, give_ownership=False, raw=False, m rv["values"] = values if method == "export": parent.gb_obj = ffi.NULL + if parent.dtype._is_udt: + rv["dtype"] = parent.dtype return rv @classmethod diff --git a/grblas/_ss/vector.py b/grblas/_ss/vector.py index feb7a8c1a..b78d6c9ec 100644 --- a/grblas/_ss/vector.py +++ b/grblas/_ss/vector.py @@ -362,7 +362,7 @@ def _export(self, format=None, *, sort=False, give_ownership=False, raw=False, m parent = self._parent else: parent = self._parent.dup(name=f"v_{method}") - dtype = np.dtype(parent.dtype.np_type) + dtype = parent.dtype.np_type index_dtype = np.dtype(np.uint64) if format is None: @@ -481,6 +481,8 @@ def _export(self, format=None, *, sort=False, give_ownership=False, raw=False, m ) if method == "export": parent.gb_obj = ffi.NULL + if parent.dtype._is_udt: + rv["dtype"] = parent.dtype return rv @classmethod diff --git a/grblas/base.py b/grblas/base.py index 311419f8d..bae7e32fb 100644 --- a/grblas/base.py +++ b/grblas/base.py @@ -318,7 +318,6 @@ def update(self, expr): return self._update(expr) def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None): - # TODO: check expected output type (now included in Expression object) if not isinstance(expr, BaseExpression): if isinstance(expr, AmbiguousAssignOrExtract): if expr._is_scalar and self._is_scalar: @@ -557,7 +556,7 @@ def __init__( raise ValueError(f"No default expr_repr for len(args) == {len(args)}") self.expr_repr = expr_repr if dtype is None: - self.dtype = lookup_dtype(op.return_type) + self.dtype = op.return_type else: self.dtype = dtype self._value = None diff --git a/grblas/binary/numpy.py b/grblas/binary/numpy.py index e832583e2..1af2a5574 100644 --- a/grblas/binary/numpy.py +++ b/grblas/binary/numpy.py @@ -89,6 +89,7 @@ "subtract": "minus", "true_divide": "truediv", } +# _graphblas_to_numpy = {val: key for key, val in _numpy_to_graphblas.items()} # Soon... # Not included: maximum, minimum, gcd, hypot, logaddexp, logaddexp2 # lcm, left_shift, nextafter, right_shift diff --git a/grblas/dtypes.py b/grblas/dtypes.py index cced776da..6a074a2a1 100644 --- a/grblas/dtypes.py +++ b/grblas/dtypes.py @@ -2,7 +2,7 @@ import numpy as np from numpy import find_common_type, promote_types -from . import lib +from . import ffi, lib # Default assumption unless FC32/FC64 are found in lib _supports_complex = hasattr(lib, "GrB_FC64") or hasattr(lib, "GxB_FC64") @@ -17,30 +17,100 @@ def __init__(self, name, gb_obj, gb_name, c_type, numba_type, np_type): self.gb_name = gb_name self.c_type = c_type self.numba_type = numba_type - self.np_type = np_type + self.np_type = np.dtype(np_type) def __repr__(self): return self.name def __eq__(self, other): if type(other) is DataType: - return self.gb_obj == other.gb_obj - else: - # Attempt to use `other` as a lookup key - try: - other = lookup_dtype(other) - return self == other - except ValueError: - raise TypeError(f"Invalid or unknown datatype: {other}") from None + return self is other + # Attempt to use `other` as a lookup key + try: + return self is lookup_dtype(other) + except ValueError: + raise TypeError(f"Invalid or unknown datatype: {other}") from None + + def __hash__(self): + return hash(self.gb_obj) + + def __lt__(self, other): + # Let us sort for prettier error reporting + if type(other) is DataType: + t1 = self.np_type + t2 = other.np_type + return (t1.kind, t1.itemsize, t1.name) < (t2.kind, t2.itemsize, t2.name) + # Attempt to use `other` as a lookup key + try: + return self < lookup_dtype(other) + except ValueError: + raise TypeError(f"Invalid or unknown datatype: {other}") from None def __reduce__(self): + if self._is_udt: + return (self._deserialize, (self.name, self.np_type, self._is_anonymous)) if self.gb_name == "GrB_Index": return "_INDEX" return self.name @property def _carg(self): - return self.gb_obj + return self.gb_obj[0] if self.gb_name is None else self.gb_obj + + @property + def _is_anonymous(self): + return globals().get(self.name) is not self + + @property + def _is_udt(self): + return self.gb_name is None + + @staticmethod + def _deserialize(name, dtype, is_anonymous): + if is_anonymous: + return register_anonymous(dtype, name) + if name in _registry: + return _registry[name] + return register_new(name, dtype) + + +def register_new(name, dtype): + if not name.isidentifier(): + raise ValueError(f"`name` argument must be a valid Python identifier; got: {name!r}") + if name in _registry or name in globals(): + raise ValueError(f"{name!r} name for dtype is unavailable") + rv = register_anonymous(dtype, name) + _registry[name] = rv + globals()[name] = rv + return rv + + +def register_anonymous(dtype, name=None): + from .exceptions import check_status_carg + + dtype = np.dtype(dtype) + if dtype.hasobject: + raise ValueError("dtype must not allow Python objects") + if dtype.isbuiltin != 0: + raise ValueError("dtype must not be a builtin type") + if name is None: + name = repr(dtype) + numba_type = numba.typeof(dtype).dtype + + if dtype in _registry: + gb_obj = _registry[dtype].gb_obj + else: + gb_obj = ffi.new("GrB_Type*") + status = lib.GrB_Type_new(gb_obj, dtype.itemsize) + check_status_carg(status, "Type", gb_obj[0]) + # For now, let's use "opaque" unsigned bytes for the c type. + rv = DataType(name, gb_obj, None, f"uint8_t[{dtype.itemsize}]", numba_type, dtype) + if dtype not in _registry: + _registry[gb_obj] = rv + _registry[dtype] = rv + _registry[numba_type] = rv + _registry[numba_type.name] = rv + return rv BOOL = DataType("BOOL", lib.GrB_BOOL, "GrB_BOOL", "_Bool", numba.types.bool_, np.bool_) @@ -76,23 +146,23 @@ def _carg(self): # Used for testing user-defined functions _sample_values = { - INT8.name: np.int8(1), - UINT8.name: np.uint8(1), - INT16.name: np.int16(1), - UINT16.name: np.uint16(1), - INT32.name: np.int32(1), - UINT32.name: np.uint32(1), - INT64.name: np.int64(1), - UINT64.name: np.uint64(1), - FP32.name: np.float32(0.5), - FP64.name: np.float64(0.5), - BOOL.name: np.bool_(True), + INT8: np.int8(1), + UINT8: np.uint8(1), + INT16: np.int16(1), + UINT16: np.uint16(1), + INT32: np.int32(1), + UINT32: np.uint32(1), + INT64: np.int64(1), + UINT64: np.uint64(1), + FP32: np.float32(0.5), + FP64: np.float64(0.5), + BOOL: np.bool_(True), } if _supports_complex: _sample_values.update( { - FC32.name: np.complex64(complex(0, 0.5)), - FC64.name: np.complex128(complex(0, 0.5)), + FC32: np.complex64(complex(0, 0.5)), + FC64: np.complex128(complex(0, 0.5)), } ) @@ -124,7 +194,7 @@ def _carg(self): _registry[dtype.c_type.upper()] = dtype _registry[dtype.numba_type] = dtype _registry[dtype.numba_type.name] = dtype - val = _sample_values[dtype.name] + val = _sample_values[dtype] _registry[val.dtype] = dtype _registry[val.dtype.name] = dtype # Upcast numpy float16 to float32 @@ -143,7 +213,7 @@ def _carg(self): _registry["complex"] = FC64 -def lookup_dtype(key): +def lookup_dtype(key, value=None): # Check for silly lookup where key is already a DataType if type(key) is DataType: return key @@ -156,11 +226,17 @@ def lookup_dtype(key): return _registry[key.name] except KeyError: pass + if value is not None and hasattr(value, "dtype") and value.dtype in _registry: + return _registry[value.dtype] try: return lookup_dtype(np.dtype(key)) except Exception: pass - raise ValueError(f"Unknown dtype: {key}") + try: + return lookup_dtype(key.literal_type) # For numba dtype inference + except Exception: + pass + raise ValueError(f"Unknown dtype: {key}", type(key)) def unify(type1, type2, *, is_left_scalar=False, is_right_scalar=False): diff --git a/grblas/formatting.py b/grblas/formatting.py index d4fbab700..331e3218f 100644 --- a/grblas/formatting.py +++ b/grblas/formatting.py @@ -135,9 +135,13 @@ def _update_matrix_dataframe(df, matrix, rows, row_offset, columns, column_offse if row_offset is not None or column_offset is not None: submatrix = submatrix[row_offset:, column_offset:].new(name="") rows, cols, vals = submatrix.to_values() + np_type = submatrix.dtype.np_type + if submatrix.dtype._is_udt and np_type.subdtype is not None: + vals = vals.tolist() df.values[rows, cols] = vals - nulls = np.isnan(vals) - df.values[rows[nulls], cols[nulls]] = "nan" + if np.issubdtype(np_type, np.inexact): + nulls = np.isnan(vals) + df.values[rows[nulls], cols[nulls]] = "nan" def _update_vector_dataframe(df, vector, columns, column_offset, *, mask=None): @@ -169,8 +173,12 @@ def _update_vector_dataframe(df, vector, columns, column_offset, *, mask=None): if column_offset is not None: subvector = subvector[column_offset:].new(name="") cols, vals = subvector.to_values() + np_type = subvector.dtype.np_type + if subvector.dtype._is_udt and np_type.subdtype is not None: + vals = vals.tolist() df.values[0, cols] = vals - df.values[0, cols[np.isnan(vals)]] = "nan" + if np.issubdtype(np_type, np.inexact): + df.values[0, cols[np.isnan(vals)]] = "nan" def _get_max_columns(): diff --git a/grblas/matrix.py b/grblas/matrix.py index 61cb37861..ed6f180bb 100644 --- a/grblas/matrix.py +++ b/grblas/matrix.py @@ -51,7 +51,7 @@ def __del__(self): if parent is not None: return gb_obj = getattr(self, "gb_obj", None) - if gb_obj is not None: + if gb_obj is not None and lib is not None: # it's difficult/dangerous to record the call, b/c `self.name` may not exist check_status(lib.GrB_Matrix_free(gb_obj), self) @@ -154,12 +154,12 @@ def isequal(self, other, *, check_dtype=False): if self._nvals != other._nvals: return False if check_dtype: - common_dtype = self.dtype + op = binary.eq[self.dtype] else: - common_dtype = unify(self.dtype, other.dtype) + op = get_typed_op(binary.eq, self.dtype, other.dtype, kind="binary") matches = Matrix.new(bool, self._nrows, self._ncols, name="M_isequal") - matches << self.ewise_mult(other, binary.eq[common_dtype]) + matches << self.ewise_mult(other, op) # ewise_mult performs intersection, so nvals will indicate mismatched empty values if matches._nvals != self._nvals: return False @@ -253,8 +253,9 @@ def to_values(self, dtype=None): n = ffi_new("GrB_Index*") scalar = Scalar(n, _INDEX, name="s_nvals", is_cscalar=True, empty=True) scalar.value = nvals + dtype_name = "UDT" if self.dtype._is_udt else self.dtype.name call( - f"GrB_Matrix_extractTuples_{self.dtype.name}", + f"GrB_Matrix_extractTuples_{dtype_name}", [rows, columns, values, _Pointer(scalar), self], ) values = values.array @@ -273,7 +274,7 @@ def build(self, rows, columns, values, *, dup_op=None, clear=False, nrows=None, rows = ints_to_numpy_buffer(rows, np.uint64, name="row indices") columns = ints_to_numpy_buffer(columns, np.uint64, name="column indices") values, dtype = values_to_numpy_buffer(values, self.dtype) - n = values.size + n = values.shape[0] if rows.size != n or columns.size != n: raise ValueError( f"`rows` and `columns` and `values` lengths must match: " @@ -292,7 +293,11 @@ def build(self, rows, columns, values, *, dup_op=None, clear=False, nrows=None, dup_op_given = dup_op is not None if not dup_op_given: - dup_op = binary.plus + if not self.dtype._is_udt: + dup_op = binary.plus + else: + dup_op = binary.any + # SS:SuiteSparse-specific: we could use NULL for dup_op dup_op = get_typed_op(dup_op, self.dtype, kind="binary") if dup_op.opclass == "Monoid": dup_op = dup_op.binaryop @@ -302,8 +307,9 @@ def build(self, rows, columns, values, *, dup_op=None, clear=False, nrows=None, rows = _CArray(rows) columns = _CArray(columns) values = _CArray(values, self.dtype) + dtype_name = "UDT" if self.dtype._is_udt else self.dtype.name call( - f"GrB_Matrix_build_{self.dtype.name}", + f"GrB_Matrix_build_{dtype_name}", [self, rows, columns, values, _as_scalar(n, _INDEX, is_cscalar=True), dup_op], ) # Check for duplicates when dup_op was not provided @@ -381,7 +387,7 @@ def from_values( """ rows = ints_to_numpy_buffer(rows, np.uint64, name="row indices") columns = ints_to_numpy_buffer(columns, np.uint64, name="column indices") - values, dtype = values_to_numpy_buffer(values, dtype) + values, new_dtype = values_to_numpy_buffer(values, dtype) # Compute nrows and ncols if not provided if nrows is None: if rows.size == 0: @@ -391,8 +397,11 @@ def from_values( if columns.size == 0: raise ValueError("No column indices provided. Unable to infer ncols.") ncols = int(columns.max()) + 1 + if dtype is None and values.ndim > 1: + # Look for array-subtdype + new_dtype = lookup_dtype(np.dtype((new_dtype.np_type, values.shape[1:]))) # Create the new matrix - C = cls.new(dtype, nrows, ncols, name=name) + C = cls.new(new_dtype, nrows, ncols, name=name) if values.ndim == 0: if dup_op is not None: raise ValueError( @@ -530,10 +539,11 @@ def ewise_union(self, other, op, left_default, right_default): other = self._expect_type( other, (Matrix, TransposedMatrix, Vector), within=method_name, argname="other", op=op ) + dtype = self.dtype if self.dtype._is_udt else None if type(left_default) is not Scalar: try: left = Scalar.from_value( - left_default, is_cscalar=False, name="" # pragma: is_grbscalar + left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar ) except TypeError: left = self._expect_type( @@ -545,11 +555,11 @@ def ewise_union(self, other, op, left_default, right_default): op=op, ) else: - left = _as_scalar(left_default, is_cscalar=False) # pragma: is_grbscalar + left = _as_scalar(left_default, dtype, is_cscalar=False) # pragma: is_grbscalar if type(right_default) is not Scalar: try: right = Scalar.from_value( - right_default, is_cscalar=False, name="" # pragma: is_grbscalar + right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar ) except TypeError: right = self._expect_type( @@ -561,7 +571,7 @@ def ewise_union(self, other, op, left_default, right_default): op=op, ) else: - right = _as_scalar(right_default, is_cscalar=False) # pragma: is_grbscalar + right = _as_scalar(right_default, dtype, is_cscalar=False) # pragma: is_grbscalar scalar_dtype = unify(left.dtype, right.dtype) nonscalar_dtype = unify(self.dtype, other.dtype) op = get_typed_op(op, scalar_dtype, nonscalar_dtype, is_left_scalar=True, kind="binary") @@ -691,8 +701,9 @@ 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 try: - left = Scalar.from_value(left, is_cscalar=None, name="") + left = Scalar.from_value(left, dtype, is_cscalar=None, name="") except TypeError: left = self._expect_type( left, @@ -714,15 +725,21 @@ def apply(self, op, right=None, *, left=None): extra_message=extra_message, ) if left._is_cscalar: - cfunc_name = f"GrB_Matrix_apply_BinaryOp1st_{left.dtype}" + if left.dtype._is_udt: + dtype_name = "UDT" + left = _Pointer(left) + else: + dtype_name = left.dtype.name + cfunc_name = f"GrB_Matrix_apply_BinaryOp1st_{dtype_name}" else: cfunc_name = "GrB_Matrix_apply_BinaryOp1st_Scalar" args = [left, self] 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 else None try: - right = Scalar.from_value(right, is_cscalar=None, name="") + right = Scalar.from_value(right, dtype, is_cscalar=None, name="") except TypeError: right = self._expect_type( right, @@ -744,7 +761,12 @@ def apply(self, op, right=None, *, left=None): extra_message=extra_message, ) if right._is_cscalar: - cfunc_name = f"GrB_Matrix_apply_BinaryOp2nd_{right.dtype}" + if right.dtype._is_udt: + dtype_name = "UDT" + right = _Pointer(right) + else: + dtype_name = right.dtype.name + cfunc_name = f"GrB_Matrix_apply_BinaryOp2nd_{dtype_name}" else: cfunc_name = "GrB_Matrix_apply_BinaryOp2nd_Scalar" args = [self, right] @@ -830,6 +852,8 @@ def reduce_scalar(self, op=monoid.plus, *, allow_empty=True): raise ValueError("allow_empty=False not allowed when using Aggregators") if allow_empty: cfunc_name = "GrB_Matrix_reduce_Monoid_Scalar" + elif self.dtype._is_udt: + cfunc_name = "GrB_Matrix_reduce_UDT" else: cfunc_name = "GrB_Matrix_reduce_{output_dtype}" return ScalarExpression( @@ -854,9 +878,10 @@ def _extract_element(self, resolved_indexes, dtype=None, *, is_cscalar, name=Non if result is None: result = Scalar.new(dtype, is_cscalar=is_cscalar, name=name) if is_cscalar: + dtype_name = "UDT" if dtype._is_udt else dtype.name if ( call( - f"GrB_Matrix_extractElement_{dtype}", + f"GrB_Matrix_extractElement_{dtype_name}", [_Pointer(result), self, rowidx.index, colidx.index], ) is not NoValue @@ -907,8 +932,9 @@ 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 try: - value = Scalar.from_value(value, is_cscalar=None, name="") + value = Scalar.from_value(value, dtype, is_cscalar=None, name="") except TypeError: value = self._expect_type( value, @@ -921,7 +947,12 @@ def _assign_element(self, resolved_indexes, value): if value._empty: call("GrB_Matrix_removeElement", [self, rowidx.index, colidx.index]) return - cfunc_name = f"GrB_Matrix_setElement_{value.dtype}" + if value.dtype._is_udt: + dtype_name = "UDT" + value = _Pointer(value) + else: + dtype_name = value.dtype.name + cfunc_name = f"GrB_Matrix_setElement_{dtype_name}" else: cfunc_name = "GrB_Matrix_setElement_Scalar" call(cfunc_name, [self, value, rowidx.index, colidx.index]) @@ -1102,8 +1133,9 @@ def _prep_for_assign(self, resolved_indexes, value, mask=None, is_submask=False) ) else: if type(value) is not Scalar: + dtype = self.dtype if self.dtype._is_udt else None try: - value = Scalar.from_value(value, is_cscalar=None, name="") + value = Scalar.from_value(value, dtype, is_cscalar=None, name="") except TypeError: if rowsize is None or colsize is None: types = (Scalar, Vector) @@ -1197,7 +1229,12 @@ def _prep_for_assign(self, resolved_indexes, value, mask=None, is_submask=False) # C[I, J](M) << c # SS, SuiteSparse-specific: subassign if value._is_cscalar: - cfunc_name = f"GrB_Matrix_subassign_{value.dtype}" + if value.dtype._is_udt: + dtype_name = "UDT" + value = _Pointer(value) + else: + dtype_name = value.dtype.name + cfunc_name = f"GrB_Matrix_subassign_{dtype_name}" else: cfunc_name = "GrB_Matrix_subassign_Scalar" expr_repr = ( @@ -1216,7 +1253,12 @@ def _prep_for_assign(self, resolved_indexes, value, mask=None, is_submask=False) cols = _CArray([cols.value]) colscalar = _as_scalar(1, _INDEX, is_cscalar=True) if value._is_cscalar: - cfunc_name = f"GrB_Matrix_assign_{value.dtype}" + if value.dtype._is_udt: + dtype_name = "UDT" + value = _Pointer(value) + else: + dtype_name = value.dtype.name + cfunc_name = f"GrB_Matrix_assign_{dtype_name}" else: cfunc_name = "GrB_Matrix_assign_Scalar" expr_repr = "[[{2._expr_name} rows], [{4._expr_name} cols]] = {0._expr_name}" diff --git a/grblas/monoid/numpy.py b/grblas/monoid/numpy.py index f5949de00..b9deea329 100644 --- a/grblas/monoid/numpy.py +++ b/grblas/monoid/numpy.py @@ -113,6 +113,7 @@ "logical_xor": "lxor", "multiply": "times", } +# _graphblas_to_numpy = {val: key for key, val in _numpy_to_graphblas.items()} # Soon... # Not included: maximum, minimum, gcd, hypot, logaddexp, logaddexp2 diff --git a/grblas/operator.py b/grblas/operator.py index 499e3a5f4..ebebae66c 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -2,19 +2,38 @@ import itertools import re from collections.abc import Mapping -from functools import lru_cache -from operator import getitem +from functools import lru_cache, reduce +from operator import getitem, mul from types import FunctionType, ModuleType import numba import numpy as np from . import config, ffi, lib -from .dtypes import INT8, _sample_values, _supports_complex, lookup_dtype, unify +from .dtypes import ( + BOOL, + FP32, + FP64, + INT8, + INT16, + INT32, + INT64, + UINT8, + UINT16, + UINT32, + UINT64, + _sample_values, + _supports_complex, + lookup_dtype, + unify, +) from .exceptions import UdfParseError, check_status_carg from .expr import InfixExprBase from .utils import libget, output_type +if _supports_complex: + from .dtypes import FC32, FC64 + _STANDARD_OPERATOR_NAMES = set() from . import binary, monoid, op, semiring, unary # noqa isort:skip @@ -23,10 +42,6 @@ UNKNOWN_OPCLASS = "UnknownOpClass" -def _normalize_type(type_): - return lookup_dtype(type_).name - - def _hasop(module, name): return name in module.__dict__ or name in module._delayed @@ -80,22 +95,62 @@ def _call_op(op, left, right=None, **kwargs): ) +_udt_mask_cache = {} + + +def _udt_mask(dtype): + """Create mask to determine which bytes of UDTs to use for equality check""" + if dtype in _udt_mask_cache: + return _udt_mask_cache[dtype] + if dtype.subdtype is not None: + mask = _udt_mask(dtype.subdtype[0]) + N = reduce(mul, dtype.subdtype[1]) + rv = np.concatenate([mask] * N) + elif dtype.names is not None: + prev_offset = mask = None + masks = [] + for name in dtype.names: + dtype2, offset = dtype.fields[name] + if mask is not None: + masks.append(np.pad(mask, (0, offset - prev_offset - mask.size))) + mask = _udt_mask(dtype2) + prev_offset = offset + masks.append(np.pad(mask, (0, dtype.itemsize - prev_offset - mask.size))) + rv = np.concatenate(masks) + else: + rv = np.ones(dtype.itemsize, dtype=bool) + # assert rv.size == dtype.itemsize + _udt_mask_cache[dtype] = rv + return rv + + class TypedOpBase: - __slots__ = "parent", "name", "type", "return_type", "gb_obj", "gb_name", "__weakref__" + __slots__ = ( + "parent", + "name", + "type", + "return_type", + "gb_obj", + "gb_name", + "_type2", + "__weakref__", + ) - def __init__(self, parent, name, type_, return_type, gb_obj, gb_name): + def __init__(self, parent, name, type_, return_type, gb_obj, gb_name, dtype2=None): self.parent = parent self.name = name - self.type = _normalize_type(type_) - self.return_type = _normalize_type(return_type) + self.type = type_ + self.return_type = return_type self.gb_obj = gb_obj self.gb_name = gb_name + self._type2 = dtype2 def __repr__(self): classname = self.opclass.lower() if classname.endswith("op"): classname = classname[:-2] - return f"{classname}.{self.name}[{self.type}]" + dtype2 = "" if self._type2 is None else f", {self._type2.name}" + return f"{classname}.{self.name}[{self.type.name}{dtype2}]" @property def _carg(self): @@ -106,7 +161,10 @@ def is_positional(self): return self.parent.is_positional def __reduce__(self): - return (getitem, (self.parent, self.type)) + if self._type2 is None or self.type == self._type2: + return (getitem, (self.parent, self.type)) + else: + return (getitem, (self.parent, (self.type, self._type2))) class TypedBuiltinUnaryOp(TypedOpBase): @@ -171,19 +229,23 @@ def monoid(self): @property def commutes_to(self): commutes_to = self.parent.commutes_to - if commutes_to is not None and self.type in commutes_to._typed_ops: + if commutes_to is not None and (self.type in commutes_to._typed_ops or self.type._is_udt): return commutes_to[self.type] @property def _semiring_commutes_to(self): commutes_to = self.parent._semiring_commutes_to - if commutes_to is not None and self.type in commutes_to._typed_ops: + if commutes_to is not None and (self.type in commutes_to._typed_ops or self.type._is_udt): return commutes_to[self.type] @property def is_commutative(self): return self.commutes_to is self + @property + def type2(self): + return self.type if self._type2 is None else self._type2 + class TypedBuiltinMonoid(TypedOpBase): __slots__ = "_identity" @@ -236,6 +298,10 @@ def binaryop(self): def commutes_to(self): return self + @property + def type2(self): + return self.type + class TypedBuiltinSemiring(TypedOpBase): __slots__ = () @@ -274,41 +340,49 @@ def commutes_to(self): def is_commutative(self): return self.binaryop.is_commutative + type2 = TypedBuiltinBinaryOp.type2 + class TypedUserUnaryOp(TypedOpBase): - __slots__ = "orig_func", "numba_func" + __slots__ = () opclass = "UnaryOp" - def __init__(self, parent, name, type_, return_type, gb_obj, orig_func, numba_func): + def __init__(self, parent, name, type_, return_type, gb_obj): super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}") - self.orig_func = orig_func - self.numba_func = numba_func + + @property + def orig_func(self): + return self.parent.orig_func + + @property + def _numba_func(self): + return self.parent._numba_func __call__ = TypedBuiltinUnaryOp.__call__ class TypedUserBinaryOp(TypedOpBase): - __slots__ = "orig_func", "numba_func", "_monoid" + __slots__ = "_monoid" opclass = "BinaryOp" - def __init__(self, parent, name, type_, return_type, gb_obj, orig_func, numba_func): - super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}") - self.orig_func = orig_func - self.numba_func = numba_func + def __init__(self, parent, name, type_, return_type, gb_obj, dtype2=None): + super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}", dtype2=dtype2) self._monoid = None @property def monoid(self): - if self._monoid is None and not self.parent._anonymous: - monoid = Monoid._find(self.name) - if monoid is not None and self.type in monoid._typed_ops: # pragma: no cover - # This may be used by grblas.binary.numpy objects + if self._monoid is None: + monoid = self.parent.monoid + if monoid is not None and self.type in monoid: self._monoid = monoid[self.type] return self._monoid commutes_to = TypedBuiltinBinaryOp.commutes_to _semiring_commutes_to = TypedBuiltinBinaryOp._semiring_commutes_to is_commutative = TypedBuiltinBinaryOp.is_commutative + orig_func = TypedUserUnaryOp.orig_func + _numba_func = TypedUserUnaryOp._numba_func + type2 = TypedBuiltinBinaryOp.type2 __call__ = TypedBuiltinBinaryOp.__call__ @@ -324,6 +398,7 @@ def __init__(self, parent, name, type_, return_type, gb_obj, binaryop, identity) binaryop._monoid = self commutes_to = TypedBuiltinMonoid.commutes_to + type2 = TypedBuiltinMonoid.type2 __call__ = TypedBuiltinMonoid.__call__ @@ -331,13 +406,14 @@ class TypedUserSemiring(TypedOpBase): __slots__ = "monoid", "binaryop" opclass = "Semiring" - def __init__(self, parent, name, type_, return_type, gb_obj, monoid, binaryop): - super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}") + def __init__(self, parent, name, type_, return_type, gb_obj, monoid, binaryop, dtype2=None): + super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}", dtype2=dtype2) self.monoid = monoid self.binaryop = binaryop commutes_to = TypedBuiltinSemiring.commutes_to is_commutative = TypedBuiltinSemiring.is_commutative + type2 = TypedBuiltinBinaryOp.type2 __call__ = TypedBuiltinSemiring.__call__ @@ -360,11 +436,12 @@ def _call(self, *args, **kwargs): class ParameterizedUnaryOp(ParameterizedUdf): - __slots__ = "func", "__signature__" + __slots__ = "func", "__signature__", "_is_udt" - def __init__(self, name, func, *, anonymous=False): + def __init__(self, name, func, *, anonymous=False, is_udt=False): self.func = func self.__signature__ = inspect.signature(func) + self._is_udt = is_udt if name is None: name = getattr(func, "__name__", name) super().__init__(name, anonymous) @@ -372,7 +449,7 @@ def __init__(self, name, func, *, anonymous=False): def _call(self, *args, **kwargs): unary = self.func(*args, **kwargs) unary._parameterized_info = (self, args, kwargs) - return UnaryOp.register_anonymous(unary, self.name) + return UnaryOp.register_anonymous(unary, self.name, is_udt=self._is_udt) def __reduce__(self): name = f"unary.{self.name}" @@ -391,12 +468,13 @@ def _deserialize(name, func, anonymous): class ParameterizedBinaryOp(ParameterizedUdf): - __slots__ = "func", "__signature__", "_monoid", "_cached_call", "_commutes_to" + __slots__ = "func", "__signature__", "_monoid", "_cached_call", "_commutes_to", "_is_udt" - def __init__(self, name, func, *, anonymous=False): + def __init__(self, name, func, *, anonymous=False, is_udt=False): self.func = func self.__signature__ = inspect.signature(func) self._monoid = None + self._is_udt = is_udt if name is None: name = getattr(func, "__name__", name) super().__init__(name, anonymous) @@ -408,7 +486,7 @@ def __init__(self, name, func, *, anonymous=False): def _call_to_cache(self, *args, **kwargs): binary = self.func(*args, **kwargs) binary._parameterized_info = (self, args, kwargs) - return BinaryOp.register_anonymous(binary, self.name) + return BinaryOp.register_anonymous(binary, self.name, is_udt=self._is_udt) def _call(self, *args, **kwargs): binop = self._cached_call(*args, **kwargs) @@ -422,7 +500,7 @@ def _call(self, *args, **kwargs): self._monoid(*args, **kwargs) except Exception: binop._monoid = None - assert binop._monoid is not binop + # assert binop._monoid is not binop if self.is_commutative: binop._commutes_to = binop # Don't bother yet with creating `binop.commutes_to` (but we could!) @@ -570,7 +648,16 @@ def _deserialize(name, monoid, binaryop, anonymous): class OpBase: - __slots__ = "name", "_typed_ops", "types", "coercions", "_anonymous", "__weakref__" + __slots__ = ( + "name", + "_typed_ops", + "types", + "coercions", + "_anonymous", + "_udt_types", + "_udt_ops", + "__weakref__", + ) _parse_config = None _initialized = False _module = None @@ -582,28 +669,47 @@ def __init__(self, name, *, anonymous=False): self.types = {} self.coercions = {} self._anonymous = anonymous + self._udt_types = None + self._udt_ops = None def __repr__(self): return f"{self._modname}.{self.name}" def __getitem__(self, type_): - type_ = _normalize_type(type_) - if type_ not in self._typed_ops: - raise KeyError(f"{self.name} does not work with {type_}") - return self._typed_ops[type_] + if type(type_) is tuple: + dtype1, dtype2 = type_ + dtype1 = lookup_dtype(dtype1) + dtype2 = lookup_dtype(dtype2) + return get_typed_op(self, dtype1, dtype2) + elif not self._is_udt: + type_ = lookup_dtype(type_) + if type_ not in self._typed_ops: + if self._udt_types is None: + if self.is_positional: + return self._typed_ops[UINT64] + raise KeyError(f"{self.name} does not work with {type_}") + else: + return self._typed_ops[type_] + # This is a UDT or is able to operate on UDTs such as `first` any `any` + dtype = lookup_dtype(type_) + return self._compile_udt(dtype, dtype) def _add(self, op): self._typed_ops[op.type] = op self.types[op.type] = op.return_type def __delitem__(self, type_): - type_ = _normalize_type(type_) + type_ = lookup_dtype(type_) del self._typed_ops[type_] del self.types[type_] def __contains__(self, type_): - type_ = _normalize_type(type_) - return type_ in self._typed_ops + try: + self[type_] + except (TypeError, KeyError, numba.NumbaError): + return False + else: + return True @classmethod def _remove_nesting(cls, funcname, *, module=None, modname=None, strict=True): @@ -686,9 +792,9 @@ def _initialize(cls): gb_obj = getattr(lib, varname) # Determine return type if return_prefix == "BOOL": - return_type = "BOOL" + return_type = BOOL if type_ is None: - type_ = "BOOL" + type_ = BOOL else: if type_ is None: # pragma: no cover raise TypeError(f"Unable to determine return type for {varname}") @@ -701,7 +807,12 @@ def _initialize(cls): raise TypeError(f"Unexpected number of bits: {num_bits}") return_type = f"{return_prefix}{num_bits}" builtin_op = cls._typed_class( - obj, name, type_, return_type, gb_obj, gb_name + obj, + name, + lookup_dtype(type_), + lookup_dtype(return_type), + gb_obj, + gb_name, ) obj._add(builtin_op) @@ -713,8 +824,16 @@ def _deserialize(cls, name, *args): return cls.register_new(name, *args) +def _identity(x): + return x # pragma: no cover + + +def _one(x): + return 1 # pragma: no cover + + class UnaryOp(OpBase): - __slots__ = "_func", "is_positional" + __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func" _module = unary _modname = "unary" _typed_class = TypedBuiltinUnaryOp @@ -747,55 +866,51 @@ class UnaryOp(OpBase): _positional = {"positioni", "positioni1", "positionj", "positionj1"} @classmethod - def _build(cls, name, func, *, anonymous=False): + def _build(cls, name, func, *, anonymous=False, is_udt=False): if type(func) is not FunctionType: raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") success = False - new_type_obj = cls(name, func, anonymous=anonymous) + unary_udf = numba.njit(func) + new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=unary_udf) return_types = {} nt = numba.types - for type_, sample_val in _sample_values.items(): - type_ = lookup_dtype(type_) - # Check if func can handle this data type - try: - with np.errstate(divide="ignore", over="ignore", under="ignore", invalid="ignore"): - ret = func(sample_val) - ret_type = lookup_dtype(type(ret)) + if not is_udt: + for type_, sample_val in _sample_values.items(): + sig = (type_.numba_type,) + try: + unary_udf.compile(sig) + except numba.TypingError: + continue + ret_type = lookup_dtype(unary_udf.overloads[sig].signature.return_type) if ret_type != type_ and ( ("INT" in ret_type.name and "INT" in type_.name) or ("FP" in ret_type.name and "FP" in type_.name) or ("FC" in ret_type.name and "FC" in type_.name) - or ( - type_ == "UINT64" - and ret_type == "FP64" - and return_types.get("INT64") == "INT64" - ) + or (type_ == UINT64 and ret_type == FP64 and return_types.get(INT64) == INT64) ): # Downcast `ret_type` to `type_`. # This is what users want most of the time, but we can't make a perfect rule. # There should be a way for users to be explicit. ret_type = type_ - elif type_ == "BOOL" and ret_type == "INT64" and return_types.get("INT8") == "INT8": + elif type_ == BOOL and ret_type == INT64 and return_types.get(INT8) == INT8: ret_type = INT8 # Numba is unable to handle BOOL correctly right now, but we have a workaround # See: https://github.com/numba/numba/issues/5395 # We're relying on coercion behaving correctly here - input_type = INT8 if type_ == "BOOL" else type_ - return_type = INT8 if ret_type == "BOOL" else ret_type + input_type = INT8 if type_ == BOOL else type_ + return_type = INT8 if ret_type == BOOL else ret_type - # JIT the func so it can be used from a cfunc - unary_udf = numba.njit(func) # Build wrapper because GraphBLAS wants pointers and void return wrapper_sig = nt.void( nt.CPointer(return_type.numba_type), nt.CPointer(input_type.numba_type), ) - if type_ == "BOOL": - if ret_type == "BOOL": + if type_ == BOOL: + if ret_type == BOOL: def unary_wrapper(z, x): z[0] = bool(unary_udf(bool(x[0]))) # pragma: no cover @@ -805,7 +920,7 @@ def unary_wrapper(z, x): def unary_wrapper(z, x): z[0] = unary_udf(bool(x[0])) # pragma: no cover - elif ret_type == "BOOL": + elif ret_type == BOOL: def unary_wrapper(z, x): z[0] = bool(unary_udf(x[0])) # pragma: no cover @@ -824,27 +939,47 @@ def unary_wrapper(z, x): "UnaryOp", new_unary, ) - op = TypedUserUnaryOp( - new_type_obj, name, type_.name, ret_type.name, new_unary[0], func, unary_udf - ) + op = TypedUserUnaryOp(new_type_obj, name, type_, ret_type, new_unary[0]) new_type_obj._add(op) success = True - return_types[type_.name] = ret_type.name - except Exception: - continue - if success: + return_types[type_] = ret_type + if success or is_udt: return new_type_obj else: raise UdfParseError("Unable to parse function using Numba") + def _compile_udt(self, dtype, dtype2): + if dtype in self._udt_types: + return self._udt_ops[dtype] + + numba_func = self._numba_func + sig = (dtype.numba_type,) + numba_func.compile(sig) # Should we catch and give additional error message? + ret_type = lookup_dtype(numba_func.overloads[sig].signature.return_type) + if ret_type is not dtype and ret_type.gb_obj is dtype.gb_obj: + ret_type = dtype + + unary_wrapper, wrapper_sig = _get_udt_wrapper(numba_func, ret_type, dtype) + unary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(unary_wrapper) + new_unary = ffi_new("GrB_UnaryOp*") + check_status_carg( + lib.GrB_UnaryOp_new(new_unary, unary_wrapper.cffi, ret_type._carg, dtype._carg), + "UnaryOp", + new_unary, + ) + op = TypedUserUnaryOp(self, self.name, dtype, ret_type, new_unary[0]) + self._udt_types[dtype] = ret_type + self._udt_ops[dtype] = op + return op + @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False): + def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): if parameterized: - return ParameterizedUnaryOp(name, func, anonymous=True) - return cls._build(name, func, anonymous=True) + return ParameterizedUnaryOp(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt) @classmethod - def register_new(cls, name, func, *, parameterized=False, lazy=False): + def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( @@ -852,10 +987,10 @@ def register_new(cls, name, func, *, parameterized=False, lazy=False): {"name": name, "func": func, "parameterized": parameterized}, ) elif parameterized: - unary_op = ParameterizedUnaryOp(name, func) + unary_op = ParameterizedUnaryOp(name, func, is_udt=is_udt) setattr(module, funcname, unary_op) else: - unary_op = cls._build(name, func) + unary_op = cls._build(name, func, is_udt=is_udt) setattr(module, funcname, unary_op) # Also save it to `grblas.op` if not yet defined opmodule, funcname = cls._remove_nesting(name, module=op, modname="op", strict=False) @@ -873,6 +1008,19 @@ def register_new(cls, name, func, *, parameterized=False, lazy=False): def _initialize(cls): super()._initialize() # Update type information with sane coercion + position_dtypes = [ + BOOL, + FP32, + FP64, + INT8, + INT16, + UINT8, + UINT16, + UINT32, + UINT64, + ] + if _supports_complex: + position_dtypes.extend([FC32, FC64]) for names, *types in ( # fmt: off ( @@ -882,17 +1030,14 @@ def _initialize(cls): "log", "log10", "log1p", "log2", "round", "signum", "sin", "sinh", "sqrt", "tan", "tanh", "trunc", ), - (("BOOL", "INT8", "INT16", "UINT8", "UINT16"), "FP32"), - (("INT32", "INT64", "UINT32", "UINT64"), "FP64"), + ((BOOL, INT8, INT16, UINT8, UINT16), FP32), + ((INT32, INT64, UINT32, UINT64), FP64), ), ( ("positioni", "positioni1", "positionj", "positionj1"), ( - ( - "BOOL", "FC32", "FC64", "FP32", "FP64", "INT8", "INT16", - "UINT8", "UINT16", "UINT32", "UINT64", - ), - "INT64", + position_dtypes, + INT64, ), ), # fmt: on @@ -903,61 +1048,167 @@ def _initialize(cls): typed_op = op._typed_ops[target_type] output_type = op.types[target_type] for dtype in input_types: - if dtype not in op.types and ( - _supports_complex or dtype not in {"FC32", "FC64"} - ): + if dtype not in op.types: op.types[dtype] = output_type op._typed_ops[dtype] = typed_op op.coercions[dtype] = target_type + # Allow some functions to work on UDTs + for (unop, func) in [ + (unary.identity, _identity), + (unary.one, _one), + ]: + unop.orig_func = func + unop._numba_func = numba.njit(func) + unop._udt_types = {} + unop._udt_ops = {} cls._initialized = True - def __init__(self, name, func=None, *, anonymous=False, is_positional=False): + def __init__( + self, + name, + func=None, + *, + anonymous=False, + is_positional=False, + is_udt=False, + numba_func=None, + ): super().__init__(name, anonymous=anonymous) - self._func = func + self.orig_func = func + self._numba_func = numba_func self.is_positional = is_positional + self._is_udt = is_udt + if is_udt: + self._udt_types = {} # {dtype: DataType} + self._udt_ops = {} # {dtype: TypedUserUnaryOp} def __reduce__(self): if self._anonymous: - if hasattr(self._func, "_parameterized_info"): - return (_deserialize_parameterized, self._func._parameterized_info) - return (self.register_anonymous, (self._func, self.name)) + if hasattr(self.orig_func, "_parameterized_info"): + return (_deserialize_parameterized, self.orig_func._parameterized_info) + return (self.register_anonymous, (self.orig_func, self.name)) name = f"unary.{self.name}" if name in _STANDARD_OPERATOR_NAMES: return name - return (self._deserialize, (self.name, self._func)) + return (self._deserialize, (self.name, self.orig_func)) __call__ = TypedBuiltinUnaryOp.__call__ def _floordiv(x, y): - return x // y # cast to integer + return x // y # pragma: no cover def _rfloordiv(x, y): - return y // x # cast to integer + return y // x # pragma: no cover def _absfirst(x, y): - return abs(x) + return np.abs(x) # pragma: no cover def _abssecond(x, y): - return abs(y) + return np.abs(y) # pragma: no cover def _rpow(x, y): - return y**x + return y**x # pragma: no cover def _isclose(rel_tol=1e-7, abs_tol=0.0): - def inner(x, y): + def inner(x, y): # pragma: no cover return x == y or abs(x - y) <= max(rel_tol * max(abs(x), abs(y)), abs_tol) return inner +def _first(x, y): + return x # pragma: no cover + + +def _second(x, y): + return y # pragma: no cover + + +def _pair(x, y): + return 1 # pragma: no cover + + +def _get_udt_wrapper(numba_func, return_type, dtype, dtype2=None): + ztype = INT8 if return_type == BOOL else return_type + xtype = INT8 if dtype == BOOL else dtype + nt = numba.types + if dtype2 is not None: + ytype = INT8 if dtype2 == BOOL else dtype2 + wrapper_sig = nt.void( + nt.CPointer(ztype.numba_type), + nt.CPointer(xtype.numba_type), + nt.CPointer(ytype.numba_type), + ) + else: + wrapper_sig = nt.void( + nt.CPointer(ztype.numba_type), + nt.CPointer(xtype.numba_type), + ) + + zarray = xarray = yarray = BL = BR = yarg = yname = "" + if return_type._is_udt: + if return_type.np_type.subdtype is None: + zarray = " z = numba.carray(z_ptr, 1)\n" + zname = "z[0]" + else: + zname = "z_ptr[0]" + BR = "[0]" + else: + zname = "z_ptr[0]" + if return_type == BOOL: + BL = "bool(" + BR = ")" + + if dtype._is_udt: + if dtype.np_type.subdtype is None: + xarray = " x = numba.carray(x_ptr, 1)\n" + xname = "x[0]" + else: + xname = "x_ptr" + elif dtype == BOOL: + xname = "bool(x_ptr[0])" + else: + xname = "x_ptr[0]" + + if dtype2 is not None: + yarg = ", y_ptr" + if dtype2._is_udt: + if dtype2.np_type.subdtype is None: + yarray = " y = numba.carray(y_ptr, 1)\n" + yname = ", y[0]" + else: + yname = ", y_ptr" + elif dtype2 == BOOL: + yname = ", bool(y_ptr[0])" + else: + yname = ", y_ptr[0]" + + d = {"numba": numba, "numba_func": numba_func} + text = ( + f"def wrapper(z_ptr, x_ptr{yarg}):\n" + f"{zarray}{xarray}{yarray}" + f" {zname} = {BL}numba_func({xname}{yname}){BR}\n" + ) + exec(text, d) + return d["wrapper"], wrapper_sig + + class BinaryOp(OpBase): - __slots__ = "_monoid", "_commutes_to", "_semiring_commutes_to", "_func", "is_positional" + __slots__ = ( + "_monoid", + "_commutes_to", + "_semiring_commutes_to", + "orig_func", + "is_positional", + "_is_udt", + "_numba_func", + ) _module = binary _modname = "binary" _typed_class = TypedBuiltinBinaryOp @@ -1062,47 +1313,42 @@ class BinaryOp(OpBase): } @classmethod - def _build(cls, name, func, *, anonymous=False): + def _build(cls, name, func, *, is_udt=False, anonymous=False): if not isinstance(func, FunctionType): raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") success = False - new_type_obj = cls(name, func, anonymous=anonymous) + binary_udf = numba.njit(func) + new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=binary_udf) return_types = {} nt = numba.types - for type_, sample_val in _sample_values.items(): - type_ = lookup_dtype(type_) - # Check if func can handle this data type - try: - with np.errstate(divide="ignore", over="ignore", under="ignore", invalid="ignore"): - ret = func(sample_val, sample_val) - ret_type = lookup_dtype(type(ret)) + if not is_udt: + for type_, sample_val in _sample_values.items(): + sig = (type_.numba_type, type_.numba_type) + try: + binary_udf.compile(sig) + except numba.TypingError: + continue + ret_type = lookup_dtype(binary_udf.overloads[sig].signature.return_type) if ret_type != type_ and ( ("INT" in ret_type.name and "INT" in type_.name) or ("FP" in ret_type.name and "FP" in type_.name) or ("FC" in ret_type.name and "FC" in type_.name) - or ( - type_ == "UINT64" - and ret_type == "FP64" - and return_types.get("INT64") == "INT64" - ) + or (type_ == UINT64 and ret_type == FP64 and return_types.get(INT64) == INT64) ): # Downcast `ret_type` to `type_`. # This is what users want most of the time, but we can't make a perfect rule. # There should be a way for users to be explicit. ret_type = type_ - elif type_ == "BOOL" and ret_type == "INT64" and return_types.get("INT8") == "INT8": + elif type_ == BOOL and ret_type == INT64 and return_types.get(INT8) == INT8: ret_type = INT8 # Numba is unable to handle BOOL correctly right now, but we have a workaround # See: https://github.com/numba/numba/issues/5395 # We're relying on coercion behaving correctly here - input_type = INT8 if type_ == "BOOL" else type_ - return_type = INT8 if ret_type == "BOOL" else ret_type - - # JIT the func so it can be used from a cfunc - binary_udf = numba.njit(func) + input_type = INT8 if type_ == BOOL else type_ + return_type = INT8 if ret_type == BOOL else ret_type # Build wrapper because GraphBLAS wants pointers and void return wrapper_sig = nt.void( @@ -1111,8 +1357,8 @@ def _build(cls, name, func, *, anonymous=False): nt.CPointer(input_type.numba_type), ) - if type_ == "BOOL": - if ret_type == "BOOL": + if type_ == BOOL: + if ret_type == BOOL: def binary_wrapper(z, x, y): z[0] = bool(binary_udf(bool(x[0]), bool(y[0]))) # pragma: no cover @@ -1122,7 +1368,7 @@ def binary_wrapper(z, x, y): def binary_wrapper(z, x, y): z[0] = binary_udf(bool(x[0]), bool(y[0])) # pragma: no cover - elif ret_type == "BOOL": + elif ret_type == BOOL: def binary_wrapper(z, x, y): z[0] = bool(binary_udf(x[0], y[0])) # pragma: no cover @@ -1145,27 +1391,137 @@ def binary_wrapper(z, x, y): "BinaryOp", new_binary, ) - op = TypedUserBinaryOp( - new_type_obj, name, type_.name, ret_type.name, new_binary[0], func, binary_udf - ) + op = TypedUserBinaryOp(new_type_obj, name, type_, ret_type, new_binary[0]) new_type_obj._add(op) success = True - return_types[type_.name] = ret_type.name - except Exception: - continue - if success: + return_types[type_] = ret_type + if success or is_udt: return new_type_obj else: raise UdfParseError("Unable to parse function using Numba") + def _compile_udt(self, dtype, dtype2): + if dtype2 is None: + dtype2 = dtype + dtypes = (dtype, dtype2) + if dtypes in self._udt_types: + return self._udt_ops[dtypes] + + nt = numba.types + if self.name == "eq" and not self._anonymous: + # assert dtype.np_type == dtype2.np_type + itemsize = dtype.np_type.itemsize + mask = _udt_mask(dtype.np_type) + ret_type = BOOL + wrapper_sig = nt.void( + nt.CPointer(INT8.numba_type), + nt.CPointer(UINT8.numba_type), + nt.CPointer(UINT8.numba_type), + ) + # PERF: we can probably make this faster + if mask.all(): + + def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover + x = numba.carray(x_ptr, itemsize) + y = numba.carray(y_ptr, itemsize) + # for i in range(itemsize): + # if x[i] != y[i]: + # z_ptr[0] = False + # break + # else: + # z_ptr[0] = True + z_ptr[0] = (x == y).all() + + else: + + def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover + x = numba.carray(x_ptr, itemsize) + y = numba.carray(y_ptr, itemsize) + # for i in range(itemsize): + # if mask[i] and x[i] != y[i]: + # z_ptr[0] = False + # break + # else: + # z_ptr[0] = True + z_ptr[0] = (x[mask] == y[mask]).all() + + elif self.name == "ne" and not self._anonymous: + # assert dtype.np_type == dtype2.np_type + itemsize = dtype.np_type.itemsize + mask = _udt_mask(dtype.np_type) + ret_type = BOOL + wrapper_sig = nt.void( + nt.CPointer(INT8.numba_type), + nt.CPointer(UINT8.numba_type), + nt.CPointer(UINT8.numba_type), + ) + if mask.all(): + + def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover + x = numba.carray(x_ptr, itemsize) + y = numba.carray(y_ptr, itemsize) + # for i in range(itemsize): + # if x[i] != y[i]: + # z_ptr[0] = True + # break + # else: + # z_ptr[0] = False + z_ptr[0] = (x != y).any() + + else: + + def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover + x = numba.carray(x_ptr, itemsize) + y = numba.carray(y_ptr, itemsize) + # for i in range(itemsize): + # if mask[i] and x[i] != y[i]: + # z_ptr[0] = True + # break + # else: + # z_ptr[0] = False + z_ptr[0] = (x[mask] != y[mask]).any() + + else: + numba_func = self._numba_func + sig = (dtype.numba_type, dtype2.numba_type) + numba_func.compile(sig) # Should we catch and give additional error message? + ret_type = lookup_dtype(numba_func.overloads[sig].signature.return_type) + if ret_type is not dtype and ret_type is not dtype2: + if ret_type.gb_obj is dtype.gb_obj: + ret_type = dtype + elif ret_type.gb_obj is dtype2.gb_obj: + ret_type = dtype2 + binary_wrapper, wrapper_sig = _get_udt_wrapper(numba_func, ret_type, dtype, dtype2) + + binary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(binary_wrapper) + new_binary = ffi_new("GrB_BinaryOp*") + check_status_carg( + lib.GrB_BinaryOp_new( + new_binary, binary_wrapper.cffi, ret_type._carg, dtype._carg, dtype2._carg + ), + "BinaryOp", + new_binary, + ) + op = TypedUserBinaryOp( + self, + self.name, + dtype, + ret_type, + new_binary[0], + dtype2=dtype2, + ) + self._udt_types[dtypes] = ret_type + self._udt_ops[dtypes] = op + return op + @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False): + def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): if parameterized: - return ParameterizedBinaryOp(name, func, anonymous=True) - return cls._build(name, func, anonymous=True) + return ParameterizedBinaryOp(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt) @classmethod - def register_new(cls, name, func, *, parameterized=False, lazy=False): + def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( @@ -1173,10 +1529,10 @@ def register_new(cls, name, func, *, parameterized=False, lazy=False): {"name": name, "func": func, "parameterized": parameterized}, ) elif parameterized: - binary_op = ParameterizedBinaryOp(name, func) + binary_op = ParameterizedBinaryOp(name, func, is_udt=is_udt) setattr(module, funcname, binary_op) else: - binary_op = cls._build(name, func) + binary_op = cls._build(name, func, is_udt=is_udt) setattr(module, funcname, binary_op) # Also save it to `grblas.op` if not yet defined opmodule, funcname = cls._remove_nesting(name, module=op, modname="op", strict=False) @@ -1210,10 +1566,10 @@ def _initialize(cls): rtruediv = binary.rtruediv = op.rtruediv = BinaryOp("rtruediv") for (new_op, builtin_op) in [(truediv, binary.cdiv), (rtruediv, binary.rdiv)]: for dtype in builtin_op.types: - if dtype in {"FP32", "FC32", "FC64"}: + if dtype.name in {"FP32", "FC32", "FC64"}: orig_dtype = dtype else: - orig_dtype = "FP64" + orig_dtype = FP64 orig_op = builtin_op[orig_dtype] cur_op = TypedBuiltinBinaryOp( new_op, @@ -1237,33 +1593,43 @@ def _initialize(cls): BinaryOp.register_new("isclose", _isclose, parameterized=True) # Update type information with sane coercion + position_dtypes = [ + BOOL, + FP32, + FP64, + INT8, + INT16, + UINT8, + UINT16, + UINT32, + UINT64, + ] + if _supports_complex: + position_dtypes.extend([FC32, FC64]) name_types = [ # fmt: off ( ("atan2", "copysign", "fmod", "hypot", "ldexp", "remainder"), - (("BOOL", "INT8", "INT16", "UINT8", "UINT16"), "FP32"), - (("INT32", "INT64", "UINT32", "UINT64"), "FP64"), + ((BOOL, INT8, INT16, UINT8, UINT16), FP32), + ((INT32, INT64, UINT32, UINT64), FP64), ), ( ( "firsti", "firsti1", "firstj", "firstj1", "secondi", "secondi1", "secondj", "secondj1",), ( - ( - "BOOL", "FC32", "FC64", "FP32", "FP64", "INT8", "INT16", - "UINT8", "UINT16", "UINT32", "UINT64", - ), - "INT64", + position_dtypes, + INT64, ), ), ( ["lxnor"], ( ( - "FP32", "FP64", "INT8", "INT16", "INT32", "INT64", - "UINT8", "UINT16", "UINT32", "UINT64", + FP32, FP64, INT8, INT16, INT32, INT64, + UINT8, UINT16, UINT32, UINT64, ), - "BOOL", + BOOL, ), ), # fmt: on @@ -1272,8 +1638,8 @@ def _initialize(cls): name_types.append( ( ["cmplx"], - (("BOOL", "INT8", "INT16", "UINT8", "UINT16"), "FP32"), - (("INT32", "INT64", "UINT32", "UINT64"), "FP64"), + ((BOOL, INT8, INT16, UINT8, UINT16), FP32), + ((INT32, INT64, UINT32, UINT64), FP64), ) ) for names, *types in name_types: @@ -1283,15 +1649,13 @@ def _initialize(cls): typed_op = cur_op._typed_ops[target_type] output_type = cur_op.types[target_type] for dtype in input_types: - if dtype not in cur_op.types and ( - _supports_complex or dtype not in {"FC32", "FC64"} - ): + if dtype not in cur_op.types: # pragma: no branch cur_op.types[dtype] = output_type cur_op._typed_ops[dtype] = typed_op cur_op.coercions[dtype] = target_type # Not valid input dtypes - del binary.ldexp["FP32"] - del binary.ldexp["FP64"] + del binary.ldexp[FP32] + del binary.ldexp[FP64] # Fill in commutes info for left_name, right_name in cls._commutes.items(): left = getattr(binary, left_name) @@ -1307,25 +1671,55 @@ def _initialize(cls): right = getattr(binary, right) left._semiring_commutes_to = right right._semiring_commutes_to = left + # Allow some functions to work on UDTs + for (binop, func) in [ + (binary.first, _first), + (binary.second, _second), + (binary.pair, _pair), + (binary.any, _first), + ]: + binop.orig_func = func + binop._numba_func = numba.njit(func) + binop._udt_types = {} + binop._udt_ops = {} + binary.any._numba_func = binary.first._numba_func + binary.eq._udt_types = {} + binary.eq._udt_ops = {} + binary.ne._udt_types = {} + binary.ne._udt_ops = {} cls._initialized = True - def __init__(self, name, func=None, *, anonymous=False, is_positional=False): + def __init__( + self, + name, + func=None, + *, + anonymous=False, + is_positional=False, + is_udt=False, + numba_func=None, + ): super().__init__(name, anonymous=anonymous) self._monoid = None self._commutes_to = None self._semiring_commutes_to = None - self._func = func + self.orig_func = func + self._numba_func = numba_func + self._is_udt = is_udt self.is_positional = is_positional + if is_udt: + self._udt_types = {} # {(dtype, dtype): DataType} + self._udt_ops = {} # {(dtype, dtype): TypedUserBinaryOp} def __reduce__(self): if self._anonymous: - if hasattr(self._func, "_parameterized_info"): - return (_deserialize_parameterized, self._func._parameterized_info) - return (self.register_anonymous, (self._func, self.name)) + if hasattr(self.orig_func, "_parameterized_info"): + return (_deserialize_parameterized, self.orig_func._parameterized_info) + return (self.register_anonymous, (self.orig_func, self.name)) name = f"binary.{self.name}" if name in _STANDARD_OPERATOR_NAMES: return name - return (self._deserialize, (self.name, self._func)) + return (self._deserialize, (self.name, self.orig_func)) __call__ = TypedBuiltinBinaryOp.__call__ is_commutative = TypedBuiltinBinaryOp.is_commutative @@ -1370,32 +1764,69 @@ def _build(cls, name, binaryop, identity, *, anonymous=False): if name is None: name = binaryop.name new_type_obj = cls(name, binaryop, identity, anonymous=anonymous) - if not isinstance(identity, Mapping): - identities = dict.fromkeys(binaryop.types, identity) - explicit_identities = False - else: - identities = identity - explicit_identities = True - for type_, identity in identities.items(): - type_ = lookup_dtype(type_) - ret_type = binaryop[type_].return_type - # If there is a domain mismatch, then DomainMismatch will be raised - # below if identities were explicitly given. - # Skip complex dtypes for now, because they segfault! - if type_ != ret_type and not explicit_identities or "FC" in type_.name: - continue - new_monoid = ffi_new("GrB_Monoid*") - func = libget(f"GrB_Monoid_new_{type_.name}") - zcast = ffi.cast(type_.c_type, identity) - check_status_carg( - func(new_monoid, binaryop[type_].gb_obj, zcast), "Monoid", new_monoid[0] - ) - op = TypedUserMonoid( - new_type_obj, name, type_.name, ret_type, new_monoid[0], binaryop[type_], identity - ) - new_type_obj._add(op) + if not binaryop._is_udt: + if not isinstance(identity, Mapping): + identities = dict.fromkeys(binaryop.types, identity) + explicit_identities = False + else: + identities = {lookup_dtype(key): val for key, val in identity.items()} + explicit_identities = True + for type_, identity in identities.items(): + ret_type = binaryop[type_].return_type + # If there is a domain mismatch, then DomainMismatch will be raised + # below if identities were explicitly given. + # Skip complex dtypes for now, because they segfault! + if type_ != ret_type and not explicit_identities or "FC" in type_.name: + continue + new_monoid = ffi_new("GrB_Monoid*") + func = libget(f"GrB_Monoid_new_{type_.name}") + zcast = ffi.cast(type_.c_type, identity) + check_status_carg( + func(new_monoid, binaryop[type_].gb_obj, zcast), "Monoid", new_monoid[0] + ) + op = TypedUserMonoid( + new_type_obj, + name, + type_, + ret_type, + new_monoid[0], + binaryop[type_], + identity, + ) + new_type_obj._add(op) return new_type_obj + def _compile_udt(self, dtype, dtype2): + if dtype2 is None: + dtype2 = dtype + elif dtype != dtype2: + raise TypeError( + "Monoid inputs must be the same dtype (got {dtype} and {dtype2}); " + "unable to coerce when using UDTs." + ) + if dtype in self._udt_types: + return self._udt_ops[dtype] + binaryop = self.binaryop._compile_udt(dtype, dtype2) + from .scalar import Scalar + + ret_type = binaryop.return_type + identity = Scalar.from_value(self._identity, dtype=ret_type, is_cscalar=True) + new_monoid = ffi_new("GrB_Monoid*") + status = lib.GrB_Monoid_new_UDT(new_monoid, binaryop.gb_obj, identity.gb_obj) + check_status_carg(status, "Monoid", new_monoid[0]) + op = TypedUserMonoid( + new_monoid, + self.name, + dtype, + ret_type, + new_monoid[0], + binaryop, + identity, + ) + self._udt_types[dtype] = dtype + self._udt_ops[dtype] = op + return op + @classmethod def register_anonymous(cls, binaryop, identity, name=None): if type(binaryop) is ParameterizedBinaryOp: @@ -1434,6 +1865,9 @@ def __init__(self, name, binaryop=None, identity=None, *, anonymous=False): self._identity = identity if binaryop is not None: binaryop._monoid = self + if binaryop._is_udt: + self._udt_types = {} # {dtype: DataType} + self._udt_ops = {} # {dtype: TypedUserMonoid} def __reduce__(self): if self._anonymous: @@ -1454,41 +1888,50 @@ def binaryop(self): def identities(self): return {dtype: val.identity for dtype, val in self._typed_ops.items()} + @property + def _is_udt(self): + return self._binaryop is not None and self._binaryop._is_udt + @classmethod def _initialize(cls): super()._initialize() - lor = monoid.lor._typed_ops["BOOL"] - land = monoid.land._typed_ops["BOOL"] + lor = monoid.lor._typed_ops[BOOL] + land = monoid.land._typed_ops[BOOL] for cur_op, typed_op in [ (monoid.max, lor), (monoid.min, land), # (monoid.plus, lor), # two choices: lor, or plus[int] (monoid.times, land), ]: - if "BOOL" not in cur_op.types: # pragma: no branch - cur_op.types["BOOL"] = "BOOL" - cur_op.coercions["BOOL"] = "BOOL" - cur_op._typed_ops["BOOL"] = typed_op + if BOOL not in cur_op.types: # pragma: no branch + cur_op.types[BOOL] = BOOL + cur_op.coercions[BOOL] = BOOL + cur_op._typed_ops[BOOL] = typed_op for cur_op in (monoid.lor, monoid.land, monoid.lxnor, monoid.lxor): - bool_op = cur_op._typed_ops["BOOL"] + bool_op = cur_op._typed_ops[BOOL] for dtype in ( - "FP32", - "FP64", - "INT8", - "INT16", - "INT32", - "INT64", - "UINT8", - "UINT16", - "UINT32", - "UINT64", + FP32, + FP64, + INT8, + INT16, + INT32, + INT64, + UINT8, + UINT16, + UINT32, + UINT64, ): if dtype in cur_op.types: # pragma: no cover continue - cur_op.types[dtype] = "BOOL" - cur_op.coercions[dtype] = "BOOL" + cur_op.types[dtype] = BOOL + cur_op.coercions[dtype] = BOOL cur_op._typed_ops[dtype] = bool_op + # Allow some functions to work on UDTs + any_ = monoid.any + any_._identity = 0 + any_._udt_types = {} + any_._udt_ops = {} cls._initialized = True commutes_to = TypedBuiltinMonoid.commutes_to @@ -1545,6 +1988,8 @@ def _build(cls, name, monoid, binaryop, *, anonymous=False): if name is None: name = f"{monoid.name}_{binaryop.name}".replace(".", "_") new_type_obj = cls(name, monoid, binaryop, anonymous=anonymous) + if binaryop._is_udt: + return new_type_obj for binary_in, binary_func in binaryop._typed_ops.items(): binary_out = binary_func.return_type # Unfortunately, we can't have user-defined monoids over bools yet @@ -1555,7 +2000,6 @@ def _build(cls, name, monoid, binaryop, *, anonymous=False): or monoid.coercions.get(binary_out, binary_out) != binary_out ): continue - binary_out = lookup_dtype(binary_out) new_semiring = ffi_new("GrB_Semiring*") check_status_carg( lib.GrB_Semiring_new(new_semiring, monoid[binary_out].gb_obj, binary_func.gb_obj), @@ -1575,6 +2019,32 @@ def _build(cls, name, monoid, binaryop, *, anonymous=False): new_type_obj._add(op) return new_type_obj + def _compile_udt(self, dtype, dtype2): + if dtype2 is None: + dtype2 = dtype + dtypes = (dtype, dtype2) + if dtypes in self._udt_types: + return self._udt_ops[dtypes] + binaryop = self.binaryop._compile_udt(dtype, dtype2) + monoid = self.monoid[binaryop.return_type] + ret_type = monoid.return_type + new_semiring = ffi_new("GrB_Semiring*") + status = lib.GrB_Semiring_new(new_semiring, monoid.gb_obj, binaryop.gb_obj) + check_status_carg(status, "Semiring", new_semiring) + op = TypedUserSemiring( + new_semiring, + self.name, + dtype, + ret_type, + new_semiring[0], + monoid, + binaryop, + dtype2=dtype2, + ) + self._udt_types[dtypes] = dtype + self._udt_ops[dtypes] = op + return op + @classmethod def register_anonymous(cls, monoid, binaryop, name=None): if type(monoid) is ParameterizedMonoid or type(binaryop) is ParameterizedBinaryOp: @@ -1655,13 +2125,39 @@ def _initialize(cls): if not hasattr(semiring, target_name): continue target_op = getattr(semiring, target_name) - if "BOOL" not in target_op.types: # pragma: no branch + if BOOL not in target_op.types: # pragma: no branch source_op = getattr(semiring, source_name) - typed_op = source_op._typed_ops["BOOL"] - target_op.types["BOOL"] = "BOOL" - target_op._typed_ops["BOOL"] = typed_op - target_op.coercions[dtype] = "BOOL" - + typed_op = source_op._typed_ops[BOOL] + target_op.types[BOOL] = BOOL + target_op._typed_ops[BOOL] = typed_op + target_op.coercions[dtype] = BOOL + + position_dtypes = [ + BOOL, + FP32, + FP64, + INT8, + INT16, + UINT8, + UINT16, + UINT32, + UINT64, + ] + notbool_dtypes = [ + FP32, + FP64, + INT8, + INT16, + INT32, + INT64, + UINT8, + UINT16, + UINT32, + UINT64, + ] + if _supports_complex: + position_dtypes.extend([FC32, FC64]) + notbool_dtypes.extend([FC32, FC64]) for lnames, rnames, *types in ( # fmt: off ( @@ -1671,11 +2167,8 @@ def _initialize(cls): "secondi", "secondi1", "secondj", "secondj1", ), ( - ( - "BOOL", "FC32", "FC64", "FP32", "FP64", "INT8", "INT16", - "UINT8", "UINT16", "UINT32", "UINT64", - ), - "INT64", + position_dtypes, + INT64, ), ), ( @@ -1683,30 +2176,27 @@ def _initialize(cls): ("first", "pair", "second"), # TODO: check if FC coercion works here ( - ( - "FC32", "FC64", "FP32", "FP64", "INT8", "INT16", "INT32", "INT64", - "UINT8", "UINT16", "UINT32", "UINT64", - ), - "BOOL", + notbool_dtypes, + BOOL, ), ), ( ("band", "bor", "bxnor", "bxor"), ("band", "bor", "bxnor", "bxor"), - (["INT8"], "UINT16"), - (["INT16"], "UINT32"), - (["INT32"], "UINT64"), - (["INT64"], "UINT64"), + ([INT8], UINT16), + ([INT16], UINT32), + ([INT32], UINT64), + ([INT64], UINT64), ), ( ("any", "eq", "land", "lor", "lxnor", "lxor"), ("eq", "land", "lor", "lxnor", "lxor", "ne"), ( ( - "FP32", "FP64", "INT8", "INT16", "INT32", "INT64", - "UINT8", "UINT16", "UINT32", "UINT64", + FP32, FP64, INT8, INT16, INT32, INT64, + UINT8, UINT16, UINT32, UINT64, ), - "BOOL", + BOOL, ), ), # fmt: on @@ -1720,9 +2210,7 @@ def _initialize(cls): typed_op = cur_op._typed_ops[target_type] output_type = cur_op.types[target_type] for dtype in input_types: - if dtype not in cur_op.types and ( - _supports_complex or dtype not in {"FC32", "FC64"} - ): + if dtype not in cur_op.types: cur_op.types[dtype] = output_type cur_op._typed_ops[dtype] = typed_op cur_op.coercions[dtype] = target_type @@ -1742,17 +2230,24 @@ def _initialize(cls): ): cur_op = getattr(semiring, opname) target = getattr(semiring, targetname) - if "BOOL" in cur_op.types or "BOOL" not in target.types: # pragma: no cover + if BOOL in cur_op.types or BOOL not in target.types: # pragma: no cover continue - cur_op.types["BOOL"] = target.types["BOOL"] - cur_op._typed_ops["BOOL"] = target._typed_ops["BOOL"] - cur_op.coercions["BOOL"] = "BOOL" + cur_op.types[BOOL] = target.types[BOOL] + cur_op._typed_ops[BOOL] = target._typed_ops[BOOL] + cur_op.coercions[BOOL] = BOOL cls._initialized = True def __init__(self, name, monoid=None, binaryop=None, *, anonymous=False): super().__init__(name, anonymous=anonymous) self._monoid = monoid self._binaryop = binaryop + try: + if self.binaryop._udt_types is not None: + self._udt_types = {} # {(dtype, dtype): DataType} + self._udt_ops = {} # {(dtype, dtype): TypedUserSemiring} + except AttributeError: + # `*_div` semirings raise here, but don't need `_udt_types` + pass def __reduce__(self): if self._anonymous: @@ -1780,6 +2275,10 @@ def monoid(self): def is_positional(self): return self.binaryop.is_positional + @property + def _is_udt(self): + return self._binaryop is not None and self._binaryop._is_udt + commutes_to = TypedBuiltinSemiring.commutes_to is_commutative = TypedBuiltinSemiring.is_commutative __call__ = TypedBuiltinSemiring.__call__ @@ -1787,10 +2286,19 @@ def is_positional(self): def get_typed_op(op, dtype, dtype2=None, *, is_left_scalar=False, is_right_scalar=False, kind=None): if isinstance(op, OpBase): + if op._is_udt: + return op._compile_udt(dtype, dtype2) if dtype2 is not None: - dtype = unify( - dtype, dtype2, is_left_scalar=is_left_scalar, is_right_scalar=is_right_scalar - ) + try: + dtype = unify( + dtype, dtype2, is_left_scalar=is_left_scalar, is_right_scalar=is_right_scalar + ) + except (TypeError, AttributeError): + if op.is_positional: + return op[UINT64] + if op._udt_types is None: + raise + return op._compile_udt(dtype, dtype2) return op[dtype] elif isinstance(op, ParameterizedUdf): op = op() # Use default parameters of parameterized UDFs diff --git a/grblas/scalar.py b/grblas/scalar.py index 213a2dcb1..7944b0259 100644 --- a/grblas/scalar.py +++ b/grblas/scalar.py @@ -38,7 +38,7 @@ def __init__(self, gb_obj, dtype, *, empty=False, is_cscalar=False, name=None): def __del__(self): gb_obj = getattr(self, "gb_obj", None) - if gb_obj is not None and "GB_Scalar_opaque" in str(gb_obj): + if gb_obj is not None and lib is not None and "GB_Scalar_opaque" in str(gb_obj): # it's difficult/dangerous to record the call, b/c `self.name` may not exist check_status(lib.GrB_Scalar_free(gb_obj), self) @@ -90,7 +90,7 @@ def __complex__(self): def __neg__(self): dtype = self.dtype - if dtype.name[0] == "U" or dtype == BOOL: + if dtype.name[0] == "U" or dtype == BOOL or dtype._is_udt: raise TypeError(f"The negative operator, `-`, is not supported for {dtype.name} dtype") rv = Scalar.new(dtype, is_cscalar=self._is_cscalar, name=f"-{self.name or 's_temp'}") if self._is_empty: @@ -138,7 +138,8 @@ def isequal(self, other, *, check_dtype=False): if other is None: return self._is_empty try: - other = Scalar.from_value(other, is_cscalar=None, name="s_isequal") + 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( other, @@ -157,7 +158,11 @@ def isequal(self, other, *, check_dtype=False): return False else: # For now, compare values in Python - return self.value == other.value + rv = self.value == other.value + try: + return bool(rv) + except ValueError: + return bool(rv.all()) def isclose(self, other, *, rel_tol=1e-7, abs_tol=0.0, check_dtype=False): """ @@ -197,7 +202,7 @@ def isclose(self, other, *, rel_tol=1e-7, abs_tol=0.0, check_dtype=False): is_right_scalar=True, kind="binary", ) - return isclose_func.numba_func(self.value, other.value) + return isclose_func._numba_func(self.value, other.value) def clear(self): if self._is_empty: @@ -224,11 +229,21 @@ def _is_empty(self): def value(self): if self._is_empty: return None + is_udt = self.dtype._is_udt if self._is_cscalar: - return self.gb_obj[0] + scalar = self else: scalar = Scalar.new(self.dtype, is_cscalar=True, name="s_temp") - call(f"GrB_Scalar_extractElement_{self.dtype.name}", [_Pointer(scalar), 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] + else: + return rv.view(np_type.subdtype[0]).reshape(np_type.subdtype[1]) + else: return scalar.gb_obj[0] @value.setter @@ -238,11 +253,26 @@ def value(self, val): elif self._is_cscalar: if output_type(val) is Scalar: val = val.value # raise below if wrong type (as determined by cffi) - self.gb_obj[0] = val + if self.dtype._is_udt: + np_type = self.dtype.np_type + if np_type.subdtype is None: + arr = np.empty(1, dtype=np_type) + else: + arr = np.empty(np_type.subdtype[1], dtype=np_type.subdtype[0]) + arr[:] = val + self.gb_obj[0 : self.dtype.np_type.itemsize] = arr.view(np.uint8) + # self.gb_obj[0:self.dtype.np_type.itemsize] = bytes(val) + else: + self.gb_obj[0] = val self._empty = False else: - val = _as_scalar(val, is_cscalar=True) - call(f"GrB_Scalar_setElement_{val.dtype}", [self, val]) + if self.dtype._is_udt: + val = _Pointer(_as_scalar(val, self.dtype, is_cscalar=True)) + dtype_name = "UDT" + else: + val = _as_scalar(val, is_cscalar=True) + dtype_name = val.dtype.name + call(f"GrB_Scalar_setElement_{dtype_name}", [self, val]) @property def nvals(self): @@ -287,7 +317,11 @@ def dup(self, dtype=None, *, is_cscalar=None, name=None): else: new_scalar = Scalar.new(dtype, is_cscalar=is_cscalar, name=name) if not self._is_empty: - new_scalar.value = new_scalar.dtype.np_type(self.value) + if 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: + new_scalar.value = self.value return new_scalar def wait(self): @@ -313,7 +347,10 @@ def new(cls, dtype, *, is_cscalar=False, name=None): # or to create a C scalar from a Python scalar. For example, see Matrix.apply. is_cscalar = True # pragma: to_grb if is_cscalar: - new_scalar = ffi_new(f"{dtype.c_type}*") + if dtype._is_udt: + new_scalar = ffi_new(dtype.c_type) + else: + new_scalar = ffi_new(f"{dtype.c_type}*") else: new_scalar = ffi_new("GrB_Scalar*") rv = cls(new_scalar, dtype, name=name, empty=True, is_cscalar=is_cscalar) @@ -330,7 +367,7 @@ def from_value(cls, value, dtype=None, *, is_cscalar=False, name=None): dtype = value.dtype else: try: - dtype = lookup_dtype(type(value)) + dtype = lookup_dtype(type(value), value) except ValueError: raise TypeError( f"Argument of from_value must be a known scalar type, not {type(value)}" diff --git a/grblas/tests/pickle3.pkl b/grblas/tests/pickle3.pkl new file mode 100644 index 000000000..4f60f8134 Binary files /dev/null and b/grblas/tests/pickle3.pkl differ diff --git a/grblas/tests/test_dtype.py b/grblas/tests/test_dtype.py index 680fb1af1..7967886aa 100644 --- a/grblas/tests/test_dtype.py +++ b/grblas/tests/test_dtype.py @@ -140,5 +140,45 @@ def test_pickle(): def test_unify_matches_numpy(): for type1, type2 in itertools.product(all_dtypes, all_dtypes): gb_type = dtypes.unify(type1, type2) - np_type = type(type1.np_type(0) + type2.np_type(0)) + np_type = type(type1.np_type.type(0) + type2.np_type.type(0)) assert gb_type is lookup_dtype(np_type), f"({type1}, {type2}) -> {gb_type}" + + +def test_lt_dtypes(): + expected = [ + dtypes.BOOL, + dtypes.FP32, + dtypes.FP64, + dtypes.INT8, + dtypes.INT16, + dtypes.INT32, + dtypes.INT64, + dtypes.UINT8, + dtypes.UINT16, + dtypes.UINT32, + dtypes.UINT64, + ] + if dtypes._supports_complex: + expected.insert(1, dtypes.FC32) + expected.insert(2, dtypes.FC64) + assert sorted(all_dtypes) == expected + assert dtypes.BOOL < "FP32" + with pytest.raises(TypeError): + dtypes.BOOL < 5 + + +def test_bad_register(): + record_dtype = np.dtype([("x", np.object_), ("y", np.float64)], align=True) + with pytest.raises(ValueError, match="Python object"): + dtypes.register_new("has_object", record_dtype) + record_dtype = np.dtype([("x", np.bool_), ("y", np.float64)], align=True) + with pytest.raises(ValueError, match="identifier"): + dtypes.register_new("$", record_dtype) + with pytest.raises(ValueError, match="builtin"): + dtypes.register_new("is_builtin", np.int8) + udt = dtypes.register_anonymous(record_dtype) + assert udt.name is not None + with pytest.raises(ValueError, match="name"): + dtypes.register_new("register_new", record_dtype) + with pytest.raises(ValueError, match="name"): + dtypes.register_new("UINT8", record_dtype) diff --git a/grblas/tests/test_formatting.py b/grblas/tests/test_formatting.py index 423bf8d39..26179464a 100644 --- a/grblas/tests/test_formatting.py +++ b/grblas/tests/test_formatting.py @@ -2,7 +2,7 @@ import pytest import grblas -from grblas import formatting, unary +from grblas import dtypes, formatting, unary from grblas.formatting import CSS_STYLE from .conftest import autocompute @@ -8054,3 +8054,52 @@ def test_index_expr_autocompute(v): "\n" "This expression may be used to extract or assign a Vector.
Example extract: v[[0, 1]].new()
Example assign: v[[0, 1]] << v" ) + + +def test_udt(): + record_dtype = np.dtype([("x", np.bool_), ("y", np.int64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "record_dtype") + v = Vector.new(udt, size=2) + v[:] = (False, 2) + v[1] = (True, 3) + repr_printer(v, "v") + assert repr(v) == ( + '"v_0" nvals size dtype format\n' + "grblas.Vector 2 2 record_dtype full\n" + "------------------------------------------------\n" + " 0 1\n" + " (False, 2) (True, 3)" + ) + A = Matrix.new(udt, nrows=1, ncols=3) + A[:, :2] = 0 + repr_printer(A, "A") + assert repr(A) == ( + '"M_0" nvals nrows ncols dtype format\n' + "grblas.Matrix 2 1 3 record_dtype bitmapr (iso)\n" + "---------------------------------------------------------------\n" + " 0 1 2\n" + "0 (False, 0) (False, 0) " + ) + np_dtype = np.dtype("(3,)uint16") + udt2 = dtypes.register_anonymous(np_dtype, "has_subdtype") + v = Vector.new(udt2, size=2) + v[:] = (1, 2, 3) + v[1] = (3, 2, 1) + repr_printer(v, "v") + assert repr(v) == ( + '"v_1" nvals size dtype format\n' + "grblas.Vector 2 2 has_subdtype full\n" + "------------------------------------------------\n" + " 0 1\n" + " [1, 2, 3] [3, 2, 1]" + ) + A = Matrix.new(udt2, nrows=1, ncols=3) + A[:, :2] = 1 + repr_printer(A, "A") + assert repr(A) == ( + '"M_1" nvals nrows ncols dtype format\n' + "grblas.Matrix 2 1 3 has_subdtype bitmapr (iso)\n" + "---------------------------------------------------------------\n" + " 0 1 2\n" + "0 [1, 1, 1] [1, 1, 1] " + ) diff --git a/grblas/tests/test_matrix.py b/grblas/tests/test_matrix.py index 7fc39b8a2..7d0d6de86 100644 --- a/grblas/tests/test_matrix.py +++ b/grblas/tests/test_matrix.py @@ -3129,3 +3129,96 @@ def test_delete_via_scalar(A): assert A[:, 0].new().nvals == 0 del A[:, :] assert A.nvals == 0 + + +def test_udt(): + record_dtype = np.dtype([("x", np.bool_), ("y", np.float64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "MatrixUDT") + A = Matrix.new(udt, nrows=2, ncols=2) + a = np.zeros(1, dtype=record_dtype) + A[0, 0] = a[0] + expected = Matrix.from_values([0], [0], a, nrows=2, ncols=2, dtype=udt) + assert A.isequal(expected) + A[0, 1] = (1, 2) + expected = Matrix.from_values( + [0, 0], [0, 1], np.array([(0, 0), (1, 2)], dtype=record_dtype), nrows=2, ncols=2, dtype=udt + ) + assert A.isequal(expected) + A[:, :] = 0 + zeros = Matrix.from_values([0, 0, 1, 1], [0, 1, 0, 1], 0, dtype=udt) + assert A.isequal(zeros) + A(A.S)[:, :] = 1 + ones = Matrix.from_values([0, 0, 1, 1], [0, 1, 0, 1], [1, 1, 1, 1], dtype=udt) + assert A.isequal(ones) + A[:, :](A.S) << 0 + assert A.isequal(zeros) + s = Scalar.new(udt) + s.value = (1, 1) + A(A.S)[:, :] = s + assert A.isequal(ones) + s.value = 0 + A[:, :](A.S) << s + assert A.isequal(zeros) + A[0, :] = 1 + A[[1], :] = 1 + assert A.isequal(ones) + A[:, [0]] = s + A[:, 1] = s + A[0, 0] = s + assert A.isequal(zeros) + t = A.reduce_scalar(monoid.any, allow_empty=True).new() + assert s == t + t = A.reduce_scalar(monoid.any, allow_empty=False).new() + assert s == t + A << binary.first(1, A) + assert A.isequal(ones) + A << binary.first(s, A) + assert A.isequal(zeros) + A << binary.second(A, 1) + assert A.isequal(ones) + A << binary.second(A, s) + assert A.isequal(zeros) + assert A[0, 0].new() == s + assert A[:, :].new().isequal(A) + expected = Vector.from_values([0, 1], s) + assert A[0, :].new().isequal(expected) + assert A.reduce_rowwise(monoid.any).new().isequal(expected) + rows, cols, values = A.to_values() + assert A.isequal(Matrix.from_values(rows, cols, values)) + assert A.isequal(Matrix.from_values(rows, cols, values, dtype=A.dtype)) + info = A.ss.export() + result = A.ss.import_any(**info) + assert result.isequal(A) + info = A.ss.export("cooc") + result = A.ss.import_any(**info) + assert result.isequal(A) + result = unary.positioni(A).new() + expected = Matrix.from_values([0, 0, 1, 1], [0, 1, 0, 1], [0, 0, 1, 1]) + assert result.isequal(expected) + + # Just make sure these work + for aggop in [agg.any_value, agg.first, agg.last, agg.count]: + A.reduce_rowwise(aggop).new() + A.reduce_columnwise(aggop).new() + A.reduce_scalar(aggop).new() + for aggop in [agg.first_index, agg.last_index]: + A.reduce_rowwise(aggop).new() + A.reduce_columnwise(aggop).new() + + np_dtype = np.dtype("(3,)uint16") + udt = dtypes.register_anonymous(np_dtype, "has_subdtype") + A = Matrix.new(udt, nrows=2, ncols=2) + A[:, :] = (1, 2, 3) + rows, cols, values = A.to_values() + assert_array_equal(values, np.array([[1, 2, 3]] * 4)) + result = Matrix.from_values(rows, cols, values) + assert A.isequal(result) + assert result.isequal(A) + result = Matrix.from_values(rows, cols, values, dtype=udt) + assert result.isequal(A) + info = A.ss.export() + result = A.ss.import_any(**info) + assert result.isequal(A) + info = A.ss.export("coor") + result = A.ss.import_any(**info) + assert result.isequal(A) diff --git a/grblas/tests/test_numpyops.py b/grblas/tests/test_numpyops.py index f0a3762c5..f5bea3f6d 100644 --- a/grblas/tests/test_numpyops.py +++ b/grblas/tests/test_numpyops.py @@ -53,15 +53,15 @@ def test_npunary(): data.append( [Vector.from_values(L, L, dtype="FC64"), np.array(L, dtype=np.complex128)], ) - blacklist = {"BOOL": {"negative"}, "FC64": {"ceil", "floor", "trunc"}} + blocklist = {"BOOL": {"negative", "sign"}, "FC64": {"ceil", "floor", "trunc"}} isclose = grblas.binary.isclose(1e-7, 0) for gb_input, np_input in data: for unary_name in sorted(npunary._unary_names): op = getattr(npunary, unary_name) - if gb_input.dtype.name not in op.types or unary_name in blacklist.get( + if gb_input.dtype not in op.types or unary_name in blocklist.get( gb_input.dtype.name, () ): - continue # pragma: no cover + continue if gb_input.dtype.name.startswith("FC"): # There are some nasty branch cuts as 1 gb_input = gb_input.dup() @@ -129,15 +129,15 @@ def test_npbinary(): [np.array(values1, dtype=np.complex128), np.array(values2, dtype=np.complex128)], ], ) - blacklist = { + blocklist = { "FP64": {"floor_divide"}, # numba/numpy difference for 1.0 / 0.0 - "BOOL": {"subtract"}, # not supported by numpy + "BOOL": {"gcd", "lcm", "subtract"}, # not supported by numpy } isclose = grblas.binary.isclose(1e-7, 0) for (gb_left, gb_right), (np_left, np_right) in data: for binary_name in sorted(npbinary._binary_names): op = getattr(npbinary, binary_name) - if gb_left.dtype.name not in op.types or binary_name in blacklist.get( + if gb_left.dtype not in op.types or binary_name in blocklist.get( gb_left.dtype.name, () ): continue @@ -213,18 +213,18 @@ def test_npmonoid(): # ], # ] # ) - blacklist = {} - reduction_blacklist = { + blocklist = {} + reduction_blocklist = { "BOOL": {"add"}, } for (gb_left, gb_right), (np_left, np_right) in data: for binary_name in sorted(npmonoid._monoid_identities): op = getattr(npmonoid, binary_name) assert len(op.types) > 0, op.name - if gb_left.dtype.name not in op.types or binary_name in blacklist.get( + if gb_left.dtype not in op.types or binary_name in blocklist.get( gb_left.dtype.name, () ): - continue # pragma: no cover + continue with np.errstate(divide="ignore", over="ignore", under="ignore", invalid="ignore"): gb_result = gb_left.ewise_mult(gb_right, op).new() np_result = getattr(np, binary_name)(np_left, np_right) @@ -243,7 +243,7 @@ def test_npmonoid(): assert compare # numpy reductions don't have dtype-dependent identities, so results sometimes differ - if binary_name in reduction_blacklist.get(gb_left.dtype.name, ()): + if binary_name in reduction_blocklist.get(gb_left.dtype.name, ()): continue gb_result = gb_left.reduce(op).new() diff --git a/grblas/tests/test_op.py b/grblas/tests/test_op.py index c01bd499e..8978b818f 100644 --- a/grblas/tests/test_op.py +++ b/grblas/tests/test_op.py @@ -5,9 +5,13 @@ import grblas as gb from grblas import agg, binary, dtypes, lib, monoid, op, operator, semiring, unary +from grblas.dtypes import BOOL, FP32, FP64, INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32, UINT64 from grblas.exceptions import DomainMismatch, UdfParseError from grblas.operator import BinaryOp, Monoid, Semiring, UnaryOp, get_semiring +if dtypes._supports_complex: + from grblas.dtypes import FC32, FC64 + from grblas import Matrix, Vector # isort:skip @@ -32,15 +36,15 @@ def test_op_repr(): def test_unaryop(): assert unary.ainv["INT32"].gb_obj == lib.GrB_AINV_INT32 assert unary.ainv[dtypes.UINT16].gb_obj == lib.GrB_AINV_UINT16 - assert orig_types(unary.positioni) == {"INT32", "INT64"} - assert orig_types(unary.positionj1) == {"INT32", "INT64"} + assert orig_types(unary.positioni) == {INT32, INT64} + assert orig_types(unary.positionj1) == {INT32, INT64} def test_binaryop(): assert binary.plus["INT32"].gb_obj == lib.GrB_PLUS_INT32 assert binary.plus[dtypes.UINT16].gb_obj == lib.GrB_PLUS_UINT16 - assert orig_types(binary.firsti) == {"INT32", "INT64"} - assert orig_types(binary.secondj1) == {"INT32", "INT64"} + assert orig_types(binary.firsti) == {INT32, INT64} + assert orig_types(binary.secondj1) == {INT32, INT64} def test_monoid(): @@ -51,14 +55,14 @@ def test_monoid(): def test_semiring(): assert semiring.min_plus["INT32"].gb_obj == lib.GrB_MIN_PLUS_SEMIRING_INT32 assert semiring.min_plus[dtypes.UINT16].gb_obj == lib.GrB_MIN_PLUS_SEMIRING_UINT16 - assert orig_types(semiring.min_firsti) == {"INT32", "INT64"} + assert orig_types(semiring.min_firsti) == {INT32, INT64} def test_agg(): assert repr(agg.count) == "agg.count" assert repr(agg.count["INT32"]) == "agg.count[INT32]" assert "INT64" in agg.sum_of_inverses - assert agg.sum_of_inverses["INT64"].return_type == "FP64" + assert agg.sum_of_inverses["INT64"].return_type == FP64 assert "BOOL" not in agg.sum_of_inverses with pytest.raises(KeyError, match="BOOL"): agg.sum_of_inverses["BOOL"] @@ -111,39 +115,44 @@ def test_get_typed_op(): ) with pytest.raises(ValueError, match="Unknown binary or aggregator"): operator.get_typed_op("bad_op_name", dtypes.INT64, kind="binary|aggregator") + with pytest.raises(Exception): + # get_typed_op expects dtypes to already be dtypes + operator.get_typed_op(binary.plus, dtypes.INT64, "bad dtype") def test_unaryop_udf(): def plus_one(x): - return x + 1 + return x + 1 # pragma: no cover UnaryOp.register_new("plus_one", plus_one) assert hasattr(unary, "plus_one") + assert unary.plus_one.orig_func is plus_one + assert unary.plus_one[int].orig_func is plus_one comp_set = { - "INT8", - "INT16", - "INT32", - "INT64", - "UINT8", - "UINT16", - "UINT32", - "UINT64", - "FP32", - "FP64", - "BOOL", + INT8, + INT16, + INT32, + INT64, + UINT8, + UINT16, + UINT32, + UINT64, + FP32, + FP64, + BOOL, } if dtypes._supports_complex: - comp_set.update({"FC32", "FC64"}) + comp_set.update({FC32, FC64}) assert set(unary.plus_one.types) == comp_set v = Vector.from_values([0, 1, 3], [1, 2, -4], dtype=dtypes.INT32) v << v.apply(unary.plus_one) result = Vector.from_values([0, 1, 3], [2, 3, -3], dtype=dtypes.INT32) assert v.isequal(result) assert "INT8" in unary.plus_one - assert "INT8" in unary.plus_one.types + assert INT8 in unary.plus_one.types del unary.plus_one["INT8"] assert "INT8" not in unary.plus_one - assert "INT8" not in unary.plus_one.types + assert INT8 not in unary.plus_one.types with pytest.raises(TypeError, match="UDF argument must be a function"): UnaryOp.register_new("bad", object()) assert not hasattr(unary, "bad") @@ -155,7 +164,7 @@ def plus_one(x): def test_unaryop_parameterized(): def plus_x(x=0): def inner(val): - return val + x + return val + x # pragma: no cover return inner @@ -179,7 +188,7 @@ def inner(val): def test_binaryop_parameterized(): def plus_plus_x(x=0): def inner(left, right): - return left + right + x + return left + right + x # pragma: no cover return inner @@ -225,7 +234,7 @@ def inner(left, right): BinaryOp.register_new("bad", lambda x, y: v) def my_add(x, y): - return x + y + return x + y # pragma: no cover op = BinaryOp.register_anonymous(my_add) assert op.name == "my_add" @@ -235,7 +244,7 @@ def my_add(x, y): def test_monoid_parameterized(): def plus_plus_x(x=0): def inner(left, right): - return left + right + x + return left + right + x # pragma: no cover return inner @@ -284,7 +293,7 @@ def plus_plus_x_identity(x=0): # identity may be a value def logaddexp(base): def inner(x, y): - return np.log(base**x + base**y) / np.log(base) + return np.log(base**x + base**y) / np.log(base) # pragma: no cover return inner @@ -301,7 +310,7 @@ def inner(x, y): def plus_times_x(x=0): def inner(left, right): - return (left + right) * x + return (left + right) * x # pragma: no cover return inner @@ -320,7 +329,7 @@ def bad_identity(x=0): def test_semiring_parameterized(): def plus_plus_x(x=0): def inner(left, right): - return left + right + x + return left + right + x # pragma: no cover return inner @@ -420,7 +429,7 @@ def other_identity(y=0): def plus_x(x=0): def inner(y): - return x + y + return x + y # pragma: no cover return inner @@ -447,22 +456,22 @@ def inner(y): def test_unaryop_udf_bool_result(): # numba has trouble compiling this, but we have a work-around def is_positive(x): - return x > 0 + return x > 0 # pragma: no cover UnaryOp.register_new("is_positive", is_positive) assert hasattr(unary, "is_positive") assert set(unary.is_positive.types) == { - "INT8", - "INT16", - "INT32", - "INT64", - "UINT8", - "UINT16", - "UINT32", - "UINT64", - "FP32", - "FP64", - "BOOL", + INT8, + INT16, + INT32, + INT64, + UINT8, + UINT16, + UINT32, + UINT64, + FP32, + FP64, + BOOL, } v = Vector.from_values([0, 1, 3], [1, 2, -4], dtype=dtypes.INT32) w = v.apply(unary.is_positive).new() @@ -472,24 +481,25 @@ def is_positive(x): def test_binaryop_udf(): def times_minus_sum(x, y): - return x * y - (x + y) + return x * y - (x + y) # pragma: no cover BinaryOp.register_new("bin_test_func", times_minus_sum) assert hasattr(binary, "bin_test_func") comp_set = { - "INT8", - "INT16", - "INT32", - "INT64", - "UINT8", - "UINT16", - "UINT32", - "UINT64", - "FP32", - "FP64", + BOOL, # goes to INT64 + INT8, + INT16, + INT32, + INT64, + UINT8, + UINT16, + UINT32, + UINT64, + FP32, + FP64, } if dtypes._supports_complex: - comp_set.update({"FC32", "FC64"}) + comp_set.update({FC32, FC64}) assert set(binary.bin_test_func.types) == comp_set v1 = Vector.from_values([0, 1, 3], [1, 2, -4], dtype=dtypes.INT32) v2 = Vector.from_values([0, 2, 3], [2, 3, 7], dtype=dtypes.INT32) @@ -500,22 +510,22 @@ def times_minus_sum(x, y): def test_monoid_udf(): def plus_plus_one(x, y): - return x + y + 1 + return x + y + 1 # pragma: no cover BinaryOp.register_new("plus_plus_one", plus_plus_one) Monoid.register_new("plus_plus_one", binary.plus_plus_one, -1) assert hasattr(monoid, "plus_plus_one") assert set(monoid.plus_plus_one.types) == { - "INT8", - "INT16", - "INT32", - "INT64", - "UINT8", - "UINT16", - "UINT32", - "UINT64", - "FP32", - "FP64", + INT8, + INT16, + INT32, + INT64, + UINT8, + UINT16, + UINT32, + UINT64, + FP32, + FP64, } v1 = Vector.from_values([0, 1, 3], [1, 2, -4], dtype=dtypes.INT32) v2 = Vector.from_values([0, 2, 3], [2, 3, 7], dtype=dtypes.INT32) @@ -531,7 +541,7 @@ def plus_plus_one(x, y): def test_semiring_udf(): def plus_plus_two(x, y): - return x + y + 2 + return x + y + 2 # pragma: no cover BinaryOp.register_new("plus_plus_two", plus_plus_two) Semiring.register_new("extra_twos", monoid.plus, binary.plus_plus_two) @@ -564,27 +574,27 @@ def test_binary_updates(): @pytest.mark.slow def test_nested_names(): def plus_three(x): - return x + 3 + return x + 3 # pragma: no cover UnaryOp.register_new("incrementers.plus_three", plus_three) assert hasattr(unary, "incrementers") assert type(unary.incrementers) is operator.OpPath assert hasattr(unary.incrementers, "plus_three") comp_set = { - "INT8", - "INT16", - "INT32", - "INT64", - "UINT8", - "UINT16", - "UINT32", - "UINT64", - "FP32", - "FP64", - "BOOL", + INT8, + INT16, + INT32, + INT64, + UINT8, + UINT16, + UINT32, + UINT64, + FP32, + FP64, + BOOL, } if dtypes._supports_complex: - comp_set.update({"FC32", "FC64"}) + comp_set.update({FC32, FC64}) assert set(unary.incrementers.plus_three.types) == comp_set v = Vector.from_values([0, 1, 3], [1, 2, -4], dtype=dtypes.INT32) @@ -593,7 +603,7 @@ def plus_three(x): assert v.isequal(result), v def plus_four(x): - return x + 4 + return x + 4 # pragma: no cover UnaryOp.register_new("incrementers.plus_four", plus_four) assert hasattr(unary.incrementers, "plus_four") @@ -679,7 +689,10 @@ def test_binaryop_attributes(): assert binary.numpy.add.monoid is monoid.numpy.add assert binary.numpy.subtract.monoid is None - op = BinaryOp.register_anonymous(lambda x, y: x + y, name="plus") + def plus(x, y): + return x + y # pragma: no cover + + op = BinaryOp.register_anonymous(plus, name="plus") assert op.monoid is None assert op[int].monoid is None @@ -716,7 +729,10 @@ def test_monoid_attributes(): assert monoid.numpy.add.binaryop is binary.numpy.add assert monoid.numpy.add.identities == {typ: 0 for typ in monoid.numpy.add.types} - binop = BinaryOp.register_anonymous(lambda x, y: x + y, name="plus") + def plus(x, y): # pragma: no cover + return x + y + + binop = BinaryOp.register_anonymous(plus, name="plus") op = Monoid.register_anonymous(binop, 0, name="plus") assert op.binaryop is binop assert op[int].binaryop is binop[int] @@ -749,7 +765,10 @@ def test_semiring_attributes(): assert semiring.numpy.add_subtract.monoid is monoid.numpy.add assert semiring.numpy.add_subtract.binaryop is binary.numpy.subtract - binop = BinaryOp.register_anonymous(lambda x, y: x + y, name="plus") + def plus(x, y): + return x + y # pragma: no cover + + binop = BinaryOp.register_anonymous(plus, name="plus") mymonoid = Monoid.register_anonymous(binop, 0, name="plus") op = Semiring.register_anonymous(mymonoid, binop, name="plus_plus") assert op.binaryop is binop @@ -813,7 +832,7 @@ def test_get_semiring(): get_semiring(binary.plus, monoid.times) def myplus(x, y): - return x + y + return x + y # pragma: no cover binop = BinaryOp.register_anonymous(myplus, name="myplus") st = get_semiring(monoid.plus, binop) @@ -956,10 +975,10 @@ def test_from_string(): @pytest.mark.slow def test_lazy_op(): - UnaryOp.register_new("lazy", lambda x: x, lazy=True) + UnaryOp.register_new("lazy", lambda x: x, lazy=True) # pragma: no branch assert isinstance(op.lazy, UnaryOp) assert isinstance(unary.lazy, UnaryOp) - BinaryOp.register_new("lazy", lambda x, y: x + y, lazy=True) + BinaryOp.register_new("lazy", lambda x, y: x + y, lazy=True) # pragma: no branch Monoid.register_new("lazy", "lazy", 0, lazy=True) assert isinstance(monoid.lazy, Monoid) assert isinstance(binary.lazy, BinaryOp) @@ -971,9 +990,9 @@ def test_lazy_op(): Semiring.register_new("lazy_lazy", monoid.lazy, binary.lazy, lazy=True) assert isinstance(semiring.lazy_lazy, Semiring) # numpy - UnaryOp.register_new("numpy.lazy", lambda x: x, lazy=True) + UnaryOp.register_new("numpy.lazy", lambda x: x, lazy=True) # pragma: no branch assert isinstance(unary.numpy.lazy, UnaryOp) - BinaryOp.register_new("numpy.lazy", lambda x, y: x + y, lazy=True) + BinaryOp.register_new("numpy.lazy", lambda x, y: x + y, lazy=True) # pragma: no branch Monoid.register_new("numpy.lazy", "numpy.lazy", 0, lazy=True) assert isinstance(monoid.numpy.lazy, Monoid) assert isinstance(binary.numpy.lazy, BinaryOp) @@ -986,7 +1005,7 @@ def test_lazy_op(): Semiring.register_new("numpy.lazy_lazy", monoid.numpy.lazy, binary.numpy.lazy, lazy=True) assert isinstance(semiring.numpy.lazy_lazy, Semiring) # misc - UnaryOp.register_new("misc.lazy", lambda x: x, lazy=True) + UnaryOp.register_new("misc.lazy", lambda x: x, lazy=True) # pragma: no branch assert isinstance(unary.misc.lazy, UnaryOp) with pytest.raises(AttributeError): unary.misc.bad @@ -1018,6 +1037,109 @@ def test_positional(): assert not semiring.any_second[int].is_positional +def test_udt(): + record_dtype = np.dtype([("x", np.bool_), ("y", np.float64)], align=True) + udt = dtypes.register_new("TestUDT", record_dtype) + assert not udt._is_anonymous + v = Vector.new(udt, size=3) + w = Vector.new(udt, size=3) + v[:] = 0 + w[:] = 1 + + def _udt_identity(val): + return val # pragma: no cover + + udt_identity = UnaryOp.register_new("udt_identity", _udt_identity, is_udt=True) + assert udt in udt_identity + assert udt in binary.eq + result = v.apply(udt_identity).new() + assert result.isequal(v) + assert dtypes.UINT8 in udt_identity + assert udt in udt_identity + assert int in udt_identity + assert operator.get_typed_op(udt_identity, udt) is udt_identity[udt] + with pytest.raises(ValueError): + "badname" in binary.eq + with pytest.raises(ValueError): + "badname" in udt_identity + + def _udt_getx(val): + return val["x"] # pragma: no cover + + udt_getx = UnaryOp.register_anonymous(_udt_getx, "udt_getx", is_udt=True) + assert udt in udt_getx + result = v.apply(udt_getx).new() + expected = Vector.from_values([0, 1, 2], 0) + assert result.isequal(expected) + + def _udt_first(x, y): + return x # pragma: no cover + + udt_first = BinaryOp.register_anonymous(_udt_first, "udt_first", is_udt=True) + 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) + 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 + assert udt_first[dtypes.BOOL, udt].return_type == dtypes.BOOL + udt_dup = dtypes.register_anonymous(record_dtype) + assert udt_first[udt, udt_dup].return_type == udt + # assert udt_first[udt_dup, udt].return_type == udt ? + + udt_any = Monoid.register_new("udt_any", udt_first, (0, 0)) + assert udt in udt_any + assert (udt, udt) in udt_any + assert (udt, dtypes.INT8) not in udt_any + assert operator.get_typed_op(udt_any, udt) is udt_any[udt] + assert udt_any(v | w).new().isequal(v) + + udt_semiring = Semiring.register_new("udt_semiring", udt_any, udt_first) + assert udt in udt_semiring + assert operator.get_typed_op(udt_semiring, udt) is udt_semiring[udt] + assert udt_semiring(v @ v).new() == (0, 0) + + result = v.apply(gb.unary.identity).new() + assert result.isequal(v) + result = v.apply(gb.unary.one).new() + assert result.dtype == dtypes.INT64 + expected = Vector.new(int, size=v.size) + expected(result.S) << 1 + assert result.isequal(expected) + result = v.apply(gb.unary.positioni).new() + expected = expected.apply(gb.unary.positioni).new() + assert result.isequal(expected) + + class BreakCompile: + pass + + def badfunc(x): # pragma: no cover + return BreakCompile(x) + + badunary = UnaryOp.register_anonymous(badfunc, is_udt=True) + assert udt not in badunary + assert int not in badunary + + def badfunc(x, y): # pragma: no cover + return BreakCompile(x) + + badbinary = BinaryOp.register_anonymous(badfunc, is_udt=True) + assert udt not in badbinary + assert int not in badbinary + + assert binary.first[udt].return_type is udt + assert binary.first[udt].commutes_to is binary.second[udt] + assert semiring.any_firsti[int].commutes_to is semiring.any_secondj[int] + assert semiring.any_firsti[udt].commutes_to is semiring.any_secondj[udt] + + assert binary.second[udt].type is udt + assert binary.second[udt].type2 is udt + assert binary.second[udt, dtypes.INT8].type is udt + assert binary.second[udt, dtypes.INT8].type2 is dtypes.INT8 + assert monoid.any[udt].type2 is udt + + def test_dir(): for mod in [unary, binary, monoid, semiring, op]: assert not set(mod._delayed) - set(dir(mod)) diff --git a/grblas/tests/test_operator_types.py b/grblas/tests/test_operator_types.py index 7aa26bbc5..e825b90b8 100644 --- a/grblas/tests/test_operator_types.py +++ b/grblas/tests/test_operator_types.py @@ -2,18 +2,25 @@ from collections import defaultdict from grblas import binary, dtypes, monoid, operator, semiring, unary +from grblas.dtypes import BOOL, FP32, FP64, INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32, UINT64 -BOOL = frozenset({"BOOL"}) -UINT = frozenset({"UINT8", "UINT16", "UINT32", "UINT64"}) -INT = frozenset({"INT8", "INT16", "INT32", "INT64", "UINT8", "UINT16", "UINT32", "UINT64"}) +if dtypes._supports_complex: + from grblas.dtypes import FC32, FC64 +else: + FC32 = "FC32" + FC64 = "FC64" + +BOOL = frozenset({BOOL}) +UINT = frozenset({UINT8, UINT16, UINT32, UINT64}) +INT = frozenset({INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32, UINT64}) BOOLINT = frozenset(BOOL | INT) -FP = frozenset({"FP32", "FP64"}) +FP = frozenset({FP32, FP64}) FPINT = frozenset(FP | INT) NOFC = frozenset(BOOL | FPINT) -FC = frozenset({"FC32", "FC64"}) +FC = frozenset({FC32, FC64}) FCFP = frozenset(FC | FP) ALL = frozenset(NOFC | FC) -POS = frozenset({"INT32", "INT64"}) +POS = frozenset({INT32, INT64}) NOBOOL = frozenset(ALL - BOOL) # fmt: off @@ -158,6 +165,9 @@ "myplus", "plus_myplus", "plus_numpy_copysign", + "udt_identity", + "udt_any", + "udt_semiring", "any_any", # numpy-graphblas commutation (can we clean this up?) "band_land", diff --git a/grblas/tests/test_pickle.py b/grblas/tests/test_pickle.py index e986ecfc0..e7d501320 100644 --- a/grblas/tests/test_pickle.py +++ b/grblas/tests/test_pickle.py @@ -1,25 +1,26 @@ import os import pickle +import numpy as np import pytest import grblas as gb def unarypickle(x): - return x + 1 + return x + 1 # pragma: no cover def binarypickle(x, y): - return x + y + return x + y # pragma: no cover def unaryanon(x): - return x + 2 + return x + 2 # pragma: no cover def binaryanon(x, y): - return x + y + return x + y # pragma: no cover @pytest.mark.slow @@ -182,28 +183,28 @@ def check_values(d): def unarypickle_par(x): def inner(y): - return x + y + return x + y # pragma: no cover return inner def binarypickle_par(z): def inner(x, y): - return x + y + z + return x + y + z # pragma: no cover return inner def unaryanon_par(x): def inner(y): - return y + x + return y + x # pragma: no cover return inner def binaryanon_par(z): def inner(x, y): - return x + y + z + return x + y + z # pragma: no cover return inner @@ -273,3 +274,42 @@ def test_deserialize_parameterized(): # Again! with open(os.path.join(thisdir, "pickle2.pkl"), "rb") as f: pickle.load(f) # TODO: check results + + +def test_udt(): + record_dtype = np.dtype([("x", np.bool_), ("y", np.int64)], align=True) + udt = gb.dtypes.register_new("PickleUDT", record_dtype) + assert not udt._is_anonymous + assert pickle.loads(pickle.dumps(udt)) is udt + + np_dtype = np.dtype("(3,)uint16") + udt2 = gb.dtypes.register_anonymous(np_dtype, "pickling") + assert udt2._is_anonymous + assert pickle.loads(pickle.dumps(udt2)).np_type == udt2.np_type + + thisdir = os.path.dirname(__file__) + with open(os.path.join(thisdir, "pickle3.pkl"), "rb") as f: + d = pickle.load(f) + udt3 = d["PickledUDT"] + v = d["v"] + assert udt3.name == "PickledUDT" + assert udt3 is gb.dtypes.PickledUDT + assert v.dtype == udt3 + expected = gb.Vector.new(udt3, size=2) + expected[0] = (False, 1) + expected[1] = (True, 3) + assert expected.isequal(v) + + udt4 = d["pickled_subdtype"] + A = d["A"] + assert udt4.name == "pickled_subdtype" + assert not hasattr(gb.dtypes, udt4.name) + assert A.dtype == udt4 + expected = gb.Matrix.new(udt4, nrows=1, ncols=2) + expected[0, 0] = (1, 2) + expected[0, 1] = (3, 4) + assert expected.isequal(A) + + any_udt = d["any[udt]"] + assert any_udt is gb.binary.any[udt3] + assert pickle.loads(pickle.dumps(gb.binary.first[udt, int])) is gb.binary.first[udt, int] diff --git a/grblas/tests/test_scalar.py b/grblas/tests/test_scalar.py index e68c2a191..358dccb1f 100644 --- a/grblas/tests/test_scalar.py +++ b/grblas/tests/test_scalar.py @@ -44,23 +44,35 @@ def test_dup(s): s4 = Scalar.from_value(-2.5, dtype=dtypes.FP64) s_empty = Scalar.new(dtypes.FP64) s_unempty = Scalar.from_value(0.0) + if s_empty.is_cscalar: + # NumPy wraps around + uint_data = [ + ("UINT8", 2**8 - 2), + ("UINT16", 2**16 - 2), + ("UINT32", 2**32 - 2), + ("UINT64", 2**64 - 2), + ] + else: + # SuiteSparse clips + uint_data = [ + ("UINT8", 0), + ("UINT16", 0), + ("UINT32", 0), + ("UINT64", 0), + ] for dtype, val in [ ("INT8", -2), ("INT16", -2), ("INT32", -2), - ("UINT8", 2**8 - 2), - ("UINT16", 2**16 - 2), - ("UINT32", 2**32 - 2), - ("UINT64", 2**64 - 2), ("BOOL", True), ("FP32", -2.5), - ]: - s5 = s4.dup(dtype=dtype) + ] + uint_data: + s5 = s4.dup(dtype=dtype, name="s5") assert s5.dtype == dtype and s5.value == val - s6 = s_empty.dup(dtype=dtype) + s6 = s_empty.dup(dtype=dtype, name="s6") assert s6.is_empty assert compute(s6.value) is None - s7 = s_unempty.dup(dtype=dtype) + s7 = s_unempty.dup(dtype=dtype, name="s7") assert not s7.is_empty assert compute(s7.value) is not None @@ -258,7 +270,7 @@ def test_neg(): ): s = Scalar.from_value(1, dtype=dtype) empty = Scalar.new(dtype) - if dtype.name == "BOOL" or dtype.name.startswith("U"): + if dtype.name == "BOOL" or dtype.name.startswith("U") or dtype._is_udt: with pytest.raises(TypeError, match="The negative operator, `-`, is not supported"): -s with pytest.raises(TypeError, match="The negative operator, `-`, is not supported"): diff --git a/grblas/tests/test_vector.py b/grblas/tests/test_vector.py index 198cdd29b..30965d3a4 100644 --- a/grblas/tests/test_vector.py +++ b/grblas/tests/test_vector.py @@ -1692,6 +1692,117 @@ def test_delete_via_scalar(v): assert v.nvals == 0 +def test_udt(): + record_dtype = np.dtype([("x", np.bool_), ("y", np.float64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "VectorUDT") + assert udt._is_anonymous + v = Vector.new(udt, size=3) + a = np.zeros(1, dtype=record_dtype) + v[0] = a[0] + expected = Vector.from_values([0], a, size=3, dtype=udt) + assert v.isequal(expected) + v[1] = (1, 2) + expected = Vector.from_values( + [0, 1], np.array([(0, 0), (1, 2)], dtype=record_dtype), size=3, dtype=udt + ) + assert v.isequal(expected) + v[:] = 0 + zeros = Vector.from_values([0, 1, 2], np.zeros(3, dtype=record_dtype), dtype=udt) + assert v.isequal(zeros) + v(v.S)[:] = 1 + ones = Vector.from_values([0, 1, 2], np.ones(3, dtype=record_dtype), dtype=udt) + assert v.isequal(ones) + v[:](v.S) << 0 + assert v.isequal(zeros) + s = Scalar.new(udt) + s.value = (1, 1) + v(v.S)[:] << s + assert v.isequal(ones) + s.value = 0 + v[:](v.S) << s + assert v.isequal(zeros) + t = v.reduce(monoid.any, allow_empty=True).new() + assert s == t + t = v.reduce(monoid.any, allow_empty=False).new() + assert s == t + v << binary.first(1, v) + assert v.isequal(ones) + v << binary.first(s, v) + assert v.isequal(zeros) + v << binary.second(v, (1, 1)) + assert v.isequal(ones) + v << binary.second(v, s) + assert v.isequal(zeros) + assert v[0].new() == s + assert v[:].new().isequal(v) + indices, values = v.to_values() + assert v.isequal(Vector.from_values(indices, values)) + result = unary.positioni(v).new() + expected = Vector.from_values([0, 1, 2], [0, 1, 2]) + assert result.isequal(expected) + result = binary.firsti(v & v).new() + assert result.isequal(expected) + result = semiring.min_firsti(v @ v).new() + assert result == 0 + assert v.reduce(agg.any_value).new() == s + assert v.reduce(agg.first).new() == s + assert v.reduce(agg.last).new() == s + assert v.reduce(agg.exists).new() == 1 + assert v.reduce(agg.first_index).new() == 0 + assert v.reduce(agg.last_index).new() == 2 + assert v.reduce(agg.count).new() == 3 + info = v.ss.export() + result = Vector.ss.import_any(**info) + assert result.isequal(v) + for aggop in [agg.any_value, agg.first, agg.last, agg.count, agg.first_index, agg.last_index]: + v.reduce(aggop).new() + + # arrays as dtypes! + np_dtype = np.dtype("(3,)uint16") + udt2 = dtypes.register_anonymous(np_dtype, "has_subdtype") + s = Scalar.new(udt2) + s.value = [0, 0, 0] + + v = Vector.new(udt2, size=2) + v[0] = 0 + w = Vector.new(udt2, size=2) + w[0] = (0, 0, 0) + assert v.isequal(w) + assert v.ewise_mult(w, binary.eq).new().reduce(monoid.land).new() + assert not v.ewise_mult(w, binary.ne).new().reduce(monoid.land).new() + + w[0] = np.array([0, 0, 1], dtype=np.uint8) + assert not v.isequal(w) + assert not v.ewise_mult(w, binary.eq).new().reduce(monoid.land).new() + assert v.ewise_mult(w, binary.ne).new().reduce(monoid.land).new() + v[:] = 1 + w[:] = np.array([1, 1, 1], dtype=np.uint8) + assert v.isequal(w) + + indices, values = v.to_values() + assert_array_equal(values, np.ones((2, 3), dtype=np.uint16)) + assert v.isequal(Vector.from_values(indices, values, dtype=v.dtype)) + assert v.isequal(Vector.from_values(indices, values)) + info = v.ss.export() + result = Vector.ss.import_any(**info) + assert result.isequal(v) + + s = v.reduce(monoid.any).new() + assert (s.value == [1, 1, 1]).all() + assert (s.value == 1).all() + assert s == 1 + assert s != 0 + assert s == (1, 1, 1) + assert s == [1, 1, 1] + assert s == s + t = Scalar.from_value(1) + assert s == t + assert t == s + # Just make sure these work + for aggop in [agg.any_value, agg.first, agg.last, agg.count, agg.first_index, agg.last_index]: + v.reduce(aggop).new() + + def test_infix_outer(): v = Vector.new(int, 2) v += 1 diff --git a/grblas/unary/numpy.py b/grblas/unary/numpy.py index 0cc87289f..9a910cf1d 100644 --- a/grblas/unary/numpy.py +++ b/grblas/unary/numpy.py @@ -106,6 +106,7 @@ _unary_names.update({"conj", "conjugate"}) _numpy_to_graphblas["conj"] = "conj" _numpy_to_graphblas["conjugate"] = "conj" +# _graphblas_to_numpy = {val: key for key, val in _numpy_to_graphblas.items()} # Soon... _operator._STANDARD_OPERATOR_NAMES.update(f"unary.numpy.{name}" for name in _unary_names) __all__ = list(_unary_names) diff --git a/grblas/utils.py b/grblas/utils.py index 7b19c7293..c22dace4e 100644 --- a/grblas/utils.py +++ b/grblas/utils.py @@ -67,14 +67,20 @@ def ints_to_numpy_buffer(array, dtype, *, name="array", copy=False, ownable=Fals return array +def _get_subdtype(dtype): + while dtype.subdtype is not None: + dtype = dtype.subdtype[0] + return dtype + + def values_to_numpy_buffer(array, dtype=None, *, copy=False, ownable=False, order="C"): if dtype is not None: dtype = lookup_dtype(dtype) - array = np.array(array, dtype.np_type, copy=copy, order=order) + array = np.array(array, _get_subdtype(dtype.np_type), copy=copy, order=order) else: is_input_np = isinstance(array, np.ndarray) array = np.array(array, copy=copy, order=order) - if array.dtype == object: + if array.dtype.hasobject: raise ValueError("object dtype for values is not allowed") if not is_input_np and array.dtype == np.int32: # pragma: no cover # fix for win64 numpy handling of ints @@ -127,8 +133,9 @@ def __init__(self, array=None, dtype=_INDEX, *, size=None, name=None): if size is not None: self.array = np.empty(size, dtype=dtype.np_type) else: - self.array = np.array(array, dtype=dtype.np_type, copy=False, order="C") - self._carg = ffi.cast(f"{dtype.c_type}*", ffi.from_buffer(self.array)) + self.array = np.array(array, dtype=_get_subdtype(dtype.np_type), copy=False, order="C") + c_type = dtype.c_type if dtype._is_udt else f"{dtype.c_type}*" + self._carg = ffi.cast(c_type, ffi.from_buffer(self.array)) self.dtype = dtype self._name = name diff --git a/grblas/vector.py b/grblas/vector.py index fc1353b1a..18309e71e 100644 --- a/grblas/vector.py +++ b/grblas/vector.py @@ -48,7 +48,7 @@ def __del__(self): if parent is not None: return gb_obj = getattr(self, "gb_obj", None) - if gb_obj is not None: + if gb_obj is not None and lib is not None: # it's difficult/dangerous to record the call, b/c `self.name` may not exist check_status(lib.GrB_Vector_free(gb_obj), self) @@ -163,12 +163,12 @@ def isequal(self, other, *, check_dtype=False): return False if check_dtype: # dtypes are equivalent, so not need to unify - common_dtype = self.dtype + op = binary.eq[self.dtype] else: - common_dtype = unify(self.dtype, other.dtype) + op = get_typed_op(binary.eq, self.dtype, other.dtype, kind="binary") matches = Vector.new(bool, self._size, name="v_isequal") - matches << self.ewise_mult(other, binary.eq[common_dtype]) + matches << self.ewise_mult(other, op) # ewise_mult performs intersection, so nvals will indicate mismatched empty values if matches._nvals != self._nvals: return False @@ -244,9 +244,8 @@ def to_values(self, dtype=None): n = ffi_new("GrB_Index*") scalar = Scalar(n, _INDEX, name="s_nvals", is_cscalar=True, empty=True) scalar.value = nvals - call( - f"GrB_Vector_extractTuples_{self.dtype.name}", [indices, values, _Pointer(scalar), self] - ) + dtype_name = "UDT" if self.dtype._is_udt else self.dtype.name + call(f"GrB_Vector_extractTuples_{dtype_name}", [indices, values, _Pointer(scalar), self]) values = values.array if dtype is not None: dtype = lookup_dtype(dtype) @@ -261,7 +260,7 @@ def build(self, indices, values, *, dup_op=None, clear=False, size=None): # TODO: accept `dtype` keyword to match the dtype of `values`? indices = ints_to_numpy_buffer(indices, np.uint64, name="indices") values, dtype = values_to_numpy_buffer(values, self.dtype) - n = values.size + n = values.shape[0] if indices.size != n: raise ValueError( f"`indices` and `values` lengths must match: {indices.size} != {values.size}" @@ -275,7 +274,11 @@ def build(self, indices, values, *, dup_op=None, clear=False, size=None): dup_op_given = dup_op is not None if not dup_op_given: - dup_op = binary.plus + if not self.dtype._is_udt: + dup_op = binary.plus + else: + dup_op = binary.any + # SS:SuiteSparse-specific: we could use NULL for dup_op dup_op = get_typed_op(dup_op, self.dtype, kind="binary") if dup_op.opclass == "Monoid": dup_op = dup_op.binaryop @@ -284,8 +287,9 @@ def build(self, indices, values, *, dup_op=None, clear=False, size=None): indices = _CArray(indices) values = _CArray(values, self.dtype) + dtype_name = "UDT" if self.dtype._is_udt else self.dtype.name call( - f"GrB_Vector_build_{self.dtype.name}", + f"GrB_Vector_build_{dtype_name}", [self, indices, values, _as_scalar(n, _INDEX, is_cscalar=True), dup_op], ) @@ -352,14 +356,17 @@ def from_values(cls, indices, values, dtype=None, *, size=None, dup_op=None, nam values may be a scalar, in which case duplicate indices are ignored. """ indices = ints_to_numpy_buffer(indices, np.uint64, name="indices") - values, dtype = values_to_numpy_buffer(values, dtype) + values, new_dtype = values_to_numpy_buffer(values, dtype) # Compute size if not provided if size is None: if indices.size == 0: raise ValueError("No indices provided. Unable to infer size.") size = int(indices.max()) + 1 + if dtype is None and values.ndim > 1: + # Look for array-subtdype + new_dtype = lookup_dtype(np.dtype((new_dtype.np_type, values.shape[1:]))) # Create the new vector - w = cls.new(dtype, size, name=name) + w = cls.new(new_dtype, size, name=name) if values.ndim == 0: if dup_op is not None: raise ValueError( @@ -495,10 +502,11 @@ def ewise_union(self, other, op, left_default, right_default): other = self._expect_type( other, (Vector, Matrix, TransposedMatrix), within=method_name, argname="other", op=op ) + dtype = self.dtype if self.dtype._is_udt else None if type(left_default) is not Scalar: try: left = Scalar.from_value( - left_default, is_cscalar=False, name="" # pragma: is_grbscalar + left_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar ) except TypeError: left = self._expect_type( @@ -510,11 +518,11 @@ def ewise_union(self, other, op, left_default, right_default): op=op, ) else: - left = _as_scalar(left_default, is_cscalar=False) # pragma: is_grbscalar + left = _as_scalar(left_default, dtype, is_cscalar=False) # pragma: is_grbscalar if type(right_default) is not Scalar: try: right = Scalar.from_value( - right_default, is_cscalar=False, name="" # pragma: is_grbscalar + right_default, dtype, is_cscalar=False, name="" # pragma: is_grbscalar ) except TypeError: right = self._expect_type( @@ -526,7 +534,7 @@ def ewise_union(self, other, op, left_default, right_default): op=op, ) else: - right = _as_scalar(right_default, is_cscalar=False) # pragma: is_grbscalar + right = _as_scalar(right_default, dtype, is_cscalar=False) # pragma: is_grbscalar scalar_dtype = unify(left.dtype, right.dtype) nonscalar_dtype = unify(self.dtype, other.dtype) op = get_typed_op(op, scalar_dtype, nonscalar_dtype, is_left_scalar=True, kind="binary") @@ -609,8 +617,9 @@ 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 try: - left = Scalar.from_value(left, is_cscalar=None, name="") + left = Scalar.from_value(left, dtype, is_cscalar=None, name="") except TypeError: left = self._expect_type( left, @@ -632,15 +641,21 @@ def apply(self, op, right=None, *, left=None): extra_message=extra_message, ) if left._is_cscalar: - cfunc_name = f"GrB_Vector_apply_BinaryOp1st_{left.dtype}" + if left.dtype._is_udt: + dtype_name = "UDT" + left = _Pointer(left) + else: + dtype_name = left.dtype.name + cfunc_name = f"GrB_Vector_apply_BinaryOp1st_{dtype_name}" else: cfunc_name = "GrB_Vector_apply_BinaryOp1st_Scalar" args = [left, self] 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 else None try: - right = Scalar.from_value(right, is_cscalar=None, name="") + right = Scalar.from_value(right, dtype, is_cscalar=None, name="") except TypeError: right = self._expect_type( right, @@ -662,7 +677,12 @@ def apply(self, op, right=None, *, left=None): extra_message=extra_message, ) if right._is_cscalar: - cfunc_name = f"GrB_Vector_apply_BinaryOp2nd_{right.dtype}" + if right.dtype._is_udt: + dtype_name = "UDT" + right = _Pointer(right) + else: + dtype_name = right.dtype.name + cfunc_name = f"GrB_Vector_apply_BinaryOp2nd_{dtype_name}" else: cfunc_name = "GrB_Vector_apply_BinaryOp2nd_Scalar" args = [self, right] @@ -698,6 +718,8 @@ def reduce(self, op=monoid.plus, *, allow_empty=True): raise ValueError("allow_empty=False not allowed when using Aggregators") if allow_empty: cfunc_name = "GrB_Vector_reduce_Monoid_Scalar" + elif self.dtype._is_udt: + cfunc_name = "GrB_Vector_reduce_UDT" else: cfunc_name = "GrB_Vector_reduce_{output_dtype}" return ScalarExpression( @@ -772,8 +794,9 @@ def _extract_element(self, resolved_indexes, dtype=None, *, is_cscalar, name=Non if result is None: result = Scalar.new(dtype, is_cscalar=is_cscalar, name=name) if is_cscalar: + dtype_name = "UDT" if dtype._is_udt else dtype.name if ( - call(f"GrB_Vector_extractElement_{dtype}", [_Pointer(result), self, idx.index]) + call(f"GrB_Vector_extractElement_{dtype_name}", [_Pointer(result), self, idx.index]) is not NoValue ): result._empty = False @@ -795,8 +818,9 @@ 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 try: - value = Scalar.from_value(value, is_cscalar=None, name="") + value = Scalar.from_value(value, dtype, is_cscalar=None, name="") except TypeError: value = self._expect_type( value, @@ -809,7 +833,12 @@ def _assign_element(self, resolved_indexes, value): if value._empty: call("GrB_Vector_removeElement", [self, idx.index]) return - cfunc_name = f"GrB_Vector_setElement_{value.dtype}" + if value.dtype._is_udt: + dtype_name = "UDT" + value = _Pointer(value) + else: + dtype_name = value.dtype.name + cfunc_name = f"GrB_Vector_setElement_{dtype_name}" else: cfunc_name = "GrB_Vector_setElement_Scalar" call(cfunc_name, [self, value, idx.index]) @@ -842,8 +871,9 @@ def _prep_for_assign(self, resolved_indexes, value, mask=None, is_submask=False) 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 try: - value = Scalar.from_value(value, is_cscalar=None, name="") + value = Scalar.from_value(value, dtype, is_cscalar=None, name="") except TypeError: value = self._expect_type( value, @@ -858,7 +888,12 @@ def _prep_for_assign(self, resolved_indexes, value, mask=None, is_submask=False) raise TypeError("Single element assign does not accept a submask") # v[I](m) << c if value._is_cscalar: - cfunc_name = f"GrB_Vector_subassign_{value.dtype}" + if value.dtype._is_udt: + dtype_name = "UDT" + value = _Pointer(value) + else: + dtype_name = value.dtype.name + cfunc_name = f"GrB_Vector_subassign_{dtype_name}" else: cfunc_name = "GrB_Vector_subassign_Scalar" expr_repr = "[[{2._expr_name} elements]](%s) = {0._expr_name}" % mask.name @@ -869,7 +904,12 @@ def _prep_for_assign(self, resolved_indexes, value, mask=None, is_submask=False) index = _CArray([index.value]) cscalar = _as_scalar(1, _INDEX, is_cscalar=True) if value._is_cscalar: - cfunc_name = f"GrB_Vector_assign_{value.dtype}" + if value.dtype._is_udt: + dtype_name = "UDT" + value = _Pointer(value) + else: + dtype_name = value.dtype.name + cfunc_name = f"GrB_Vector_assign_{dtype_name}" else: cfunc_name = "GrB_Vector_assign_Scalar" expr_repr = "[[{2._expr_name} elements]] = {0._expr_name}"