From 8cfd869ef9322f453c4f063e7f37062f32418439 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 15 Mar 2022 17:37:03 -0500 Subject: [PATCH 01/14] WIP: begin adding user-defined types UDTs can be *any* fixed-length numpy type that doesn't have Python objects. Basically, UDFs that operate on the UDTs need to be able to be compiled with `@numba.njit`. --- grblas/dtypes.py | 66 ++++++++++++-- grblas/formatting.py | 8 +- grblas/matrix.py | 72 ++++++++++----- grblas/operator.py | 202 ++++++++++++++++++++++++++++++++----------- grblas/scalar.py | 37 ++++++-- grblas/utils.py | 5 +- grblas/vector.py | 76 ++++++++++------ 7 files changed, 352 insertions(+), 114 deletions(-) diff --git a/grblas/dtypes.py b/grblas/dtypes.py index 845a5a483..076d42500 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") @@ -24,23 +24,77 @@ def __repr__(self): def __eq__(self, other): if type(other) is DataType: - return self.gb_obj == other.gb_obj + return self.np_type == other.np_type else: # Attempt to use `other` as a lookup key try: other = lookup_dtype(other) - return self == other + return self.np_type == other.np_type except ValueError: raise TypeError(f"Invalid or unknown datatype: {other}") 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 self.gb_obj in _registry + + @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 + _registry[dtype] = rv + _registry[rv.gb_obj] = rv + _registry[rv.numba_type] = rv + _registry[rv.numba_type.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) + if type(numba_type) is numba.types.DType: + numba_type = numba_type.dtype + 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. + # grb_name probably isn't useful, right? + return DataType(name, gb_obj, None, f"uint8_t[{dtype.itemsize}]", numba_type, dtype) BOOL = DataType("BOOL", lib.GrB_BOOL, "GrB_BOOL", "_Bool", numba.types.bool_, np.bool_) @@ -143,7 +197,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,6 +210,8 @@ 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: diff --git a/grblas/formatting.py b/grblas/formatting.py index d4fbab700..121b976da 100644 --- a/grblas/formatting.py +++ b/grblas/formatting.py @@ -136,8 +136,9 @@ def _update_matrix_dataframe(df, matrix, rows, row_offset, columns, column_offse submatrix = submatrix[row_offset:, column_offset:].new(name="") rows, cols, vals = submatrix.to_values() df.values[rows, cols] = vals - nulls = np.isnan(vals) - df.values[rows[nulls], cols[nulls]] = "nan" + if np.issubdtype(vals.dtype, np.inexact): + nulls = np.isnan(vals) + df.values[rows[nulls], cols[nulls]] = "nan" def _update_vector_dataframe(df, vector, columns, column_offset, *, mask=None): @@ -170,7 +171,8 @@ def _update_vector_dataframe(df, vector, columns, column_offset, *, mask=None): subvector = subvector[column_offset:].new(name="") cols, vals = subvector.to_values() df.values[0, cols] = vals - df.values[0, cols[np.isnan(vals)]] = "nan" + if np.issubdtype(vals.dtype, 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 4e3969848..dc557e927 100644 --- a/grblas/matrix.py +++ b/grblas/matrix.py @@ -52,7 +52,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) @@ -254,8 +254,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 @@ -292,19 +293,23 @@ def build(self, rows, columns, values, *, dup_op=None, clear=False, nrows=None, return dup_op_given = dup_op is not None - if not dup_op_given: - dup_op = binary.plus - dup_op = get_typed_op(dup_op, self.dtype, kind="binary") - if dup_op.opclass == "Monoid": - dup_op = dup_op.binaryop + if self.dtype._is_udt: # XXX + dup_op = None else: - self._expect_op(dup_op, "BinaryOp", within="build", argname="dup_op") + if not dup_op_given: + dup_op = binary.plus + dup_op = get_typed_op(dup_op, self.dtype, kind="binary") + if dup_op.opclass == "Monoid": + dup_op = dup_op.binaryop + else: + self._expect_op(dup_op, "BinaryOp", within="build", argname="dup_op") 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 @@ -509,10 +514,11 @@ def ewise_union(self, other, op, left_default, right_default): # SS, SuiteSparse-specific: eWiseUnion method_name = "ewise_union" other = self._expect_type(other, Matrix, 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( @@ -524,11 +530,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( @@ -540,7 +546,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") @@ -655,8 +661,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, @@ -678,13 +685,15 @@ 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}" + dtype_name = "UDT" if left.dtype._is_udt else 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="") except TypeError: @@ -708,7 +717,8 @@ 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}" + dtype_name = "UDT" if right.dtype._is_udt else right.dtype.name + cfunc_name = f"GrB_Matrix_apply_BinaryOp2nd_{dtype_name}" else: cfunc_name = "GrB_Matrix_apply_BinaryOp2nd_Scalar" args = [self, right] @@ -846,9 +856,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 @@ -899,8 +910,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, @@ -913,7 +925,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]) @@ -1094,8 +1111,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) @@ -1189,7 +1207,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 = ( @@ -1208,7 +1231,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/operator.py b/grblas/operator.py index 6323e2ac4..814a0ebb5 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -241,25 +241,29 @@ def is_commutative(self): 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._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): + 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 self._monoid = None @property @@ -275,6 +279,8 @@ def monoid(self): _semiring_commutes_to = TypedBuiltinBinaryOp._semiring_commutes_to is_commutative = TypedBuiltinBinaryOp.is_commutative __call__ = TypedBuiltinBinaryOp.__call__ + orig_func = TypedUserUnaryOp.orig_func + numba_func = TypedUserUnaryOp.numba_func class TypedUserMonoid(TypedOpBase): @@ -325,11 +331,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) @@ -337,7 +344,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}" @@ -356,12 +363,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) @@ -373,7 +381,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) @@ -673,7 +681,7 @@ def _deserialize(cls, name, *args): class UnaryOp(OpBase): - __slots__ = "_func", "is_positional" + __slots__ = "_func", "is_positional", "_is_udt", "_udt_types", "_udt_ops", "_numba_func" _module = unary _modname = "unary" _typed_class = TypedBuiltinUnaryOp @@ -706,13 +714,14 @@ 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(): @@ -721,7 +730,7 @@ def _build(cls, name, func, *, anonymous=False): try: with np.errstate(divide="ignore", over="ignore", under="ignore", invalid="ignore"): ret = func(sample_val) - ret_type = lookup_dtype(type(ret)) + ret_type = lookup_dtype(type(ret), ret) 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) @@ -745,8 +754,6 @@ def _build(cls, name, func, *, anonymous=False): 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), @@ -783,27 +790,83 @@ 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_.name, ret_type.name, new_unary[0]) new_type_obj._add(op) success = True return_types[type_.name] = ret_type.name except Exception: continue - if success: + if success or is_udt: return new_type_obj else: raise UdfParseError("Unable to parse function using Numba") + def _getitem_udt(self, dtype, dtype2): + assert dtype2 is None # ignore dtype2 + np_type = dtype.np_type + if np_type in self._udt_types: + return self._udt_ops[np_type] + + sample_val = np.zeros(1, dtype=np_type)[0] + # if not np.issubdtype(np_type, np.flexible): + # # dtypes such as `np.dtype('S5')` are complicated + # sample_val = sample_val[0] + with np.errstate(divide="ignore", over="ignore", under="ignore", invalid="ignore"): + ret = self._func(sample_val) + try: + ret_type = lookup_dtype(type(ret), ret) + except Exception: + if np.issubdtype(getattr(ret_type, "dtype", np_type), np.flexible): + ret_type = dtype # Guess for now + else: + raise + + # 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 + return_type = INT8 if ret_type == "BOOL" else ret_type + + # Build wrapper because GraphBLAS wants pointers and void return + nt = numba.types + wrapper_sig = nt.void( + nt.CPointer(return_type.numba_type), + nt.CPointer(dtype.numba_type), + ) + numba_func = self._numba_func + if ret_type == "BOOL": + + def unary_wrapper(z_ptr, x_ptr): + z = numba.carray(z_ptr, 1) + x = numba.carray(x_ptr, 1) + z[0] = bool(numba_func(x[0])) # pragma: no cover + + else: + + def unary_wrapper(z_ptr, x_ptr): + z = numba.carray(z_ptr, 1) + x = numba.carray(x_ptr, 1) + z[0] = numba_func(x[0]) # pragma: no cover + + 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.name, ret_type.name, new_unary[0]) + self._udt_types[np_type] = ret_type + self._udt_ops[np_type] = 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] = ( @@ -811,10 +874,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) @@ -868,10 +931,24 @@ def _initialize(cls): op.coercions[dtype] = target_type 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._numba_func = numba_func self.is_positional = is_positional + self._is_udt = is_udt + if is_udt: + self._udt_types = {} # {np_type: DataType} + self._udt_ops = {} # {np_type: TypedUserUnaryOp} def __reduce__(self): if self._anonymous: @@ -910,7 +987,17 @@ def inner(x, y): class BinaryOp(OpBase): - __slots__ = "_monoid", "commutes_to", "_semiring_commutes_to", "_func", "is_positional" + __slots__ = ( + "_monoid", + "commutes_to", + "_semiring_commutes_to", + "_func", + "is_positional", + "_is_udt", + "_udt_types", + "_udt_ops", + "_numba_func", + ) _module = binary _modname = "binary" _typed_class = TypedBuiltinBinaryOp @@ -1014,13 +1101,14 @@ 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(): @@ -1029,7 +1117,7 @@ def _build(cls, name, func, *, anonymous=False): try: with np.errstate(divide="ignore", over="ignore", under="ignore", invalid="ignore"): ret = func(sample_val, sample_val) - ret_type = lookup_dtype(type(ret)) + ret_type = lookup_dtype(type(ret), ret) 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) @@ -1053,9 +1141,6 @@ def _build(cls, name, func, *, anonymous=False): 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) - # Build wrapper because GraphBLAS wants pointers and void return wrapper_sig = nt.void( nt.CPointer(return_type.numba_type), @@ -1097,27 +1182,25 @@ 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_.name, ret_type.name, new_binary[0]) new_type_obj._add(op) success = True return_types[type_.name] = ret_type.name except Exception: continue - if success: + if success or is_udt: return new_type_obj else: raise UdfParseError("Unable to parse function using Numba") @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] = ( @@ -1125,10 +1208,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) @@ -1257,12 +1340,23 @@ def _initialize(cls): right._semiring_commutes_to = left 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._numba_func = numba_func + self._is_udt = is_udt self.is_positional = is_positional def __reduce__(self): @@ -1401,6 +1495,10 @@ 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() @@ -1720,6 +1818,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__ @@ -1727,6 +1829,8 @@ 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._getitem_udt(dtype, dtype2) if dtype2 is not None: dtype = unify( dtype, dtype2, is_left_scalar=is_left_scalar, is_right_scalar=is_right_scalar diff --git a/grblas/scalar.py b/grblas/scalar.py index 43c601680..b27d270b7 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) @@ -224,11 +224,17 @@ 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 + return np.array(ffi.buffer(scalar.gb_obj[0 : np_type.itemsize])).view(np_type)[0] + else: return scalar.gb_obj[0] @value.setter @@ -238,11 +244,23 @@ 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: + self.gb_obj[0 : self.dtype.np_type.itemsize] = np.array( + [val], dtype=self.dtype.np_type + ).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]) + # call(f"GrB_Scalar_setElement_{val.dtype}", [self, val]) @property def nvals(self): @@ -313,7 +331,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 +351,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/utils.py b/grblas/utils.py index 7b19c7293..d297d911a 100644 --- a/grblas/utils.py +++ b/grblas/utils.py @@ -74,7 +74,7 @@ def values_to_numpy_buffer(array, dtype=None, *, copy=False, ownable=False, orde 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 @@ -128,7 +128,8 @@ def __init__(self, array=None, dtype=_INDEX, *, size=None, name=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)) + 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 0ffa36f70..587d2a8c9 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) @@ -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) @@ -274,18 +273,22 @@ def build(self, indices, values, *, dup_op=None, clear=False, size=None): return dup_op_given = dup_op is not None - if not dup_op_given: - dup_op = binary.plus - dup_op = get_typed_op(dup_op, self.dtype, kind="binary") - if dup_op.opclass == "Monoid": - dup_op = dup_op.binaryop + if self.dtype._is_udt: # XXX + dup_op = None else: - self._expect_op(dup_op, "BinaryOp", within="build", argname="dup_op") + if not dup_op_given: + dup_op = binary.plus + dup_op = get_typed_op(dup_op, self.dtype, kind="binary") + if dup_op.opclass == "Monoid": + dup_op = dup_op.binaryop + else: + self._expect_op(dup_op, "BinaryOp", within="build", argname="dup_op") 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], ) @@ -458,10 +461,11 @@ def ewise_union(self, other, op, left_default, right_default): # SS, SuiteSparse-specific: eWiseUnion method_name = "ewise_union" other = self._expect_type(other, 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( @@ -473,11 +477,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( @@ -489,7 +493,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") @@ -558,8 +562,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, @@ -581,15 +586,17 @@ 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}" + dtype_name = "UDT" if left.dtype._is_udt else 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, @@ -611,7 +618,8 @@ 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}" + dtype_name = "UDT" if right.dtype._is_udt else right.dtype.name + cfunc_name = f"GrB_Vector_apply_BinaryOp2nd_{dtype_name}" else: cfunc_name = "GrB_Vector_apply_BinaryOp2nd_Scalar" args = [self, right] @@ -721,8 +729,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 @@ -744,8 +753,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, @@ -758,7 +768,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]) @@ -791,8 +806,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, @@ -807,7 +823,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 @@ -818,7 +839,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}" From c3268a53caa4b8c75581e8344c3d2fb79d498f68 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 17 Mar 2022 23:40:01 -0500 Subject: [PATCH 02/14] UDFs on UDTs! Also, get return type of UDFs from numba. --- grblas/dtypes.py | 4 +- grblas/operator.py | 265 ++++++++++++++++++++++++++-------- grblas/tests/test_numpyops.py | 18 +-- grblas/tests/test_op.py | 1 + 4 files changed, 214 insertions(+), 74 deletions(-) diff --git a/grblas/dtypes.py b/grblas/dtypes.py index 076d42500..721a91834 100644 --- a/grblas/dtypes.py +++ b/grblas/dtypes.py @@ -86,9 +86,7 @@ def register_anonymous(dtype, name=None): raise ValueError("dtype must not be a builtin type") if name is None: name = repr(dtype) - numba_type = numba.typeof(dtype) - if type(numba_type) is numba.types.DType: - numba_type = numba_type.dtype + numba_type = numba.typeof(dtype).dtype gb_obj = ffi.new("GrB_Type*") status = lib.GrB_Type_new(gb_obj, dtype.itemsize) check_status_carg(status, "Type", gb_obj[0]) diff --git a/grblas/operator.py b/grblas/operator.py index 814a0ebb5..a2f601a24 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -537,7 +537,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 @@ -554,6 +563,14 @@ def __repr__(self): return f"{self._modname}.{self.name}" def __getitem__(self, type_): + if self._is_udt: + if type(type_) is tuple: + dtype1, dtype2 = type_ + dtype1 = lookup_dtype(dtype1) + dtype2 = lookup_dtype(dtype2) + else: + dtype1 = dtype2 = lookup_dtype(type_) + return self._getitem_udt(dtype1, dtype2) type_ = _normalize_type(type_) if type_ not in self._typed_ops: raise KeyError(f"{self.name} does not work with {type_}") @@ -681,7 +698,7 @@ def _deserialize(cls, name, *args): class UnaryOp(OpBase): - __slots__ = "_func", "is_positional", "_is_udt", "_udt_types", "_udt_ops", "_numba_func" + __slots__ = "_func", "is_positional", "_is_udt", "_numba_func" _module = unary _modname = "unary" _typed_class = TypedBuiltinUnaryOp @@ -724,13 +741,15 @@ def _build(cls, name, func, *, anonymous=False, is_udt=False): 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), ret) + if not is_udt: + for type_, sample_val in _sample_values.items(): + type_ = lookup_dtype(type_) + 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) @@ -794,32 +813,23 @@ def unary_wrapper(z, x): new_type_obj._add(op) success = True return_types[type_.name] = ret_type.name - except Exception: - continue if success or is_udt: return new_type_obj else: raise UdfParseError("Unable to parse function using Numba") def _getitem_udt(self, dtype, dtype2): - assert dtype2 is None # ignore dtype2 np_type = dtype.np_type if np_type in self._udt_types: return self._udt_ops[np_type] - sample_val = np.zeros(1, dtype=np_type)[0] - # if not np.issubdtype(np_type, np.flexible): - # # dtypes such as `np.dtype('S5')` are complicated - # sample_val = sample_val[0] - with np.errstate(divide="ignore", over="ignore", under="ignore", invalid="ignore"): - ret = self._func(sample_val) + numba_func = self._numba_func + sig = (dtype.numba_type,) try: - ret_type = lookup_dtype(type(ret), ret) - except Exception: - if np.issubdtype(getattr(ret_type, "dtype", np_type), np.flexible): - ret_type = dtype # Guess for now - else: - raise + numba_func.compile(sig) + except numba.TypingError: + raise TypeError("TODO") + ret_type = lookup_dtype(numba_func.overloads[sig].signature.return_type) # Numba is unable to handle BOOL correctly right now, but we have a workaround # See: https://github.com/numba/numba/issues/5395 @@ -832,7 +842,6 @@ def _getitem_udt(self, dtype, dtype2): nt.CPointer(return_type.numba_type), nt.CPointer(dtype.numba_type), ) - numba_func = self._numba_func if ret_type == "BOOL": def unary_wrapper(z_ptr, x_ptr): @@ -972,11 +981,11 @@ def _rfloordiv(x, y): def _absfirst(x, y): - return abs(x) + return np.abs(x) def _abssecond(x, y): - return abs(y) + return np.abs(y) def _isclose(rel_tol=1e-7, abs_tol=0.0): @@ -994,8 +1003,6 @@ class BinaryOp(OpBase): "_func", "is_positional", "_is_udt", - "_udt_types", - "_udt_ops", "_numba_func", ) _module = binary @@ -1111,13 +1118,15 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): 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), ret) + if not is_udt: + for type_, sample_val in _sample_values.items(): + type_ = lookup_dtype(type_) + 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) @@ -1186,13 +1195,74 @@ def binary_wrapper(z, x, y): new_type_obj._add(op) success = True return_types[type_.name] = ret_type.name - except Exception: - continue if success or is_udt: return new_type_obj else: raise UdfParseError("Unable to parse function using Numba") + def _getitem_udt(self, dtype, dtype2): + if dtype2 is None: + dtype2 = dtype + np_types = (dtype.np_type, dtype2.np_type) + if np_types in self._udt_types: + return self._udt_ops[np_types] + + numba_func = self._numba_func + sig = (dtype.numba_type, dtype2.numba_type) + try: + numba_func.compile(sig) + except numba.TypingError: + raise TypeError("TODO") + ret_type = lookup_dtype(numba_func.overloads[sig].signature.return_type) + + # 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 + return_type = INT8 if ret_type == "BOOL" else ret_type + + # Build wrapper because GraphBLAS wants pointers and void return + nt = numba.types + wrapper_sig = nt.void( + nt.CPointer(return_type.numba_type), + nt.CPointer(dtype.numba_type), + nt.CPointer(dtype2.numba_type), + ) + if ret_type == "BOOL": + + def binary_wrapper(z_ptr, x_ptr, y_ptr): + z = numba.carray(z_ptr, 1) + x = numba.carray(x_ptr, 1) + y = numba.carray(y_ptr, 1) + z[0] = bool(numba_func(x[0], y[0])) # pragma: no cover + + else: + + def binary_wrapper(z_ptr, x_ptr, y_ptr): + z = numba.carray(z_ptr, 1) + x = numba.carray(x_ptr, 1) + y = numba.carray(y_ptr, 1) + z[0] = numba_func(x[0], y[0]) # pragma: no cover + + 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.name, # Should we do a tuple for both inputs? + ret_type.name, + new_binary[0], + ) + self._udt_types[np_types] = ret_type + self._udt_ops[np_types] = op + return op + @classmethod def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): if parameterized: @@ -1358,6 +1428,9 @@ def __init__( self._numba_func = numba_func self._is_udt = is_udt self.is_positional = is_positional + if is_udt: + self._udt_types = {} # {(np_type, np_type): DataType} + self._udt_ops = {} # {(np_type, np_type): TypedUserBinaryOp} def __reduce__(self): if self._anonymous: @@ -1411,32 +1484,67 @@ 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 = 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) return new_type_obj + def _getitem_udt(self, dtype, dtype2): + if dtype2 is None: + dtype2 = dtype + elif dtype != dtype2: + raise TypeError("TODO") + np_type = dtype.np_type + if np_type in self._udt_types: + return self._udt_ops[np_type] + binaryop = self._binaryop._getitem_udt(dtype, dtype2) + from .scalar import Scalar + + identity = Scalar.from_value(self._identity, dtype=dtype, 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.name, + dtype.name, + new_monoid[0], + binaryop, + identity, + ) + self._udt_types[np_type] = dtype + self._udt_ops[np_type] = op + return op + @classmethod def register_anonymous(cls, binaryop, identity, name=None): if type(binaryop) is ParameterizedBinaryOp: @@ -1475,6 +1583,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 = {} # {np_type: DataType} + self._udt_ops = {} # {np_type: TypedUserMonoid} def __reduce__(self): if self._anonymous: @@ -1590,6 +1701,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 @@ -1620,6 +1733,31 @@ def _build(cls, name, monoid, binaryop, *, anonymous=False): new_type_obj._add(op) return new_type_obj + def _getitem_udt(self, dtype, dtype2): + if dtype2 is None: + dtype2 = dtype + np_types = (dtype.np_type, dtype2.np_type) + if np_types in self._udt_types: + return self._udt_ops[np_types] + binaryop = self._binaryop._getitem_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.name, # Should we do a tuple for both inputs? + ret_type, + new_semiring[0], + monoid, + binaryop, + ) + self._udt_types[np_types] = dtype + self._udt_ops[np_types] = op + return op + @classmethod def register_anonymous(cls, monoid, binaryop, name=None): if type(monoid) is ParameterizedMonoid or type(binaryop) is ParameterizedBinaryOp: @@ -1791,6 +1929,9 @@ def __init__(self, name, monoid=None, binaryop=None, *, anonymous=False): super().__init__(name, anonymous=anonymous) self._monoid = monoid self._binaryop = binaryop + if binaryop is not None and binaryop._is_udt: + self._udt_types = {} # {(np_type, np_type): DataType} + self._udt_ops = {} # {(np_type, np_type): TypedUserSemiring} def __reduce__(self): if self._anonymous: diff --git a/grblas/tests/test_numpyops.py b/grblas/tests/test_numpyops.py index bc476efe6..acc688367 100644 --- a/grblas/tests/test_numpyops.py +++ b/grblas/tests/test_numpyops.py @@ -53,12 +53,12 @@ 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.name not in op.types or unary_name in blocklist.get( gb_input.dtype.name, () ): continue # pragma: no cover @@ -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.name not in op.types or binary_name in blocklist.get( gb_left.dtype.name, () ): continue @@ -210,15 +210,15 @@ 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.name not in op.types or binary_name in blocklist.get( gb_left.dtype.name, () ): continue # pragma: no cover @@ -240,7 +240,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 a3a63d300..0e6689221 100644 --- a/grblas/tests/test_op.py +++ b/grblas/tests/test_op.py @@ -469,6 +469,7 @@ def times_minus_sum(x, y): BinaryOp.register_new("bin_test_func", times_minus_sum) assert hasattr(binary, "bin_test_func") comp_set = { + "BOOL", # goes to INT64 "INT8", "INT16", "INT32", From b03daea2f3694800ccbd5e63866cc9806455fddc Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Sat, 19 Mar 2022 13:59:05 -0500 Subject: [PATCH 03/14] Add recipes for first, any, eq, etc. on UDTs --- grblas/matrix.py | 2 + grblas/operator.py | 184 +++++++++++++++++++++++++++++++++------------ grblas/vector.py | 14 +++- 3 files changed, 148 insertions(+), 52 deletions(-) diff --git a/grblas/matrix.py b/grblas/matrix.py index dc557e927..fde94e3a3 100644 --- a/grblas/matrix.py +++ b/grblas/matrix.py @@ -832,6 +832,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( diff --git a/grblas/operator.py b/grblas/operator.py index a2f601a24..3e71ee95e 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -10,7 +10,7 @@ import numpy as np from . import config, ffi, lib -from .dtypes import INT8, _sample_values, _supports_complex, lookup_dtype, unify +from .dtypes import BOOL, INT8, UINT8, _sample_values, _supports_complex, lookup_dtype, unify from .exceptions import UdfParseError, check_status_carg from .expr import InfixExprBase from .utils import libget, output_type @@ -154,12 +154,14 @@ def commutes_to(self): commutes_to = self.parent.commutes_to if commutes_to is not None and self.type in commutes_to._typed_ops: return commutes_to[self.type] + # TODO: what about UDTs and `gb.binary.any`? @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: return commutes_to[self.type] + # TODO: what about UDTs and `gb.binary.any`? @property def is_commutative(self): @@ -268,10 +270,9 @@ def __init__(self, parent, name, type_, return_type, gb_obj): @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 @@ -558,23 +559,28 @@ 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_): - if self._is_udt: - if type(type_) is tuple: - dtype1, dtype2 = type_ - dtype1 = lookup_dtype(dtype1) - dtype2 = lookup_dtype(dtype2) + if not self._is_udt: + type_ = _normalize_type(type_) + if type_ not in self._typed_ops: + if self._udt_types is None: + raise KeyError(f"{self.name} does not work with {type_}") else: - dtype1 = dtype2 = lookup_dtype(type_) - return self._getitem_udt(dtype1, dtype2) - 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_] + return self._typed_ops[type_] + # This is a UDT or is able to operate on UDTs such as `first` any `any` + if type(type_) is tuple: + dtype1, dtype2 = type_ + dtype1 = lookup_dtype(dtype1) + dtype2 = lookup_dtype(dtype2) + else: + dtype1 = dtype2 = lookup_dtype(type_) + return self._getitem_udt(dtype1, dtype2) def _add(self, op): self._typed_ops[op.type] = op @@ -586,8 +592,24 @@ def __delitem__(self, type_): del self.types[type_] def __contains__(self, type_): - type_ = _normalize_type(type_) - return type_ in self._typed_ops + if not self._is_udt: + type_ = _normalize_type(type_) + if type_ in self._typed_ops: + return True + elif self._udt_types is None: + return False + if type(type_) is tuple: + dtype1, dtype2 = type_ + dtype1 = lookup_dtype(dtype1) + dtype2 = lookup_dtype(dtype2) + else: + dtype1 = dtype2 = lookup_dtype(type_) + try: + self._getitem_udt(dtype1, dtype2) + except TypeError: + return False + else: + return True @classmethod def _remove_nesting(cls, funcname, *, module=None, modname=None, strict=True): @@ -995,6 +1017,18 @@ def inner(x, y): return inner +def _first(x, y): + return x + + +def _second(x, y): + return y + + +def _pair(x, y): + return 1 + + class BinaryOp(OpBase): __slots__ = ( "_monoid", @@ -1207,41 +1241,72 @@ def _getitem_udt(self, dtype, dtype2): if np_types in self._udt_types: return self._udt_ops[np_types] - numba_func = self._numba_func - sig = (dtype.numba_type, dtype2.numba_type) - try: - numba_func.compile(sig) - except numba.TypingError: - raise TypeError("TODO") - ret_type = lookup_dtype(numba_func.overloads[sig].signature.return_type) - - # 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 - return_type = INT8 if ret_type == "BOOL" else ret_type - - # Build wrapper because GraphBLAS wants pointers and void return nt = numba.types - wrapper_sig = nt.void( - nt.CPointer(return_type.numba_type), - nt.CPointer(dtype.numba_type), - nt.CPointer(dtype2.numba_type), - ) - if ret_type == "BOOL": + if self.name == "eq": + assert dtype == dtype2 # XXX: must be same size? + itemsize = dtype.np_type.itemsize + ret_type = BOOL + wrapper_sig = nt.void( + nt.CPointer(INT8.numba_type), + nt.CPointer(UINT8.numba_type), + nt.CPointer(UINT8.numba_type), + ) def binary_wrapper(z_ptr, x_ptr, y_ptr): - z = numba.carray(z_ptr, 1) - x = numba.carray(x_ptr, 1) - y = numba.carray(y_ptr, 1) - z[0] = bool(numba_func(x[0], y[0])) # pragma: no cover + x = numba.carray(x_ptr, itemsize) + y = numba.carray(y_ptr, itemsize) + z_ptr[0] = (x == y).all() + + elif self.name == "ne": + assert dtype == dtype2 # XXX: must be same size? + itemsize = dtype.np_type.itemsize + ret_type = BOOL + wrapper_sig = nt.void( + nt.CPointer(INT8.numba_type), + nt.CPointer(UINT8.numba_type), + nt.CPointer(UINT8.numba_type), + ) + + def binary_wrapper(z_ptr, x_ptr, y_ptr): + x = numba.carray(x_ptr, itemsize) + y = numba.carray(y_ptr, itemsize) + z_ptr[0] = (x != y).all() else: + numba_func = self._numba_func + sig = (dtype.numba_type, dtype2.numba_type) + try: + numba_func.compile(sig) + except numba.TypingError: + raise TypeError("TODO") + ret_type = lookup_dtype(numba_func.overloads[sig].signature.return_type) + + # 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 + return_type = INT8 if ret_type == "BOOL" else ret_type + + # Build wrapper because GraphBLAS wants pointers and void return + wrapper_sig = nt.void( + nt.CPointer(return_type.numba_type), + nt.CPointer(dtype.numba_type), + nt.CPointer(dtype2.numba_type), + ) + if ret_type == "BOOL": - def binary_wrapper(z_ptr, x_ptr, y_ptr): - z = numba.carray(z_ptr, 1) - x = numba.carray(x_ptr, 1) - y = numba.carray(y_ptr, 1) - z[0] = numba_func(x[0], y[0]) # pragma: no cover + def binary_wrapper(z_ptr, x_ptr, y_ptr): + z = numba.carray(z_ptr, 1) + x = numba.carray(x_ptr, 1) + y = numba.carray(y_ptr, 1) + z[0] = bool(numba_func(x[0], y[0])) # pragma: no cover + + else: + + def binary_wrapper(z_ptr, x_ptr, y_ptr): + z = numba.carray(z_ptr, 1) + x = numba.carray(x_ptr, 1) + y = numba.carray(y_ptr, 1) + z[0] = numba_func(x[0], y[0]) # pragma: no cover binary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(binary_wrapper) new_binary = ffi_new("GrB_BinaryOp*") @@ -1408,6 +1473,20 @@ 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._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 = {} cls._initialized = True def __init__( @@ -1525,7 +1604,7 @@ def _getitem_udt(self, dtype, dtype2): np_type = dtype.np_type if np_type in self._udt_types: return self._udt_ops[np_type] - binaryop = self._binaryop._getitem_udt(dtype, dtype2) + binaryop = self.binaryop._getitem_udt(dtype, dtype2) from .scalar import Scalar identity = Scalar.from_value(self._identity, dtype=dtype, is_cscalar=True) @@ -1645,6 +1724,11 @@ def _initialize(cls): 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 @@ -1739,8 +1823,8 @@ def _getitem_udt(self, dtype, dtype2): np_types = (dtype.np_type, dtype2.np_type) if np_types in self._udt_types: return self._udt_ops[np_types] - binaryop = self._binaryop._getitem_udt(dtype, dtype2) - monoid = self._monoid[binaryop.return_type] + binaryop = self.binaryop._getitem_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) diff --git a/grblas/vector.py b/grblas/vector.py index 587d2a8c9..67921bcab 100644 --- a/grblas/vector.py +++ b/grblas/vector.py @@ -586,7 +586,11 @@ def apply(self, op, right=None, *, left=None): extra_message=extra_message, ) if left._is_cscalar: - dtype_name = "UDT" if left.dtype._is_udt else left.dtype.name + 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" @@ -618,7 +622,11 @@ def apply(self, op, right=None, *, left=None): extra_message=extra_message, ) if right._is_cscalar: - dtype_name = "UDT" if right.dtype._is_udt else right.dtype.name + 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" @@ -655,6 +663,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( From 1e3317f609541b2533665b41cad87fc08ca8206e Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 29 Mar 2022 20:06:16 -0500 Subject: [PATCH 04/14] Test UDTs Also: - make dtypes hashable - change operator attributes to use dtypes, not dtype names - update `==` between UDTs (ignore bytes where appropriate) --- grblas/_agg.py | 9 +- grblas/dtypes.py | 52 +++- grblas/matrix.py | 14 +- grblas/operator.py | 441 +++++++++++++++++----------- grblas/scalar.py | 5 +- grblas/tests/test_matrix.py | 62 ++++ grblas/tests/test_op.py | 177 +++++++---- grblas/tests/test_operator_types.py | 22 +- grblas/tests/test_scalar.py | 2 +- grblas/tests/test_vector.py | 57 ++++ 10 files changed, 579 insertions(+), 262 deletions(-) diff --git a/grblas/_agg.py b/grblas/_agg.py index cfc87107a..e0876dd6a 100644 --- a/grblas/_agg.py +++ b/grblas/_agg.py @@ -5,7 +5,6 @@ from . import agg, binary, monoid, semiring, unary from .dtypes import lookup_dtype, unify -from .operator import _normalize_type from .scalar import Scalar @@ -15,7 +14,7 @@ def _get_types(ops, initdtype): 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()} + prev = {key: unify(lookup_dtype(val), initdtype) for key, val in ops[0].types.items()} for op in ops[1:]: cur = {} types = op.types @@ -47,7 +46,7 @@ def __init__( 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 @@ -82,7 +81,7 @@ def types(self): return self._types def __getitem__(self, dtype): - dtype = _normalize_type(dtype) + dtype = lookup_dtype(dtype) if dtype not in self.types: raise KeyError(f"{self.name} does not work with {dtype}") if dtype not in self._typed_ops: @@ -90,7 +89,7 @@ def __getitem__(self, dtype): return self._typed_ops[dtype] def __contains__(self, dtype): - dtype = _normalize_type(dtype) + dtype = lookup_dtype(dtype) return dtype in self.types def __repr__(self): diff --git a/grblas/dtypes.py b/grblas/dtypes.py index f97659f9a..b970868cb 100644 --- a/grblas/dtypes.py +++ b/grblas/dtypes.py @@ -24,14 +24,27 @@ def __repr__(self): def __eq__(self, other): if type(other) is DataType: - return self.np_type == other.np_type - else: - # Attempt to use `other` as a lookup key - try: - other = lookup_dtype(other) - return self.np_type == other.np_type - except ValueError: - raise TypeError(f"Invalid or unknown datatype: {other}") from None + return self.gb_obj is other.gb_obj + # Attempt to use `other` as a lookup key + try: + other = lookup_dtype(other) + return self.gb_obj is other.gb_obj + 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: + return str(self.np_type) < str(other.np_type) + # Attempt to use `other` as a lookup key + try: + other = lookup_dtype(other) + return str(self.np_type) < str(other.np_type) + except ValueError: + raise TypeError(f"Invalid or unknown datatype: {other}") from None def __reduce__(self): if self._is_udt: @@ -68,10 +81,6 @@ def register_new(name, dtype): raise ValueError(f"{name!r} name for dtype is unavailable") rv = register_anonymous(dtype, name) _registry[name] = rv - _registry[dtype] = rv - _registry[rv.gb_obj] = rv - _registry[rv.numba_type] = rv - _registry[rv.numba_type.name] = rv globals()[name] = rv return rv @@ -87,12 +96,23 @@ def register_anonymous(dtype, name=None): if name is None: name = repr(dtype) numba_type = numba.typeof(dtype).dtype - gb_obj = ffi.new("GrB_Type*") - status = lib.GrB_Type_new(gb_obj, dtype.itemsize) - check_status_carg(status, "Type", gb_obj[0]) + + if dtype in _registry: + other = _registry[dtype] + gb_obj = other.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. # grb_name probably isn't useful, right? - return DataType(name, gb_obj, None, f"uint8_t[{dtype.itemsize}]", numba_type, dtype) + 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_) diff --git a/grblas/matrix.py b/grblas/matrix.py index 510ce4a30..0efa25d38 100644 --- a/grblas/matrix.py +++ b/grblas/matrix.py @@ -690,7 +690,11 @@ def apply(self, op, right=None, *, left=None): extra_message=extra_message, ) if left._is_cscalar: - dtype_name = "UDT" if left.dtype._is_udt else left.dtype.name + 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" @@ -700,7 +704,7 @@ def apply(self, op, right=None, *, left=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, @@ -722,7 +726,11 @@ def apply(self, op, right=None, *, left=None): extra_message=extra_message, ) if right._is_cscalar: - dtype_name = "UDT" if right.dtype._is_udt else right.dtype.name + 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" diff --git a/grblas/operator.py b/grblas/operator.py index 8f4ab1f9c..3103c4e53 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 BOOL, INT8, UINT8, _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,14 +95,43 @@ 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, (rv.size, dtype.itemsize) + _udt_mask_cache[dtype] = rv + return rv + + class TypedOpBase: __slots__ = "parent", "name", "type", "return_type", "gb_obj", "gb_name", "__weakref__" def __init__(self, parent, name, type_, return_type, gb_obj, gb_name): self.parent = parent self.name = name - self.type = _normalize_type(type_) - self.return_type = _normalize_type(return_type) + self.type = lookup_dtype(type_) + self.return_type = lookup_dtype(return_type) self.gb_obj = gb_obj self.gb_name = gb_name @@ -567,7 +611,7 @@ def __repr__(self): def __getitem__(self, type_): if not self._is_udt: - type_ = _normalize_type(type_) + type_ = lookup_dtype(type_) if type_ not in self._typed_ops: if self._udt_types is None: raise KeyError(f"{self.name} does not work with {type_}") @@ -587,13 +631,13 @@ def _add(self, 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_): if not self._is_udt: - type_ = _normalize_type(type_) + type_ = lookup_dtype(type_) if type_ in self._typed_ops: return True elif self._udt_types is None: @@ -776,24 +820,20 @@ def _build(cls, name, func, *, anonymous=False, is_udt=False): ("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 # Build wrapper because GraphBLAS wants pointers and void return wrapper_sig = nt.void( @@ -801,8 +841,8 @@ def _build(cls, name, func, *, anonymous=False, is_udt=False): 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 @@ -812,7 +852,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 @@ -831,19 +871,18 @@ def unary_wrapper(z, x): "UnaryOp", new_unary, ) - op = TypedUserUnaryOp(new_type_obj, name, type_.name, ret_type.name, new_unary[0]) + 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 + return_types[type_] = ret_type if success or is_udt: return new_type_obj else: raise UdfParseError("Unable to parse function using Numba") def _getitem_udt(self, dtype, dtype2): - np_type = dtype.np_type - if np_type in self._udt_types: - return self._udt_ops[np_type] + if dtype in self._udt_types: + return self._udt_ops[dtype] numba_func = self._numba_func sig = (dtype.numba_type,) @@ -856,7 +895,7 @@ def _getitem_udt(self, dtype, dtype2): # 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 - return_type = INT8 if ret_type == "BOOL" else ret_type + return_type = INT8 if ret_type == BOOL else ret_type # Build wrapper because GraphBLAS wants pointers and void return nt = numba.types @@ -864,7 +903,7 @@ def _getitem_udt(self, dtype, dtype2): nt.CPointer(return_type.numba_type), nt.CPointer(dtype.numba_type), ) - if ret_type == "BOOL": + if ret_type == BOOL: def unary_wrapper(z_ptr, x_ptr): z = numba.carray(z_ptr, 1) @@ -885,9 +924,9 @@ def unary_wrapper(z_ptr, x_ptr): "UnaryOp", new_unary, ) - op = TypedUserUnaryOp(self, self.name, dtype.name, ret_type.name, new_unary[0]) - self._udt_types[np_type] = ret_type - self._udt_ops[np_type] = op + 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 @@ -926,6 +965,19 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal 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 ( @@ -935,17 +987,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 @@ -956,9 +1005,7 @@ 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 @@ -980,8 +1027,8 @@ def __init__( self.is_positional = is_positional self._is_udt = is_udt if is_udt: - self._udt_types = {} # {np_type: DataType} - self._udt_ops = {} # {np_type: TypedUserUnaryOp} + self._udt_types = {} # {dtype: DataType} + self._udt_ops = {} # {dtype: TypedUserUnaryOp} def __reduce__(self): if self._anonymous: @@ -1167,24 +1214,20 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): ("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 # Build wrapper because GraphBLAS wants pointers and void return wrapper_sig = nt.void( @@ -1193,8 +1236,8 @@ def _build(cls, name, func, *, is_udt=False, 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 @@ -1204,7 +1247,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 @@ -1227,10 +1270,10 @@ def binary_wrapper(z, x, y): "BinaryOp", new_binary, ) - op = TypedUserBinaryOp(new_type_obj, name, type_.name, ret_type.name, new_binary[0]) + 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 + return_types[type_] = ret_type if success or is_udt: return new_type_obj else: @@ -1239,29 +1282,52 @@ def binary_wrapper(z, x, y): def _getitem_udt(self, dtype, dtype2): if dtype2 is None: dtype2 = dtype - np_types = (dtype.np_type, dtype2.np_type) - if np_types in self._udt_types: - return self._udt_ops[np_types] + dtypes = (dtype, dtype2) + if dtypes in self._udt_types: + return self._udt_ops[dtypes] nt = numba.types - if self.name == "eq": + if self.name == "eq" and not self._anonymous: assert dtype == dtype2 # XXX: must be same size? 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): - x = numba.carray(x_ptr, itemsize) - y = numba.carray(y_ptr, itemsize) - z_ptr[0] = (x == y).all() + def binary_wrapper(z_ptr, x_ptr, y_ptr): + 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: - elif self.name == "ne": + def binary_wrapper(z_ptr, x_ptr, y_ptr): + 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 == dtype2 # XXX: must be same size? itemsize = dtype.np_type.itemsize + mask = _udt_mask(dtype.np_type) ret_type = BOOL wrapper_sig = nt.void( nt.CPointer(INT8.numba_type), @@ -1269,10 +1335,31 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): nt.CPointer(UINT8.numba_type), ) - def binary_wrapper(z_ptr, x_ptr, y_ptr): - x = numba.carray(x_ptr, itemsize) - y = numba.carray(y_ptr, itemsize) - z_ptr[0] = (x != y).all() + if mask.all(): + + def binary_wrapper(z_ptr, x_ptr, y_ptr): + 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): + 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 @@ -1322,12 +1409,12 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): op = TypedUserBinaryOp( self, self.name, - dtype.name, # Should we do a tuple for both inputs? - ret_type.name, + dtype, # Should we do a tuple for both inputs? + ret_type, new_binary[0], ) - self._udt_types[np_types] = ret_type - self._udt_ops[np_types] = op + self._udt_types[dtypes] = ret_type + self._udt_ops[dtypes] = op return op @classmethod @@ -1382,10 +1469,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, @@ -1408,33 +1495,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 @@ -1443,8 +1540,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: @@ -1454,15 +1551,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: 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, right in cls._commutes_to.items(): left = getattr(binary, left) @@ -1512,8 +1607,8 @@ def __init__( self._is_udt = is_udt self.is_positional = is_positional if is_udt: - self._udt_types = {} # {(np_type, np_type): DataType} - self._udt_ops = {} # {(np_type, np_type): TypedUserBinaryOp} + self._udt_types = {} # {(dtype, dtype): DataType} + self._udt_ops = {} # {(dtype, dtype): TypedUserBinaryOp} def __reduce__(self): if self._anonymous: @@ -1591,7 +1686,7 @@ def _build(cls, name, binaryop, identity, *, anonymous=False): op = TypedUserMonoid( new_type_obj, name, - type_.name, + type_, ret_type, new_monoid[0], binaryop[type_], @@ -1605,27 +1700,27 @@ def _getitem_udt(self, dtype, dtype2): dtype2 = dtype elif dtype != dtype2: raise TypeError("TODO") - np_type = dtype.np_type - if np_type in self._udt_types: - return self._udt_ops[np_type] + if dtype in self._udt_types: + return self._udt_ops[dtype] binaryop = self.binaryop._getitem_udt(dtype, dtype2) from .scalar import Scalar - identity = Scalar.from_value(self._identity, dtype=dtype, is_cscalar=True) + 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.name, - dtype.name, + dtype, + ret_type, new_monoid[0], binaryop, identity, ) - self._udt_types[np_type] = dtype - self._udt_ops[np_type] = op + self._udt_types[dtype] = dtype + self._udt_ops[dtype] = op return op @classmethod @@ -1667,8 +1762,8 @@ def __init__(self, name, binaryop=None, identity=None, *, anonymous=False): if binaryop is not None: binaryop._monoid = self if binaryop._is_udt: - self._udt_types = {} # {np_type: DataType} - self._udt_ops = {} # {np_type: TypedUserMonoid} + self._udt_types = {} # {dtype: DataType} + self._udt_ops = {} # {dtype: TypedUserMonoid} def __reduce__(self): if self._anonymous: @@ -1696,37 +1791,37 @@ def _is_udt(self): @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 @@ -1824,9 +1919,9 @@ def _build(cls, name, monoid, binaryop, *, anonymous=False): def _getitem_udt(self, dtype, dtype2): if dtype2 is None: dtype2 = dtype - np_types = (dtype.np_type, dtype2.np_type) - if np_types in self._udt_types: - return self._udt_ops[np_types] + dtypes = (dtype, dtype2) + if dtypes in self._udt_types: + return self._udt_ops[dtypes] binaryop = self.binaryop._getitem_udt(dtype, dtype2) monoid = self.monoid[binaryop.return_type] ret_type = monoid.return_type @@ -1836,14 +1931,14 @@ def _getitem_udt(self, dtype, dtype2): op = TypedUserSemiring( new_semiring, self.name, - dtype.name, # Should we do a tuple for both inputs? + dtype, # Should we do a tuple for both inputs? ret_type, new_semiring[0], monoid, binaryop, ) - self._udt_types[np_types] = dtype - self._udt_ops[np_types] = op + self._udt_types[dtypes] = dtype + self._udt_ops[dtypes] = op return op @classmethod @@ -1921,13 +2016,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 ( @@ -1937,11 +2058,8 @@ def _initialize(cls): "secondi", "secondi1", "secondj", "secondj1", ), ( - ( - "BOOL", "FC32", "FC64", "FP32", "FP64", "INT8", "INT16", - "UINT8", "UINT16", "UINT32", "UINT64", - ), - "INT64", + position_dtypes, + INT64, ), ), ( @@ -1949,30 +2067,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 @@ -1986,9 +2101,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 @@ -2008,11 +2121,11 @@ 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): @@ -2020,8 +2133,8 @@ def __init__(self, name, monoid=None, binaryop=None, *, anonymous=False): self._monoid = monoid self._binaryop = binaryop if binaryop is not None and binaryop._is_udt: - self._udt_types = {} # {(np_type, np_type): DataType} - self._udt_ops = {} # {(np_type, np_type): TypedUserSemiring} + self._udt_types = {} # {(dtype, dtype): DataType} + self._udt_ops = {} # {(dtype, dtype): TypedUserSemiring} def __reduce__(self): if self._anonymous: diff --git a/grblas/scalar.py b/grblas/scalar.py index f8e010194..928e73188 100644 --- a/grblas/scalar.py +++ b/grblas/scalar.py @@ -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, diff --git a/grblas/tests/test_matrix.py b/grblas/tests/test_matrix.py index cd22cb54c..2dcea6fba 100644 --- a/grblas/tests/test_matrix.py +++ b/grblas/tests/test_matrix.py @@ -3119,3 +3119,65 @@ 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)) + + # 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) + + # agg: any_value, first, last, first_index, last_index, count diff --git a/grblas/tests/test_op.py b/grblas/tests/test_op.py index 9f231c7ce..f5b5722bf 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"] @@ -120,30 +124,30 @@ def plus_one(x): UnaryOp.register_new("plus_one", plus_one) assert hasattr(unary, "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") @@ -452,17 +456,17 @@ def is_positive(x): 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() @@ -477,20 +481,20 @@ def times_minus_sum(x, y): BinaryOp.register_new("bin_test_func", times_minus_sum) assert hasattr(binary, "bin_test_func") comp_set = { - "BOOL", # goes to INT64 - "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) @@ -507,16 +511,16 @@ def plus_plus_one(x, y): 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) @@ -572,20 +576,20 @@ def plus_three(x): 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) @@ -1015,3 +1019,46 @@ def test_positional(): assert semiring.any_secondj[int].is_positional assert not semiring.any_first.is_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) + v = Vector.new(udt, size=3) + w = Vector.new(udt, size=3) + v[:] = 0 + w[:] = 1 + + def _udt_identity(val): + return val + + 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) + + def _udt_getx(val): + return val["x"] + + 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 + + udt_first = BinaryOp.register_anonymous(_udt_first, "udt_first", is_udt=True) + assert udt in udt_first + assert udt_first(v & w).new().isequal(v) + assert udt_first(v, 1).new().isequal(v) + + udt_any = Monoid.register_new("udt_any", udt_first, (0, 0)) + assert udt in udt_any + 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 udt_semiring(v @ v).new() == (0, 0) diff --git a/grblas/tests/test_operator_types.py b/grblas/tests/test_operator_types.py index f476ab495..415e543b4 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 @@ -153,6 +160,9 @@ "myplus", "plus_myplus", "plus_numpy_copysign", + "udt_identity", + "udt_any", + "udt_semiring", # Semirings created by commutation "any_rfloordiv", "any_rtruediv", diff --git a/grblas/tests/test_scalar.py b/grblas/tests/test_scalar.py index e68c2a191..e08f41d3f 100644 --- a/grblas/tests/test_scalar.py +++ b/grblas/tests/test_scalar.py @@ -258,7 +258,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 7025fb9fd..2f82dad4f 100644 --- a/grblas/tests/test_vector.py +++ b/grblas/tests/test_vector.py @@ -1667,3 +1667,60 @@ def test_delete_via_scalar(v): assert v.isequal(Vector.from_values([4, 6], [2, 0])) del 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") + 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) + + # 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.first_index).new() == 0 + # assert v.reduce(agg.last_index).new() == 2 + # assert v.reduce(agg.count).new() == 3 From 9e4038166b225c35f1c6c7b4d9110bdb32feae64 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 29 Mar 2022 22:02:20 -0500 Subject: [PATCH 05/14] Allow positional operators to work on UDTs; also, coverage --- grblas/dtypes.py | 9 +++-- grblas/operator.py | 65 ++++++++++++++++++++++++----------- grblas/tests/test_matrix.py | 8 ++--- grblas/tests/test_numpyops.py | 10 +++--- grblas/tests/test_op.py | 59 +++++++++++++++++-------------- grblas/tests/test_pickle.py | 16 ++++----- grblas/tests/test_vector.py | 12 ++++--- 7 files changed, 105 insertions(+), 74 deletions(-) diff --git a/grblas/dtypes.py b/grblas/dtypes.py index b970868cb..d2ae847c0 100644 --- a/grblas/dtypes.py +++ b/grblas/dtypes.py @@ -24,11 +24,11 @@ def __repr__(self): def __eq__(self, other): if type(other) is DataType: - return self.gb_obj is other.gb_obj + return self is other # Attempt to use `other` as a lookup key try: other = lookup_dtype(other) - return self.gb_obj is other.gb_obj + return self is other except ValueError: raise TypeError(f"Invalid or unknown datatype: {other}") from None @@ -59,7 +59,7 @@ def _carg(self): @property def _is_anonymous(self): - return self.gb_obj in _registry + return globals().get(self.name) is not self @property def _is_udt(self): @@ -98,8 +98,7 @@ def register_anonymous(dtype, name=None): numba_type = numba.typeof(dtype).dtype if dtype in _registry: - other = _registry[dtype] - gb_obj = other.gb_obj + gb_obj = _registry[dtype].gb_obj else: gb_obj = ffi.new("GrB_Type*") status = lib.GrB_Type_new(gb_obj, dtype.itemsize) diff --git a/grblas/operator.py b/grblas/operator.py index 3103c4e53..c5e8620a6 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -103,6 +103,7 @@ def _udt_mask(dtype): if dtype in _udt_mask_cache: return _udt_mask_cache[dtype] if dtype.subdtype is not None: + 1 / 0 mask = _udt_mask(dtype.subdtype[0]) N = reduce(mul, dtype.subdtype[1]) rv = np.concatenate([mask] * N) @@ -295,6 +296,7 @@ def __init__(self, parent, name, type_, return_type, gb_obj): @property def orig_func(self): + 1 / 0 return self.parent._func @property @@ -614,11 +616,14 @@ def __getitem__(self, type_): 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` if type(type_) is tuple: + 1 / 0 dtype1, dtype2 = type_ dtype1 = lookup_dtype(dtype1) dtype2 = lookup_dtype(dtype2) @@ -638,11 +643,12 @@ def __delitem__(self, type_): def __contains__(self, type_): if not self._is_udt: type_ = lookup_dtype(type_) - if type_ in self._typed_ops: + if type_ in self._typed_ops or self.is_positional: return True elif self._udt_types is None: return False if type(type_) is tuple: + 1 / 0 dtype1, dtype2 = type_ dtype1 = lookup_dtype(dtype1) dtype2 = lookup_dtype(dtype2) @@ -651,6 +657,7 @@ def __contains__(self, type_): try: self._getitem_udt(dtype1, dtype2) except TypeError: + 1 / 0 return False else: return True @@ -889,8 +896,11 @@ def _getitem_udt(self, dtype, dtype2): try: numba_func.compile(sig) except numba.TypingError: + 1 / 0 raise TypeError("TODO") 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 # Numba is unable to handle BOOL correctly right now, but we have a workaround # See: https://github.com/numba/numba/issues/5395 @@ -905,17 +915,17 @@ def _getitem_udt(self, dtype, dtype2): ) if ret_type == BOOL: - def unary_wrapper(z_ptr, x_ptr): + def unary_wrapper(z_ptr, x_ptr): # pragma: no cover z = numba.carray(z_ptr, 1) x = numba.carray(x_ptr, 1) - z[0] = bool(numba_func(x[0])) # pragma: no cover + z[0] = bool(numba_func(x[0])) else: - def unary_wrapper(z_ptr, x_ptr): + def unary_wrapper(z_ptr, x_ptr): # pragma: no cover z = numba.carray(z_ptr, 1) x = numba.carray(x_ptr, 1) - z[0] = numba_func(x[0]) # pragma: no cover + z[0] = numba_func(x[0]) unary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(unary_wrapper) new_unary = ffi_new("GrB_UnaryOp*") @@ -1044,38 +1054,38 @@ def __reduce__(self): 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 np.abs(x) + return np.abs(x) # pragma: no cover def _abssecond(x, y): - return np.abs(y) + return np.abs(y) # 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 + return x # pragma: no cover def _second(x, y): - return y + return y # pragma: no cover def _pair(x, y): - return 1 + return 1 # pragma: no cover class BinaryOp(OpBase): @@ -1281,6 +1291,7 @@ def binary_wrapper(z, x, y): def _getitem_udt(self, dtype, dtype2): if dtype2 is None: + 1 / 0 dtype2 = dtype dtypes = (dtype, dtype2) if dtypes in self._udt_types: @@ -1300,7 +1311,7 @@ def _getitem_udt(self, dtype, dtype2): # PERF: we can probably make this faster if mask.all(): - def binary_wrapper(z_ptr, x_ptr, y_ptr): + 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): @@ -1313,7 +1324,7 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): else: - def binary_wrapper(z_ptr, x_ptr, y_ptr): + 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): @@ -1337,7 +1348,7 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): if mask.all(): - def binary_wrapper(z_ptr, x_ptr, y_ptr): + 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): @@ -1350,7 +1361,7 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): else: - def binary_wrapper(z_ptr, x_ptr, y_ptr): + 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): @@ -1367,8 +1378,17 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): try: numba_func.compile(sig) except numba.TypingError: + 1 / 0 raise TypeError("TODO") 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 + 1 / 0 + else: + 1 / 0 # Numba is unable to handle BOOL correctly right now, but we have a workaround # See: https://github.com/numba/numba/issues/5395 @@ -1383,19 +1403,19 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): ) if ret_type == "BOOL": - def binary_wrapper(z_ptr, x_ptr, y_ptr): + def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover z = numba.carray(z_ptr, 1) x = numba.carray(x_ptr, 1) y = numba.carray(y_ptr, 1) - z[0] = bool(numba_func(x[0], y[0])) # pragma: no cover + z[0] = bool(numba_func(x[0], y[0])) else: - def binary_wrapper(z_ptr, x_ptr, y_ptr): + def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover z = numba.carray(z_ptr, 1) x = numba.carray(x_ptr, 1) y = numba.carray(y_ptr, 1) - z[0] = numba_func(x[0], y[0]) # pragma: no cover + z[0] = numba_func(x[0], y[0]) binary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(binary_wrapper) new_binary = ffi_new("GrB_BinaryOp*") @@ -1697,8 +1717,10 @@ def _build(cls, name, binaryop, identity, *, anonymous=False): def _getitem_udt(self, dtype, dtype2): if dtype2 is None: + 1 / 0 dtype2 = dtype elif dtype != dtype2: + 1 / 0 raise TypeError("TODO") if dtype in self._udt_types: return self._udt_ops[dtype] @@ -1918,6 +1940,7 @@ def _build(cls, name, monoid, binaryop, *, anonymous=False): def _getitem_udt(self, dtype, dtype2): if dtype2 is None: + 1 / 0 dtype2 = dtype dtypes = (dtype, dtype2) if dtypes in self._udt_types: diff --git a/grblas/tests/test_matrix.py b/grblas/tests/test_matrix.py index 2dcea6fba..501820598 100644 --- a/grblas/tests/test_matrix.py +++ b/grblas/tests/test_matrix.py @@ -3175,9 +3175,7 @@ def test_udt(): assert A.reduce_rowwise(monoid.any).new().isequal(expected) rows, cols, values = A.to_values() assert A.isequal(Matrix.from_values(rows, cols, values)) - - # 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) - + 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) # agg: any_value, first, last, first_index, last_index, count diff --git a/grblas/tests/test_numpyops.py b/grblas/tests/test_numpyops.py index 4531732d0..f5bea3f6d 100644 --- a/grblas/tests/test_numpyops.py +++ b/grblas/tests/test_numpyops.py @@ -58,10 +58,10 @@ def test_npunary(): 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 blocklist.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() @@ -137,7 +137,7 @@ def test_npbinary(): 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 blocklist.get( + if gb_left.dtype not in op.types or binary_name in blocklist.get( gb_left.dtype.name, () ): continue @@ -221,10 +221,10 @@ def test_npmonoid(): 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 blocklist.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) diff --git a/grblas/tests/test_op.py b/grblas/tests/test_op.py index f5b5722bf..c070ef054 100644 --- a/grblas/tests/test_op.py +++ b/grblas/tests/test_op.py @@ -119,7 +119,7 @@ def test_get_typed_op(): 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") @@ -159,7 +159,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 @@ -183,7 +183,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 @@ -229,7 +229,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" @@ -239,7 +239,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 @@ -288,7 +288,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 @@ -305,7 +305,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 @@ -324,7 +324,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 @@ -424,7 +424,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 @@ -451,7 +451,7 @@ 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") @@ -476,7 +476,7 @@ 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") @@ -505,7 +505,7 @@ 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) @@ -536,7 +536,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) @@ -569,7 +569,7 @@ 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") @@ -598,7 +598,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") @@ -684,7 +684,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 @@ -754,7 +757,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 @@ -817,7 +823,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) @@ -959,10 +965,10 @@ def test_from_string(): 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) @@ -974,9 +980,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) @@ -989,7 +995,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 @@ -1024,13 +1030,14 @@ def test_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 + return val # pragma: no cover udt_identity = UnaryOp.register_new("udt_identity", _udt_identity, is_udt=True) assert udt in udt_identity @@ -1039,7 +1046,7 @@ def _udt_identity(val): assert result.isequal(v) def _udt_getx(val): - return val["x"] + return val["x"] # pragma: no cover udt_getx = UnaryOp.register_anonymous(_udt_getx, "udt_getx", is_udt=True) assert udt in udt_getx @@ -1048,7 +1055,7 @@ def _udt_getx(val): assert result.isequal(expected) def _udt_first(x, y): - return x + return x # pragma: no cover udt_first = BinaryOp.register_anonymous(_udt_first, "udt_first", is_udt=True) assert udt in udt_first diff --git a/grblas/tests/test_pickle.py b/grblas/tests/test_pickle.py index e986ecfc0..1de6d3727 100644 --- a/grblas/tests/test_pickle.py +++ b/grblas/tests/test_pickle.py @@ -7,19 +7,19 @@ 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 +182,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 diff --git a/grblas/tests/test_vector.py b/grblas/tests/test_vector.py index 2f82dad4f..195c9471a 100644 --- a/grblas/tests/test_vector.py +++ b/grblas/tests/test_vector.py @@ -1672,6 +1672,7 @@ def test_delete_via_scalar(v): 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] @@ -1713,10 +1714,13 @@ def test_udt(): 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 = 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 From 77433c42f70f669c89382146ab5af2d0709583d9 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 29 Mar 2022 23:49:40 -0500 Subject: [PATCH 06/14] Aggregators on UDT --- grblas/_agg.py | 33 ++++++++++++++++++++++++--------- grblas/dtypes.py | 14 +++++++++++++- grblas/operator.py | 28 +++++++++++++++++++--------- grblas/tests/test_vector.py | 14 +++++++------- 4 files changed, 63 insertions(+), 26 deletions(-) diff --git a/grblas/_agg.py b/grblas/_agg.py index e0876dd6a..4be9fbb82 100644 --- a/grblas/_agg.py +++ b/grblas/_agg.py @@ -4,7 +4,8 @@ import numpy as np from . import agg, binary, monoid, semiring, unary -from .dtypes import lookup_dtype, unify +from .dtypes import INT64, lookup_dtype, unify +from .operator import get_typed_op from .scalar import Scalar @@ -42,6 +43,7 @@ def __init__( composite=None, custom=None, types=None, + any_dtype=None, ): self.name = name self._initval_orig = initval @@ -67,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,7 +85,7 @@ def types(self): def __getitem__(self, dtype): dtype = lookup_dtype(dtype) - if dtype not in self.types: + 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) @@ -90,7 +93,7 @@ def __getitem__(self, dtype): def __contains__(self, dtype): dtype = lookup_dtype(dtype) - return dtype in self.types + return self._any_dtype or dtype in self.types def __repr__(self): return f"agg.{self.name}" @@ -106,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}]" @@ -159,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] @@ -241,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 ) @@ -263,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( @@ -563,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, ) @@ -600,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/dtypes.py b/grblas/dtypes.py index d2ae847c0..56fa1c9fb 100644 --- a/grblas/dtypes.py +++ b/grblas/dtypes.py @@ -36,6 +36,7 @@ def __hash__(self): return hash(self.gb_obj) def __lt__(self, other): + 1 / 0 # Let us sort for prettier error reporting if type(other) is DataType: return str(self.np_type) < str(other.np_type) @@ -48,6 +49,7 @@ def __lt__(self, other): def __reduce__(self): if self._is_udt: + 1 / 0 return (self._deserialize, (self.name, self.np_type, self._is_anonymous)) if self.gb_name == "GrB_Index": return "_INDEX" @@ -67,6 +69,7 @@ def _is_udt(self): @staticmethod def _deserialize(name, dtype, is_anonymous): + 1 / 0 if is_anonymous: return register_anonymous(dtype, name) if name in _registry: @@ -76,8 +79,10 @@ def _deserialize(name, dtype, is_anonymous): def register_new(name, dtype): if not name.isidentifier(): + 1 / 0 raise ValueError(f"`name` argument must be a valid Python identifier; got: {name!r}") if name in _registry or name in globals(): + 1 / 0 raise ValueError(f"{name!r} name for dtype is unavailable") rv = register_anonymous(dtype, name) _registry[name] = rv @@ -90,10 +95,13 @@ def register_anonymous(dtype, name=None): dtype = np.dtype(dtype) if dtype.hasobject: + 1 / 0 raise ValueError("dtype must not allow Python objects") if dtype.isbuiltin != 0: + 1 / 0 raise ValueError("dtype must not be a builtin type") if name is None: + 1 / 0 name = repr(dtype) numba_type = numba.typeof(dtype).dtype @@ -233,7 +241,11 @@ def lookup_dtype(key, value=None): 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/operator.py b/grblas/operator.py index c5e8620a6..89095f6dd 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -1336,6 +1336,7 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover z_ptr[0] = (x[mask] == y[mask]).all() elif self.name == "ne" and not self._anonymous: + 1 / 0 assert dtype == dtype2 # XXX: must be same size? itemsize = dtype.np_type.itemsize mask = _udt_mask(dtype.np_type) @@ -1387,8 +1388,6 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover elif ret_type.gb_obj is dtype2.gb_obj: ret_type = dtype2 1 / 0 - else: - 1 / 0 # Numba is unable to handle BOOL correctly right now, but we have a workaround # See: https://github.com/numba/numba/issues/5395 @@ -1571,7 +1570,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: + 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 @@ -2155,9 +2154,13 @@ def __init__(self, name, monoid=None, binaryop=None, *, anonymous=False): super().__init__(name, anonymous=anonymous) self._monoid = monoid self._binaryop = binaryop - if binaryop is not None and binaryop._is_udt: - self._udt_types = {} # {(dtype, dtype): DataType} - self._udt_ops = {} # {(dtype, dtype): TypedUserSemiring} + 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: @@ -2199,9 +2202,16 @@ def get_typed_op(op, dtype, dtype2=None, *, is_left_scalar=False, is_right_scala if op._is_udt: return op._getitem_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: + if op.is_positional: + return op[UINT64] + if op._udt_types is None: + raise + return op._getitem_udt(dtype, dtype2) return op[dtype] elif isinstance(op, ParameterizedUdf): op = op() # Use default parameters of parameterized UDFs diff --git a/grblas/tests/test_vector.py b/grblas/tests/test_vector.py index 195c9471a..a7fd94ad4 100644 --- a/grblas/tests/test_vector.py +++ b/grblas/tests/test_vector.py @@ -1721,10 +1721,10 @@ def test_udt(): 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.first_index).new() == 0 - # assert v.reduce(agg.last_index).new() == 2 - # assert v.reduce(agg.count).new() == 3 + 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 From 0b898ab4a60c10718c2120f74abb0c1e58099d57 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 30 Mar 2022 16:10:42 -0500 Subject: [PATCH 07/14] Clean up usage of `lookup_dtype` and `unify` Some usage of `lookup_dtype` is not unnecessary now that ops use dtypes instead of dtype names. Also, using `get_typed_op` is preferred over `unify` where possible. --- grblas/_agg.py | 6 +++--- grblas/base.py | 2 +- grblas/dtypes.py | 28 ++++++++++++++-------------- grblas/matrix.py | 6 +++--- grblas/operator.py | 21 +++++++++++---------- grblas/vector.py | 6 +++--- 6 files changed, 35 insertions(+), 34 deletions(-) diff --git a/grblas/_agg.py b/grblas/_agg.py index 4be9fbb82..d45052be4 100644 --- a/grblas/_agg.py +++ b/grblas/_agg.py @@ -4,7 +4,7 @@ import numpy as np from . import agg, binary, monoid, semiring, unary -from .dtypes import INT64, lookup_dtype, unify +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) 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 diff --git a/grblas/base.py b/grblas/base.py index da1aa1459..61e753f25 100644 --- a/grblas/base.py +++ b/grblas/base.py @@ -547,7 +547,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/dtypes.py b/grblas/dtypes.py index 56fa1c9fb..337811852 100644 --- a/grblas/dtypes.py +++ b/grblas/dtypes.py @@ -155,23 +155,23 @@ def register_anonymous(dtype, name=None): # 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)), } ) @@ -203,7 +203,7 @@ def register_anonymous(dtype, name=None): _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 diff --git a/grblas/matrix.py b/grblas/matrix.py index bb7f04adb..fe40a1787 100644 --- a/grblas/matrix.py +++ b/grblas/matrix.py @@ -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 diff --git a/grblas/operator.py b/grblas/operator.py index 89095f6dd..df2938677 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -131,8 +131,8 @@ class TypedOpBase: def __init__(self, parent, name, type_, return_type, gb_obj, gb_name): self.parent = parent self.name = name - self.type = lookup_dtype(type_) - self.return_type = lookup_dtype(return_type) + self.type = type_ + self.return_type = return_type self.gb_obj = gb_obj self.gb_name = gb_name @@ -743,9 +743,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}") @@ -758,7 +758,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) @@ -816,7 +821,6 @@ def _build(cls, name, func, *, anonymous=False, is_udt=False): nt = numba.types if not is_udt: for type_, sample_val in _sample_values.items(): - type_ = lookup_dtype(type_) sig = (type_.numba_type,) try: unary_udf.compile(sig) @@ -1213,7 +1217,6 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): nt = numba.types if not is_udt: for type_, sample_val in _sample_values.items(): - type_ = lookup_dtype(type_) sig = (type_.numba_type, type_.numba_type) try: binary_udf.compile(sig) @@ -1686,10 +1689,9 @@ def _build(cls, name, binaryop, identity, *, anonymous=False): identities = dict.fromkeys(binaryop.types, identity) explicit_identities = False else: - identities = identity + identities = {lookup_dtype(key): val for key, val in identity.items()} 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. @@ -1917,7 +1919,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), diff --git a/grblas/vector.py b/grblas/vector.py index 9c3b905bf..1698da77c 100644 --- a/grblas/vector.py +++ b/grblas/vector.py @@ -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 From 6e92f4f61b525459cef2b28ed70a2a6c5e6df9d3 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 7 Apr 2022 00:22:25 -0500 Subject: [PATCH 08/14] A little bit of progress --- grblas/dtypes.py | 15 +++++-------- grblas/formatting.py | 12 +++++++++-- grblas/operator.py | 35 +++++++++++++++---------------- grblas/scalar.py | 16 ++++++++------ grblas/tests/test_dtype.py | 42 ++++++++++++++++++++++++++++++++++++- grblas/tests/test_op.py | 2 ++ grblas/tests/test_scalar.py | 13 ++++++++---- grblas/tests/test_vector.py | 26 +++++++++++++++++++++++ 8 files changed, 120 insertions(+), 41 deletions(-) diff --git a/grblas/dtypes.py b/grblas/dtypes.py index 337811852..fb05f92a5 100644 --- a/grblas/dtypes.py +++ b/grblas/dtypes.py @@ -17,7 +17,7 @@ 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 @@ -36,14 +36,14 @@ def __hash__(self): return hash(self.gb_obj) def __lt__(self, other): - 1 / 0 # Let us sort for prettier error reporting if type(other) is DataType: - return str(self.np_type) < str(other.np_type) + 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: - other = lookup_dtype(other) - return str(self.np_type) < str(other.np_type) + return self < lookup_dtype(other) except ValueError: raise TypeError(f"Invalid or unknown datatype: {other}") from None @@ -79,10 +79,8 @@ def _deserialize(name, dtype, is_anonymous): def register_new(name, dtype): if not name.isidentifier(): - 1 / 0 raise ValueError(f"`name` argument must be a valid Python identifier; got: {name!r}") if name in _registry or name in globals(): - 1 / 0 raise ValueError(f"{name!r} name for dtype is unavailable") rv = register_anonymous(dtype, name) _registry[name] = rv @@ -95,13 +93,10 @@ def register_anonymous(dtype, name=None): dtype = np.dtype(dtype) if dtype.hasobject: - 1 / 0 raise ValueError("dtype must not allow Python objects") if dtype.isbuiltin != 0: - 1 / 0 raise ValueError("dtype must not be a builtin type") if name is None: - 1 / 0 name = repr(dtype) numba_type = numba.typeof(dtype).dtype diff --git a/grblas/formatting.py b/grblas/formatting.py index 121b976da..919108a71 100644 --- a/grblas/formatting.py +++ b/grblas/formatting.py @@ -135,8 +135,12 @@ 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: + 1 / 0 + vals = vals.tolist() df.values[rows, cols] = vals - if np.issubdtype(vals.dtype, np.inexact): + if np.issubdtype(np_type, np.inexact): nulls = np.isnan(vals) df.values[rows[nulls], cols[nulls]] = "nan" @@ -170,8 +174,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: + 1 / 0 + vals = vals.tolist() df.values[0, cols] = vals - if np.issubdtype(vals.dtype, np.inexact): + if np.issubdtype(np_type, np.inexact): df.values[0, cols[np.isnan(vals)]] = "nan" diff --git a/grblas/operator.py b/grblas/operator.py index 7e7889fe8..722feae6d 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -103,7 +103,6 @@ def _udt_mask(dtype): if dtype in _udt_mask_cache: return _udt_mask_cache[dtype] if dtype.subdtype is not None: - 1 / 0 mask = _udt_mask(dtype.subdtype[0]) N = reduce(mul, dtype.subdtype[1]) rv = np.concatenate([mask] * N) @@ -296,11 +295,10 @@ def __init__(self, parent, name, type_, return_type, gb_obj): @property def orig_func(self): - 1 / 0 - return self.parent._func + return self.parent.orig_func @property - def numba_func(self): + def _numba_func(self): return self.parent._numba_func __call__ = TypedBuiltinUnaryOp.__call__ @@ -327,7 +325,7 @@ def monoid(self): is_commutative = TypedBuiltinBinaryOp.is_commutative __call__ = TypedBuiltinBinaryOp.__call__ orig_func = TypedUserUnaryOp.orig_func - numba_func = TypedUserUnaryOp.numba_func + _numba_func = TypedUserUnaryOp._numba_func class TypedUserMonoid(TypedOpBase): @@ -782,7 +780,7 @@ def _deserialize(cls, name, *args): class UnaryOp(OpBase): - __slots__ = "_func", "is_positional", "_is_udt", "_numba_func" + __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func" _module = unary _modname = "unary" _typed_class = TypedBuiltinUnaryOp @@ -1042,7 +1040,7 @@ def __init__( 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 @@ -1052,13 +1050,13 @@ def __init__( 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__ @@ -1107,7 +1105,7 @@ class BinaryOp(OpBase): "_monoid", "_commutes_to", "_semiring_commutes_to", - "_func", + "orig_func", "is_positional", "_is_udt", "_numba_func", @@ -1414,6 +1412,7 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover nt.CPointer(dtype.numba_type), nt.CPointer(dtype2.numba_type), ) + if ret_type == "BOOL": def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover @@ -1614,7 +1613,7 @@ def _initialize(cls): (binary.pair, _pair), (binary.any, _first), ]: - binop._func = func + binop.orig_func = func binop._numba_func = numba.njit(func) binop._udt_types = {} binop._udt_ops = {} @@ -1637,7 +1636,7 @@ def __init__( 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 @@ -1647,13 +1646,13 @@ def __init__( 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 diff --git a/grblas/scalar.py b/grblas/scalar.py index 5c981be64..fe8821db4 100644 --- a/grblas/scalar.py +++ b/grblas/scalar.py @@ -198,7 +198,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: @@ -246,9 +246,13 @@ def value(self, val): if output_type(val) is Scalar: val = val.value # raise below if wrong type (as determined by cffi) if self.dtype._is_udt: - self.gb_obj[0 : self.dtype.np_type.itemsize] = np.array( - [val], dtype=self.dtype.np_type - ).view(np.uint8) + 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 @@ -261,7 +265,6 @@ def value(self, val): val = _as_scalar(val, is_cscalar=True) dtype_name = val.dtype.name call(f"GrB_Scalar_setElement_{dtype_name}", [self, val]) - # call(f"GrB_Scalar_setElement_{val.dtype}", [self, val]) @property def nvals(self): @@ -306,7 +309,8 @@ 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) + new_scalar.value = self.value + # new_scalar.value = new_scalar.dtype.np_type(self.value) return new_scalar def wait(self): 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_op.py b/grblas/tests/test_op.py index 28ed0da86..9d1b5765d 100644 --- a/grblas/tests/test_op.py +++ b/grblas/tests/test_op.py @@ -123,6 +123,8 @@ def plus_one(x): 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, diff --git a/grblas/tests/test_scalar.py b/grblas/tests/test_scalar.py index e08f41d3f..926b9b223 100644 --- a/grblas/tests/test_scalar.py +++ b/grblas/tests/test_scalar.py @@ -48,10 +48,15 @@ def test_dup(s): ("INT8", -2), ("INT16", -2), ("INT32", -2), - ("UINT8", 2**8 - 2), - ("UINT16", 2**16 - 2), - ("UINT32", 2**32 - 2), - ("UINT64", 2**64 - 2), + # Do we want these to wrap or be clipped? Why does it clip? + # ("UINT8", 2**8 - 2), + # ("UINT16", 2**16 - 2), + # ("UINT32", 2**32 - 2), + # ("UINT64", 2**64 - 2), + ("UINT8", 0), + ("UINT16", 0), + ("UINT32", 0), + ("UINT64", 0), ("BOOL", True), ("FP32", -2.5), ]: diff --git a/grblas/tests/test_vector.py b/grblas/tests/test_vector.py index 696d1cbe0..6c9d92c48 100644 --- a/grblas/tests/test_vector.py +++ b/grblas/tests/test_vector.py @@ -1752,6 +1752,32 @@ def test_udt(): assert v.reduce(agg.last_index).new() == 2 assert v.reduce(agg.count).new() == 3 + # arrays as dtypes! + np_dtype = np.dtype("(3,)uint16") + udt = dtypes.register_anonymous(np_dtype, "has_subdtype") + s = Scalar.new(udt) + s.value = [0, 0, 0] + + v = Vector.new(udt, size=2) + v[0] = 0 + w = Vector.new(udt, size=2) + w[0] = (0, 0, 0) + assert v.isequal(w) + w[0] = np.array([0, 0, 1], dtype=np.uint8) + assert not v.isequal(w) + 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)) + + binary.eq[v.dtype] + # binary.ne[v.dtype] + # binary.any[v.dtype] + # monoid.any[v.dtype] + # binary.eq(v & w).new() + def test_infix_outer(): v = Vector.new(int, 2) From 37caac4b8cb6d79e469d58be33834202a0d76b9f Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 7 Apr 2022 11:40:33 -0500 Subject: [PATCH 09/14] Fix failing test --- grblas/scalar.py | 7 +++++-- grblas/tests/test_scalar.py | 33 ++++++++++++++++++++------------- 2 files changed, 25 insertions(+), 15 deletions(-) diff --git a/grblas/scalar.py b/grblas/scalar.py index fe8821db4..c1cca5e46 100644 --- a/grblas/scalar.py +++ b/grblas/scalar.py @@ -309,8 +309,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 = self.value - # 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): diff --git a/grblas/tests/test_scalar.py b/grblas/tests/test_scalar.py index 926b9b223..358dccb1f 100644 --- a/grblas/tests/test_scalar.py +++ b/grblas/tests/test_scalar.py @@ -44,28 +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), - # Do we want these to wrap or be clipped? Why does it clip? - # ("UINT8", 2**8 - 2), - # ("UINT16", 2**16 - 2), - # ("UINT32", 2**32 - 2), - # ("UINT64", 2**64 - 2), - ("UINT8", 0), - ("UINT16", 0), - ("UINT32", 0), - ("UINT64", 0), ("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 From 6eeabfaea4795dd611794f92447157cb866bc869 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 7 Apr 2022 15:48:46 -0500 Subject: [PATCH 10/14] Get UDFs to work on array UDTs --- grblas/operator.py | 97 ++++++++++++++++++++++++++----------- grblas/scalar.py | 6 ++- grblas/tests/test_vector.py | 14 ++++-- 3 files changed, 82 insertions(+), 35 deletions(-) diff --git a/grblas/operator.py b/grblas/operator.py index 1f6a82b8f..f03964fbf 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -1135,6 +1135,71 @@ 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])" # is this still necessary? + 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])" # is this still necessary? + 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", @@ -1383,7 +1448,6 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover z_ptr[0] = (x[mask] == y[mask]).all() elif self.name == "ne" and not self._anonymous: - 1 / 0 assert dtype == dtype2 # XXX: must be same size? itemsize = dtype.np_type.itemsize mask = _udt_mask(dtype.np_type) @@ -1435,34 +1499,7 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover elif ret_type.gb_obj is dtype2.gb_obj: ret_type = dtype2 1 / 0 - - # 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 - return_type = INT8 if ret_type == "BOOL" else ret_type - - # Build wrapper because GraphBLAS wants pointers and void return - wrapper_sig = nt.void( - nt.CPointer(return_type.numba_type), - nt.CPointer(dtype.numba_type), - nt.CPointer(dtype2.numba_type), - ) - - if ret_type == "BOOL": - - def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover - z = numba.carray(z_ptr, 1) - x = numba.carray(x_ptr, 1) - y = numba.carray(y_ptr, 1) - z[0] = bool(numba_func(x[0], y[0])) - - else: - - def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover - z = numba.carray(z_ptr, 1) - x = numba.carray(x_ptr, 1) - y = numba.carray(y_ptr, 1) - z[0] = numba_func(x[0], y[0]) + 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*") @@ -1655,6 +1692,8 @@ def _initialize(cls): 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__( diff --git a/grblas/scalar.py b/grblas/scalar.py index c1cca5e46..19364784f 100644 --- a/grblas/scalar.py +++ b/grblas/scalar.py @@ -234,7 +234,11 @@ def value(self): call(f"GrB_Scalar_extractElement_{dtype_name}", [_Pointer(scalar), self]) if is_udt: np_type = self.dtype.np_type - return np.array(ffi.buffer(scalar.gb_obj[0 : np_type.itemsize])).view(np_type)[0] + 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] diff --git a/grblas/tests/test_vector.py b/grblas/tests/test_vector.py index 7bc5ef1e7..0aaaf160f 100644 --- a/grblas/tests/test_vector.py +++ b/grblas/tests/test_vector.py @@ -1763,8 +1763,13 @@ def test_udt(): w = Vector.new(udt, 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) @@ -1772,11 +1777,10 @@ def test_udt(): indices, values = v.to_values() assert_array_equal(values, np.ones((2, 3), dtype=np.uint16)) - binary.eq[v.dtype] - # binary.ne[v.dtype] - # binary.any[v.dtype] - # monoid.any[v.dtype] - # binary.eq(v & w).new() + s = v.reduce(monoid.any).new() + assert (s.value == [1, 1, 1]).all() + assert (s.value == 1).all() + # print(s == 1) # TODO: decide what to do with scalars with array dtypes def test_infix_outer(): From c105b92962042c14a4f430a30947a4a1209ebe82 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 8 Apr 2022 18:33:20 -0500 Subject: [PATCH 11/14] Coverage --- README.md | 6 ---- docs/index.rst | 8 ------ grblas/base.py | 1 - grblas/binary/numpy.py | 1 + grblas/formatting.py | 2 -- grblas/monoid/numpy.py | 1 + grblas/operator.py | 36 ++++++++++++----------- grblas/scalar.py | 6 +++- grblas/tests/test_formatting.py | 51 ++++++++++++++++++++++++++++++++- grblas/tests/test_op.py | 43 +++++++++++++++++++++++++++ grblas/tests/test_vector.py | 9 +++++- grblas/unary/numpy.py | 1 + 12 files changed, 128 insertions(+), 37 deletions(-) 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/base.py b/grblas/base.py index 5f381888f..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: diff --git a/grblas/binary/numpy.py b/grblas/binary/numpy.py index e832583e2..72045d646 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()} # Not included: maximum, minimum, gcd, hypot, logaddexp, logaddexp2 # lcm, left_shift, nextafter, right_shift diff --git a/grblas/formatting.py b/grblas/formatting.py index 919108a71..331e3218f 100644 --- a/grblas/formatting.py +++ b/grblas/formatting.py @@ -137,7 +137,6 @@ def _update_matrix_dataframe(df, matrix, rows, row_offset, columns, column_offse rows, cols, vals = submatrix.to_values() np_type = submatrix.dtype.np_type if submatrix.dtype._is_udt and np_type.subdtype is not None: - 1 / 0 vals = vals.tolist() df.values[rows, cols] = vals if np.issubdtype(np_type, np.inexact): @@ -176,7 +175,6 @@ def _update_vector_dataframe(df, vector, columns, column_offset, *, mask=None): cols, vals = subvector.to_values() np_type = subvector.dtype.np_type if subvector.dtype._is_udt and np_type.subdtype is not None: - 1 / 0 vals = vals.tolist() df.values[0, cols] = vals if np.issubdtype(np_type, np.inexact): diff --git a/grblas/monoid/numpy.py b/grblas/monoid/numpy.py index f5949de00..84d25ff9a 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()} # Not included: maximum, minimum, gcd, hypot, logaddexp, logaddexp2 diff --git a/grblas/operator.py b/grblas/operator.py index f03964fbf..99ca75e2e 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -662,7 +662,6 @@ def __getitem__(self, type_): return self._typed_ops[type_] # This is a UDT or is able to operate on UDTs such as `first` any `any` if type(type_) is tuple: - 1 / 0 dtype1, dtype2 = type_ dtype1 = lookup_dtype(dtype1) dtype2 = lookup_dtype(dtype2) @@ -686,17 +685,9 @@ def __contains__(self, type_): return True elif self._udt_types is None: return False - if type(type_) is tuple: - 1 / 0 - dtype1, dtype2 = type_ - dtype1 = lookup_dtype(dtype1) - dtype2 = lookup_dtype(dtype2) - else: - dtype1 = dtype2 = lookup_dtype(type_) try: - self._getitem_udt(dtype1, dtype2) - except TypeError: - 1 / 0 + self[type_] + except (TypeError, KeyError): return False else: return True @@ -814,6 +805,14 @@ def _deserialize(cls, name, *args): return cls.register_new(name, *args) +def _identity(x): + return x + + +def _one(x): + return 1 + + class UnaryOp(OpBase): __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func" _module = unary @@ -939,7 +938,6 @@ def _getitem_udt(self, dtype, dtype2): try: numba_func.compile(sig) except numba.TypingError: - 1 / 0 raise TypeError("TODO") 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: @@ -1062,6 +1060,15 @@ def _initialize(cls): 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__( @@ -1403,7 +1410,6 @@ def binary_wrapper(z, x, y): def _getitem_udt(self, dtype, dtype2): if dtype2 is None: - 1 / 0 dtype2 = dtype dtypes = (dtype, dtype2) if dtypes in self._udt_types: @@ -1490,7 +1496,6 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover try: numba_func.compile(sig) except numba.TypingError: - 1 / 0 raise TypeError("TODO") ret_type = lookup_dtype(numba_func.overloads[sig].signature.return_type) if ret_type is not dtype and ret_type is not dtype2: @@ -1805,10 +1810,8 @@ def _build(cls, name, binaryop, identity, *, anonymous=False): def _getitem_udt(self, dtype, dtype2): if dtype2 is None: - 1 / 0 dtype2 = dtype elif dtype != dtype2: - 1 / 0 raise TypeError("TODO") if dtype in self._udt_types: return self._udt_ops[dtype] @@ -2027,7 +2030,6 @@ def _build(cls, name, monoid, binaryop, *, anonymous=False): def _getitem_udt(self, dtype, dtype2): if dtype2 is None: - 1 / 0 dtype2 = dtype dtypes = (dtype, dtype2) if dtypes in self._udt_types: diff --git a/grblas/scalar.py b/grblas/scalar.py index 19364784f..7944b0259 100644 --- a/grblas/scalar.py +++ b/grblas/scalar.py @@ -158,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): """ 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_op.py b/grblas/tests/test_op.py index 9d1b5765d..2ab9e0b00 100644 --- a/grblas/tests/test_op.py +++ b/grblas/tests/test_op.py @@ -1048,6 +1048,14 @@ def _udt_identity(val): 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 @@ -1063,17 +1071,52 @@ def _udt_first(x, y): 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 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): + 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): + return BreakCompile(x) + + badbinary = BinaryOp.register_anonymous(badfunc, is_udt=True) + assert udt not in badbinary + assert int not in badbinary + def test_dir(): for mod in [unary, binary, monoid, semiring, op]: diff --git a/grblas/tests/test_vector.py b/grblas/tests/test_vector.py index 0aaaf160f..8299db12b 100644 --- a/grblas/tests/test_vector.py +++ b/grblas/tests/test_vector.py @@ -1780,7 +1780,14 @@ def test_udt(): s = v.reduce(monoid.any).new() assert (s.value == [1, 1, 1]).all() assert (s.value == 1).all() - # print(s == 1) # TODO: decide what to do with scalars with array dtypes + 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 def test_infix_outer(): diff --git a/grblas/unary/numpy.py b/grblas/unary/numpy.py index 0cc87289f..a0750b38c 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()} _operator._STANDARD_OPERATOR_NAMES.update(f"unary.numpy.{name}" for name in _unary_names) __all__ = list(_unary_names) From 1ce84217b3818555288f373c829bfd5aebfd6406 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 8 Apr 2022 22:06:48 -0500 Subject: [PATCH 12/14] More coverage; better handling of array subdtypes --- grblas/_ss/matrix.py | 32 ++++++++++++++++++-------------- grblas/_ss/vector.py | 4 +++- grblas/dtypes.py | 5 +---- grblas/matrix.py | 9 ++++++--- grblas/operator.py | 17 +++++++---------- grblas/tests/pickle3.pkl | Bin 0 -> 770 bytes grblas/tests/test_pickle.py | 36 ++++++++++++++++++++++++++++++++++++ grblas/tests/test_vector.py | 2 ++ grblas/utils.py | 10 ++++++++-- grblas/vector.py | 9 ++++++--- 10 files changed, 87 insertions(+), 37 deletions(-) create mode 100644 grblas/tests/pickle3.pkl 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/dtypes.py b/grblas/dtypes.py index fb05f92a5..6787313b6 100644 --- a/grblas/dtypes.py +++ b/grblas/dtypes.py @@ -27,8 +27,7 @@ def __eq__(self, other): return self is other # Attempt to use `other` as a lookup key try: - other = lookup_dtype(other) - return self is other + return self is lookup_dtype(other) except ValueError: raise TypeError(f"Invalid or unknown datatype: {other}") from None @@ -49,7 +48,6 @@ def __lt__(self, other): def __reduce__(self): if self._is_udt: - 1 / 0 return (self._deserialize, (self.name, self.np_type, self._is_anonymous)) if self.gb_name == "GrB_Index": return "_INDEX" @@ -69,7 +67,6 @@ def _is_udt(self): @staticmethod def _deserialize(name, dtype, is_anonymous): - 1 / 0 if is_anonymous: return register_anonymous(dtype, name) if name in _registry: diff --git a/grblas/matrix.py b/grblas/matrix.py index dcaf9e0fb..63dc4c08d 100644 --- a/grblas/matrix.py +++ b/grblas/matrix.py @@ -274,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: " @@ -386,7 +386,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: @@ -396,8 +396,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( diff --git a/grblas/operator.py b/grblas/operator.py index 99ca75e2e..c1395ba03 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -687,7 +687,7 @@ def __contains__(self, type_): return False try: self[type_] - except (TypeError, KeyError): + except (TypeError, KeyError, numba.NumbaError): return False else: return True @@ -935,10 +935,7 @@ def _getitem_udt(self, dtype, dtype2): numba_func = self._numba_func sig = (dtype.numba_type,) - try: - numba_func.compile(sig) - except numba.TypingError: - raise TypeError("TODO") + 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 @@ -1493,10 +1490,7 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover else: numba_func = self._numba_func sig = (dtype.numba_type, dtype2.numba_type) - try: - numba_func.compile(sig) - except numba.TypingError: - raise TypeError("TODO") + 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: @@ -1812,7 +1806,10 @@ def _getitem_udt(self, dtype, dtype2): if dtype2 is None: dtype2 = dtype elif dtype != dtype2: - raise TypeError("TODO") + 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._getitem_udt(dtype, dtype2) diff --git a/grblas/tests/pickle3.pkl b/grblas/tests/pickle3.pkl new file mode 100644 index 0000000000000000000000000000000000000000..1bdedd7742d08a58dbbcc799f7abcba2d4ab7c22 GIT binary patch literal 770 zcmZuv&x_MQ6wah+TJ{Glir{Ha9u^OwqM`>ym+q;<29dSbFiE2?m^Ni5sjlpTmsSJs zu&~dAf2qgyO;YLNb|A@P^3C_Y@0<7C`}N0lnsrG>xbrScKIiG=!^s)p{wE#h$^?`0 ze4ZM@2PdjjXIux-BsHndRGxiH>6$=bcUDd3=Y-t`%zNmLUsC<^2fe2w*+FqZAqqt( zVp8uF$Ei!wyL85p?Td>}WD%qFb=!!Y`aBZ98 z!MwE;nJR9@fd_|{^0q+OwL8_VCh%}e3wO!t!_anV^fmraEH+^+(xfaj;r($_{2!GK z*^jews^%PM6IzfIIt`{(US>*b#UbuRI!%h1DRq^UHs3DxW|Q{kHA!blLyobHR>=O6 z;9foK3-SAAGwF&zl^c>Sdd6N77k*}<%oJ3k|FO_hRT2!)XIRCT3s2#>C59LO zXxS?4)WZWi8)E%*_-pAa>@1=IN5T1ut{H7oAUbn%5v v*zrn=#tyux;Vt&s1U13&T{>5~JT#A3x|MB*+sgLV(!qzvEdO*ZzY6{W6W}4w literal 0 HcmV?d00001 diff --git a/grblas/tests/test_pickle.py b/grblas/tests/test_pickle.py index 1de6d3727..4d1e5baf4 100644 --- a/grblas/tests/test_pickle.py +++ b/grblas/tests/test_pickle.py @@ -1,6 +1,7 @@ import os import pickle +import numpy as np import pytest import grblas as gb @@ -273,3 +274,38 @@ 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) diff --git a/grblas/tests/test_vector.py b/grblas/tests/test_vector.py index 8299db12b..6506e8f08 100644 --- a/grblas/tests/test_vector.py +++ b/grblas/tests/test_vector.py @@ -1776,6 +1776,8 @@ def test_udt(): 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)) s = v.reduce(monoid.any).new() assert (s.value == [1, 1, 1]).all() diff --git a/grblas/utils.py b/grblas/utils.py index d297d911a..c22dace4e 100644 --- a/grblas/utils.py +++ b/grblas/utils.py @@ -67,10 +67,16 @@ 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) @@ -127,7 +133,7 @@ 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.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 diff --git a/grblas/vector.py b/grblas/vector.py index cce98ab6b..1202340af 100644 --- a/grblas/vector.py +++ b/grblas/vector.py @@ -260,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}" @@ -355,14 +355,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: + # 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( From c9e8772946e812216ef93e5abf4fb283bc8f4dd0 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 8 Apr 2022 23:43:44 -0500 Subject: [PATCH 13/14] More coverage. Anything else needed? --- grblas/operator.py | 53 +++++++++---------------------------- grblas/tests/test_matrix.py | 25 +++++++++++++++++ grblas/tests/test_op.py | 22 ++++++++++++--- grblas/tests/test_vector.py | 14 +++++++--- 4 files changed, 67 insertions(+), 47 deletions(-) diff --git a/grblas/operator.py b/grblas/operator.py index c1395ba03..43c43b38a 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -119,7 +119,7 @@ def _udt_mask(dtype): rv = np.concatenate(masks) else: rv = np.ones(dtype.itemsize, dtype=bool) - # assert rv.size == dtype.itemsize, (rv.size, dtype.itemsize) + # assert rv.size == dtype.itemsize _udt_mask_cache[dtype] = rv return rv @@ -215,16 +215,14 @@ 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] - # TODO: what about UDTs and `gb.binary.any`? @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] - # TODO: what about UDTs and `gb.binary.any`? @property def is_commutative(self): @@ -475,7 +473,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!) @@ -806,11 +804,11 @@ def _deserialize(cls, name, *args): def _identity(x): - return x + return x # pragma: no cover def _one(x): - return 1 + return 1 # pragma: no cover class UnaryOp(OpBase): @@ -940,31 +938,7 @@ def _getitem_udt(self, dtype, dtype2): if ret_type is not dtype and ret_type.gb_obj is dtype.gb_obj: ret_type = dtype - # 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 - return_type = INT8 if ret_type == BOOL else ret_type - - # Build wrapper because GraphBLAS wants pointers and void return - nt = numba.types - wrapper_sig = nt.void( - nt.CPointer(return_type.numba_type), - nt.CPointer(dtype.numba_type), - ) - if ret_type == BOOL: - - def unary_wrapper(z_ptr, x_ptr): # pragma: no cover - z = numba.carray(z_ptr, 1) - x = numba.carray(x_ptr, 1) - z[0] = bool(numba_func(x[0])) - - else: - - def unary_wrapper(z_ptr, x_ptr): # pragma: no cover - z = numba.carray(z_ptr, 1) - x = numba.carray(x_ptr, 1) - z[0] = numba_func(x[0]) - + 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( @@ -1117,7 +1091,7 @@ def _abssecond(x, y): def _rpow(x, y): - return y**x + return y**x # pragma: no cover def _isclose(rel_tol=1e-7, abs_tol=0.0): @@ -1177,7 +1151,7 @@ def _get_udt_wrapper(numba_func, return_type, dtype, dtype2=None): else: xname = "x_ptr" elif dtype == BOOL: - xname = "bool(x_ptr[0])" # is this still necessary? + xname = "bool(x_ptr[0])" else: xname = "x_ptr[0]" @@ -1190,7 +1164,7 @@ def _get_udt_wrapper(numba_func, return_type, dtype, dtype2=None): else: yname = ", y_ptr" elif dtype2 == BOOL: - yname = ", bool(y_ptr[0])" # is this still necessary? + yname = ", bool(y_ptr[0])" else: yname = ", y_ptr[0]" @@ -1414,7 +1388,7 @@ def _getitem_udt(self, dtype, dtype2): nt = numba.types if self.name == "eq" and not self._anonymous: - assert dtype == dtype2 # XXX: must be same size? + # assert dtype.np_type == dtype2.np_type itemsize = dtype.np_type.itemsize mask = _udt_mask(dtype.np_type) ret_type = BOOL @@ -1451,7 +1425,7 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover z_ptr[0] = (x[mask] == y[mask]).all() elif self.name == "ne" and not self._anonymous: - assert dtype == dtype2 # XXX: must be same size? + # assert dtype.np_type == dtype2.np_type itemsize = dtype.np_type.itemsize mask = _udt_mask(dtype.np_type) ret_type = BOOL @@ -1497,7 +1471,6 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover ret_type = dtype elif ret_type.gb_obj is dtype2.gb_obj: ret_type = dtype2 - 1 / 0 binary_wrapper, wrapper_sig = _get_udt_wrapper(numba_func, ret_type, dtype, dtype2) binary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(binary_wrapper) @@ -2298,7 +2271,7 @@ def get_typed_op(op, dtype, dtype2=None, *, is_left_scalar=False, is_right_scala dtype = unify( dtype, dtype2, is_left_scalar=is_left_scalar, is_right_scalar=is_right_scalar ) - except TypeError: + except (TypeError, AttributeError): if op.is_positional: return op[UINT64] if op._udt_types is None: diff --git a/grblas/tests/test_matrix.py b/grblas/tests/test_matrix.py index 5486c1a7b..02c0aa69f 100644 --- a/grblas/tests/test_matrix.py +++ b/grblas/tests/test_matrix.py @@ -3185,7 +3185,32 @@ def test_udt(): 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) # agg: any_value, first, last, first_index, last_index, count + + 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_op.py b/grblas/tests/test_op.py index 2ab9e0b00..552023b03 100644 --- a/grblas/tests/test_op.py +++ b/grblas/tests/test_op.py @@ -115,6 +115,9 @@ 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(): @@ -726,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] @@ -1076,6 +1082,11 @@ def _udt_first(x, y): 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 @@ -1103,20 +1114,25 @@ def _udt_first(x, y): class BreakCompile: pass - def badfunc(x): + 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): + 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] + def test_dir(): for mod in [unary, binary, monoid, semiring, op]: diff --git a/grblas/tests/test_vector.py b/grblas/tests/test_vector.py index 6506e8f08..b5668da75 100644 --- a/grblas/tests/test_vector.py +++ b/grblas/tests/test_vector.py @@ -1751,16 +1751,19 @@ def test_udt(): 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) # arrays as dtypes! np_dtype = np.dtype("(3,)uint16") - udt = dtypes.register_anonymous(np_dtype, "has_subdtype") - s = Scalar.new(udt) + udt2 = dtypes.register_anonymous(np_dtype, "has_subdtype") + s = Scalar.new(udt2) s.value = [0, 0, 0] - v = Vector.new(udt, size=2) + v = Vector.new(udt2, size=2) v[0] = 0 - w = Vector.new(udt, size=2) + 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() @@ -1778,6 +1781,9 @@ def test_udt(): 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() From 21f483b71cf5d88b615e7b2ca5d5a549f792fa43 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Sat, 9 Apr 2022 13:00:32 -0500 Subject: [PATCH 14/14] A few finishing touches --- grblas/binary/numpy.py | 2 +- grblas/dtypes.py | 1 - grblas/matrix.py | 17 +++---- grblas/monoid/numpy.py | 2 +- grblas/operator.py | 90 ++++++++++++++++++++++-------------- grblas/tests/pickle3.pkl | Bin 770 -> 848 bytes grblas/tests/test_matrix.py | 10 +++- grblas/tests/test_op.py | 6 +++ grblas/tests/test_pickle.py | 4 ++ grblas/tests/test_vector.py | 5 ++ grblas/unary/numpy.py | 2 +- grblas/vector.py | 19 ++++---- 12 files changed, 102 insertions(+), 56 deletions(-) diff --git a/grblas/binary/numpy.py b/grblas/binary/numpy.py index 72045d646..1af2a5574 100644 --- a/grblas/binary/numpy.py +++ b/grblas/binary/numpy.py @@ -89,7 +89,7 @@ "subtract": "minus", "true_divide": "truediv", } -_graphblas_to_numpy = {val: key for key, val in _numpy_to_graphblas.items()} +# _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 6787313b6..6a074a2a1 100644 --- a/grblas/dtypes.py +++ b/grblas/dtypes.py @@ -104,7 +104,6 @@ def register_anonymous(dtype, name=None): 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. - # grb_name probably isn't useful, right? rv = DataType(name, gb_obj, None, f"uint8_t[{dtype.itemsize}]", numba_type, dtype) if dtype not in _registry: _registry[gb_obj] = rv diff --git a/grblas/matrix.py b/grblas/matrix.py index 63dc4c08d..ed6f180bb 100644 --- a/grblas/matrix.py +++ b/grblas/matrix.py @@ -292,16 +292,17 @@ def build(self, rows, columns, values, *, dup_op=None, clear=False, nrows=None, return dup_op_given = dup_op is not None - if self.dtype._is_udt: # XXX - dup_op = None - else: - if not dup_op_given: + if not dup_op_given: + if not self.dtype._is_udt: dup_op = binary.plus - dup_op = get_typed_op(dup_op, self.dtype, kind="binary") - if dup_op.opclass == "Monoid": - dup_op = dup_op.binaryop else: - self._expect_op(dup_op, "BinaryOp", within="build", argname="dup_op") + 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 + else: + self._expect_op(dup_op, "BinaryOp", within="build", argname="dup_op") rows = _CArray(rows) columns = _CArray(columns) diff --git a/grblas/monoid/numpy.py b/grblas/monoid/numpy.py index 84d25ff9a..b9deea329 100644 --- a/grblas/monoid/numpy.py +++ b/grblas/monoid/numpy.py @@ -113,7 +113,7 @@ "logical_xor": "lxor", "multiply": "times", } -_graphblas_to_numpy = {val: key for key, val in _numpy_to_graphblas.items()} +# _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 43c43b38a..ebebae66c 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -125,21 +125,32 @@ def _udt_mask(dtype): 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 = 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): @@ -150,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): @@ -228,6 +242,10 @@ def _semiring_commutes_to(self): 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" @@ -280,6 +298,10 @@ def binaryop(self): def commutes_to(self): return self + @property + def type2(self): + return self.type + class TypedBuiltinSemiring(TypedOpBase): __slots__ = () @@ -318,6 +340,8 @@ def commutes_to(self): def is_commutative(self): return self.binaryop.is_commutative + type2 = TypedBuiltinBinaryOp.type2 + class TypedUserUnaryOp(TypedOpBase): __slots__ = () @@ -341,8 +365,8 @@ class TypedUserBinaryOp(TypedOpBase): __slots__ = "_monoid" opclass = "BinaryOp" - def __init__(self, parent, name, type_, return_type, gb_obj): - super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}") + 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 @@ -356,9 +380,10 @@ def monoid(self): commutes_to = TypedBuiltinBinaryOp.commutes_to _semiring_commutes_to = TypedBuiltinBinaryOp._semiring_commutes_to is_commutative = TypedBuiltinBinaryOp.is_commutative - __call__ = TypedBuiltinBinaryOp.__call__ orig_func = TypedUserUnaryOp.orig_func _numba_func = TypedUserUnaryOp._numba_func + type2 = TypedBuiltinBinaryOp.type2 + __call__ = TypedBuiltinBinaryOp.__call__ class TypedUserMonoid(TypedOpBase): @@ -373,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__ @@ -380,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__ @@ -649,7 +676,12 @@ def __repr__(self): return f"{self._modname}.{self.name}" def __getitem__(self, type_): - if not self._is_udt: + 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: @@ -659,13 +691,8 @@ def __getitem__(self, type_): else: return self._typed_ops[type_] # This is a UDT or is able to operate on UDTs such as `first` any `any` - if type(type_) is tuple: - dtype1, dtype2 = type_ - dtype1 = lookup_dtype(dtype1) - dtype2 = lookup_dtype(dtype2) - else: - dtype1 = dtype2 = lookup_dtype(type_) - return self._getitem_udt(dtype1, dtype2) + dtype = lookup_dtype(type_) + return self._compile_udt(dtype, dtype) def _add(self, op): self._typed_ops[op.type] = op @@ -677,12 +704,6 @@ def __delitem__(self, type_): del self.types[type_] def __contains__(self, type_): - if not self._is_udt: - type_ = lookup_dtype(type_) - if type_ in self._typed_ops or self.is_positional: - return True - elif self._udt_types is None: - return False try: self[type_] except (TypeError, KeyError, numba.NumbaError): @@ -927,7 +948,7 @@ def unary_wrapper(z, x): else: raise UdfParseError("Unable to parse function using Numba") - def _getitem_udt(self, dtype, dtype2): + def _compile_udt(self, dtype, dtype2): if dtype in self._udt_types: return self._udt_ops[dtype] @@ -1379,7 +1400,7 @@ def binary_wrapper(z, x, y): else: raise UdfParseError("Unable to parse function using Numba") - def _getitem_udt(self, dtype, dtype2): + def _compile_udt(self, dtype, dtype2): if dtype2 is None: dtype2 = dtype dtypes = (dtype, dtype2) @@ -1434,7 +1455,6 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover 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 @@ -1485,9 +1505,10 @@ def binary_wrapper(z_ptr, x_ptr, y_ptr): # pragma: no cover op = TypedUserBinaryOp( self, self.name, - dtype, # Should we do a tuple for both inputs? + dtype, ret_type, new_binary[0], + dtype2=dtype2, ) self._udt_types[dtypes] = ret_type self._udt_ops[dtypes] = op @@ -1775,7 +1796,7 @@ def _build(cls, name, binaryop, identity, *, anonymous=False): new_type_obj._add(op) return new_type_obj - def _getitem_udt(self, dtype, dtype2): + def _compile_udt(self, dtype, dtype2): if dtype2 is None: dtype2 = dtype elif dtype != dtype2: @@ -1785,7 +1806,7 @@ def _getitem_udt(self, dtype, dtype2): ) if dtype in self._udt_types: return self._udt_ops[dtype] - binaryop = self.binaryop._getitem_udt(dtype, dtype2) + binaryop = self.binaryop._compile_udt(dtype, dtype2) from .scalar import Scalar ret_type = binaryop.return_type @@ -1998,13 +2019,13 @@ def _build(cls, name, monoid, binaryop, *, anonymous=False): new_type_obj._add(op) return new_type_obj - def _getitem_udt(self, dtype, dtype2): + 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._getitem_udt(dtype, dtype2) + binaryop = self.binaryop._compile_udt(dtype, dtype2) monoid = self.monoid[binaryop.return_type] ret_type = monoid.return_type new_semiring = ffi_new("GrB_Semiring*") @@ -2013,11 +2034,12 @@ def _getitem_udt(self, dtype, dtype2): op = TypedUserSemiring( new_semiring, self.name, - dtype, # Should we do a tuple for both inputs? + dtype, ret_type, new_semiring[0], monoid, binaryop, + dtype2=dtype2, ) self._udt_types[dtypes] = dtype self._udt_ops[dtypes] = op @@ -2265,7 +2287,7 @@ def _is_udt(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._getitem_udt(dtype, dtype2) + return op._compile_udt(dtype, dtype2) if dtype2 is not None: try: dtype = unify( @@ -2276,7 +2298,7 @@ def get_typed_op(op, dtype, dtype2=None, *, is_left_scalar=False, is_right_scala return op[UINT64] if op._udt_types is None: raise - return op._getitem_udt(dtype, dtype2) + 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/tests/pickle3.pkl b/grblas/tests/pickle3.pkl index 1bdedd7742d08a58dbbcc799f7abcba2d4ab7c22..4f60f8134f11672cef9a36a4e8e93e3289ebacd4 100644 GIT binary patch delta 116 zcmZo-yTHcUz%td9c_ZryMrJN=rO6i<(^xXtGGr&4F-0&-0tF^A<@@w-B<58{m!_1& zPU+!{&o4+VN-W7Qn$p9bo?4PwlA1eZ@{}I_^rEDk#9}>!ELT!yUSd(D9#Ah(L56JG Jl%Of4dH~wjDrx`# delta 37 qcmcb>*2Kozz%una(?-@4jLb9tdriK;m^RspDFV!%%9PJkss{iIDGlxb diff --git a/grblas/tests/test_matrix.py b/grblas/tests/test_matrix.py index 02c0aa69f..7d0d6de86 100644 --- a/grblas/tests/test_matrix.py +++ b/grblas/tests/test_matrix.py @@ -3195,7 +3195,15 @@ def test_udt(): 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) - # agg: any_value, first, last, first_index, last_index, count + + # 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") diff --git a/grblas/tests/test_op.py b/grblas/tests/test_op.py index 552023b03..8978b818f 100644 --- a/grblas/tests/test_op.py +++ b/grblas/tests/test_op.py @@ -1133,6 +1133,12 @@ def badfunc(x, y): # pragma: no cover 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]: diff --git a/grblas/tests/test_pickle.py b/grblas/tests/test_pickle.py index 4d1e5baf4..e7d501320 100644 --- a/grblas/tests/test_pickle.py +++ b/grblas/tests/test_pickle.py @@ -309,3 +309,7 @@ def test_udt(): 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_vector.py b/grblas/tests/test_vector.py index b5668da75..30965d3a4 100644 --- a/grblas/tests/test_vector.py +++ b/grblas/tests/test_vector.py @@ -1754,6 +1754,8 @@ def test_udt(): 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") @@ -1796,6 +1798,9 @@ def test_udt(): 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(): diff --git a/grblas/unary/numpy.py b/grblas/unary/numpy.py index a0750b38c..9a910cf1d 100644 --- a/grblas/unary/numpy.py +++ b/grblas/unary/numpy.py @@ -106,7 +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()} +# _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/vector.py b/grblas/vector.py index 1202340af..18309e71e 100644 --- a/grblas/vector.py +++ b/grblas/vector.py @@ -273,16 +273,17 @@ def build(self, indices, values, *, dup_op=None, clear=False, size=None): return dup_op_given = dup_op is not None - if self.dtype._is_udt: # XXX - dup_op = None - else: - if not dup_op_given: + if not dup_op_given: + if not self.dtype._is_udt: dup_op = binary.plus - dup_op = get_typed_op(dup_op, self.dtype, kind="binary") - if dup_op.opclass == "Monoid": - dup_op = dup_op.binaryop else: - self._expect_op(dup_op, "BinaryOp", within="build", argname="dup_op") + 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 + else: + self._expect_op(dup_op, "BinaryOp", within="build", argname="dup_op") indices = _CArray(indices) values = _CArray(values, self.dtype) @@ -362,7 +363,7 @@ def from_values(cls, indices, values, dtype=None, *, size=None, dup_op=None, nam raise ValueError("No indices provided. Unable to infer size.") size = int(indices.max()) + 1 if dtype is None and values.ndim > 1: - # Array-subtdype + # Look for array-subtdype new_dtype = lookup_dtype(np.dtype((new_dtype.np_type, values.shape[1:]))) # Create the new vector w = cls.new(new_dtype, size, name=name)