diff --git a/graphblas/array_api/__init__.py b/graphblas/array_api/__init__.py new file mode 100644 index 000000000..b35c16924 --- /dev/null +++ b/graphblas/array_api/__init__.py @@ -0,0 +1,86 @@ +"""An Array API (2025.12) namespace over ``graphblas.Tensor``: ``import graphblas.array_api as xp``. + +**THIS API IS EXPERIMENTAL AND MAY CHANGE.** + +An Array is logically dense: its value at a position is its stored entry there, or its +``fill_value`` where it stores none (pydata/sparse's model). Results equal numpy's on the same +dense data, and keep the stored entries sparse where the fill values allow. +""" + +from . import ( + _constants, + _creation, + _dtypes, + _elementwise, + _indexing, + _info, + _linalg, + _manipulation, + _searching, + _set, + _sorting, + _statistical, + linalg, +) +from ._array import Array +from ._constants import * # noqa: F403 +from ._creation import * # noqa: F403 +from ._dtypes import * # noqa: F403 +from ._dtypes import DTYPES as _DTYPES +from ._elementwise import * # noqa: F403 +from ._indexing import * # noqa: F403 +from ._info import * # noqa: F403 +from ._linalg import * # noqa: F403 +from ._manipulation import * # noqa: F403 +from ._searching import * # noqa: F403 +from ._set import * # noqa: F403 +from ._sorting import * # noqa: F403 +from ._statistical import * # noqa: F403 +from ._utils import DensifyError + +__array_api_version__ = "2025.12" + +bool = _DTYPES["bool"] # noqa: A001 (the Array API's name) +int8 = _DTYPES["int8"] +int16 = _DTYPES["int16"] +int32 = _DTYPES["int32"] +int64 = _DTYPES["int64"] +uint8 = _DTYPES["uint8"] +uint16 = _DTYPES["uint16"] +uint32 = _DTYPES["uint32"] +uint64 = _DTYPES["uint64"] +float32 = _DTYPES["float32"] +float64 = _DTYPES["float64"] +complex64 = _DTYPES["complex64"] +complex128 = _DTYPES["complex128"] + +__all__ = [ + "Array", + "DensifyError", + "__array_api_version__", + "bool", + "complex64", + "complex128", + "float32", + "float64", + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", +] +__all__ += _constants.__all__ +__all__ += _creation.__all__ +__all__ += _dtypes.__all__ +__all__ += _elementwise.__all__ +__all__ += _indexing.__all__ +__all__ += _info.__all__ +__all__ += _linalg.__all__ +__all__ += _manipulation.__all__ +__all__ += _searching.__all__ +__all__ += _set.__all__ +__all__ += _sorting.__all__ +__all__ += _statistical.__all__ diff --git a/graphblas/array_api/_array.py b/graphblas/array_api/_array.py new file mode 100644 index 000000000..8f165ac69 --- /dev/null +++ b/graphblas/array_api/_array.py @@ -0,0 +1,405 @@ +"""The Array: a Tensor of stored entries plus a fill value for every other position.""" + +import math + +import numpy as np + +from . import _utils +from ._utils import np_dtype + +CPU_DEVICE = "cpu" +API_VERSIONS = ("2021.12", "2022.12", "2023.12", "2024.12", "2025.12") + + +def check_device(device): + if device is not None and device != CPU_DEVICE: + raise ValueError(f"graphblas.array_api has only the device {CPU_DEVICE!r}, not {device!r}") + + +class Array: + """An n-dimensional array whose value at a position is its stored entry there, or + ``fill_value`` where it stores none (pydata/sparse's model). + + Create Arrays with the namespace's functions (``asarray``, ``zeros``, ``from_coo``, ...). + The backing ``graphblas.Tensor`` holds the stored entries; stored entries may equal the + fill value. + """ + + # The shape and dtype are kept here, not read from the Tensor, so that an Array whose + # Tensor is not computed yet (``_pending``, a deferred contraction) can answer them + __slots__ = "_tensor", "_pending", "_fill", "_facts", "_shape", "_dtype", "__weakref__" + # numpy hands its operators over to ours (``ndarray + Array``), as it does for other + # array types, rather than converting the Array to a dense ndarray + __array_ufunc__ = None + + def __init__(self, *args, **kwargs): + raise TypeError( + "Array is not created directly; use graphblas.array_api.asarray (or zeros, full, " + "from_coo, ...)" + ) + + @classmethod + def _new(cls, tensor, fill): + """An Array of ``tensor``'s entries and ``fill`` (cast to its dtype) elsewhere. + + A 0-d Array stores nothing: its value is its fill. + """ + if tensor.ndim == 0 and tensor.nvals: + fill = tensor.to_scalar().value + tensor = _utils.empty_0d(tensor.dtype) + self = object.__new__(cls) + self._tensor = tensor + self._pending = None + self._shape = tensor.shape + self._dtype = np_dtype(tensor.dtype) + self._fill = _utils.scalar(fill, self._dtype) + self._facts = None + return self + + @classmethod + def _deferred(cls, pending, shape, dtype, fill): + """An Array whose Tensor ``pending.compute()`` gives on first use (see ``_lazy``).""" + self = object.__new__(cls) + self._tensor = None + self._pending = pending + self._shape = tuple(shape) + self._dtype = np.dtype(dtype) + self._fill = _utils.scalar(fill, self._dtype) + self._facts = None + return self + + def _T(self): + """The Tensor of stored entries. Every read of it goes through here: a pending Array + computes it here, once. + """ + if (T := self._tensor) is None: + T = self._tensor = self._pending.compute() + self._pending = None # (releases the operands' snapshots) + return T + + def _replace(self, other): + """Make this Array hold ``other``'s values (for in-place operators and assignment).""" + self._tensor = other._T() + self._pending = None # (a pending product it was is not its value any more) + self._shape = other._shape + self._dtype = other._dtype + self._fill = other._fill + self._facts = None + + def _fact(self, name, compute): + """A cached property of the stored values. Code that changes the stored entries in + place must clear the cache (``_facts = None``), as ``_replace`` does. + """ + facts = self._facts + if facts is None: + facts = self._facts = {} + if name not in facts: + facts[name] = compute(self) + return facts[name] + + def _copy(self): + return Array._new(self._T().dup(), self._fill) + + def __reduce__(self): + # (a Tensor holds weak references, which do not pickle: rebuild from coordinates) + coords, values = self._T().to_coo() + return _unpickle, (coords, values, self._shape, self._dtype.str, self._fill) + + # (copy.copy would share the Tensor, which __setitem__ changes in place) + def __copy__(self): + return self._copy() + + def __deepcopy__(self, memo): + return self._copy() + + # Attributes + @property + def dtype(self): + return self._dtype + + @property + def device(self): + return CPU_DEVICE + + @property + def shape(self): + return self._shape + + @property + def ndim(self): + return len(self._shape) + + @property + def size(self): + return math.prod(self._shape) + + @property + def T(self): + """The axes reversed (numpy's ``.T``; the Array API defines it for 2-d arrays).""" + from ._manipulation import permute_dims + + return permute_dims(self, tuple(range(self.ndim - 1, -1, -1))) + + @property + def mT(self): + from ._manipulation import permute_dims + + if self.ndim < 2: + raise ValueError("mT needs at least 2 dimensions") + axes = list(range(self.ndim)) + axes[-2], axes[-1] = axes[-1], axes[-2] + return permute_dims(self, tuple(axes)) + + # Extensions + @property + def fill_value(self): + """The value at every position the Array does not store (a numpy scalar).""" + return self._fill + + @property + def nnz(self): + """The number of stored entries, entries equal to the fill value included (as in + pydata/sparse); ``prune`` drops those. No other function's value depends on it. + """ + return self._T().nvals + + @property + def tensor(self): + """The stored entries as a ``graphblas.Tensor`` (positions holding the fill value are + missing); a copy that shares the data until either side changes. + """ + return self._T().dup() + + def todense(self): + """A numpy array of every position (raises past ``array_api_densify_limit``).""" + _utils.check_dense(self.size, f"densifying an array of shape {self.shape}") + dt = self.dtype + T = self._T() + if T.ndim == 0 or T.nvals == 0: + return np.full(self.shape, self._fill, dtype=dt) + if T.nvals == self.size: + return T.to_dense(dtype=T.dtype).astype(dt, copy=False) + coords, values = T.to_coo() + out = np.full(self.shape, self._fill, dtype=dt) + out[coords] = values + return out + + def to_coo(self): + """``(coords, values)`` of the stored entries: a tuple of index arrays (one per axis) + and the values, in C order. + """ + return self._T().to_coo() + + def __repr__(self): + fill = self._fill.item() + if self.ndim == 0: + return f"Array({fill!r}, dtype={self.dtype})" + if self._tensor is None: # (its nnz would compute it) + return f"Array(shape={self.shape}, dtype={self.dtype}, fill_value={fill!r}, pending)" + return f"Array(shape={self.shape}, dtype={self.dtype}, nnz={self.nnz}, fill_value={fill!r})" + + # Protocols + def __array_namespace__(self, *, api_version=None): + if api_version is not None and api_version not in API_VERSIONS: + raise ValueError(f"unsupported Array API version {api_version!r}") + from .. import array_api + + return array_api + + def __array__(self, dtype=None, copy=None): + if copy is False: + raise ValueError("converting an Array to a numpy array always copies") + out = self.todense() + return out if dtype is None else out.astype(dtype, copy=False) + + def __dlpack__(self, *, stream=None, max_version=None, dl_device=None, copy=None): + if copy is False: + raise BufferError("exporting an Array through DLPack always copies") + kwargs = {} if max_version is None else {"max_version": max_version} + if dl_device is not None: + kwargs["dl_device"] = dl_device + return self.todense().__dlpack__(stream=stream, **kwargs) + + def __dlpack_device__(self): + return np.empty(0).__dlpack_device__() + + def to_device(self, device, /, *, stream=None): + check_device(device) + return self + + def _value(self, what): + if self.ndim != 0: + raise TypeError(f"{what} needs a 0-d array; this one has shape {self.shape}") + return self._fill + + def __bool__(self): + return bool(self._value("bool()")) + + def __int__(self): + value = self._value("int()") + if self.dtype.kind == "c": + raise TypeError("int() of a complex array") + return int(value) + + def __float__(self): + value = self._value("float()") + if self.dtype.kind == "c": + raise TypeError("float() of a complex array") + return float(value) + + def __complex__(self): + return complex(self._value("complex()")) + + def __index__(self): + value = self._value("an index") + if self.dtype.kind not in "iu": + raise TypeError(f"only integer arrays are indices, not {self.dtype}") + return int(value) + + def __len__(self): + if self.ndim == 0: + raise TypeError("len() of a 0-d array") + return self.shape[0] + + def __iter__(self): + if self.ndim == 0: + raise TypeError("iteration over a 0-d array") + # (computed first: each row of a pending product would be a product of its own) + self._T() + return (self[i] for i in range(self.shape[0])) + + __hash__ = None + + # Indexing (graphblas.array_api._indexing) + def __getitem__(self, key): + from ._indexing import getitem + + return getitem(self, key) + + def __setitem__(self, key, value): + from ._indexing import setitem + + setitem(self, key, value) + + # Operators + def __matmul__(self, other): + if not _is_operand(other): + return NotImplemented + from ._linalg import matmul + + return matmul(self, other) + + def __rmatmul__(self, other): + if not _is_operand(other): + return NotImplemented + from ._linalg import matmul + + return matmul(other, self) + + def __imatmul__(self, other): + from ._linalg import matmul + + return self._inplace(matmul(self, other)) + + def _inplace(self, result): + if result.shape != self.shape: + raise ValueError( + f"an in-place operation would change the shape {self.shape} to {result.shape}" + ) + if result.dtype != self.dtype: + raise TypeError( + f"an in-place operation would change the data type {self.dtype} to " + f"{result.dtype}" + ) + self._replace(result) + return self + + def __neg__(self): + from ._elementwise import negative + + return negative(self) + + def __pos__(self): + from ._elementwise import positive + + return positive(self) + + def __abs__(self): + from ._elementwise import abs as abs_ + + return abs_(self) + + def __invert__(self): + from ._elementwise import bitwise_invert + + return bitwise_invert(self) + + +def _unpickle(coords, values, shape, dtype, fill): + if not shape: + return Array._new(_utils.empty_0d(dtype), fill) + from ._creation import from_coo + + return from_coo(coords, values, shape, fill_value=fill, dtype=np.dtype(dtype)) + + +def _is_operand(x): + return isinstance(x, (Array, bool, int, float, complex, np.generic, np.ndarray)) + + +def _binary_methods(name, func_name): + def forward(self, other): + if not _is_operand(other): + return NotImplemented + from . import _elementwise + + return getattr(_elementwise, func_name)(self, other) + + def reflected(self, other): + if not _is_operand(other): + return NotImplemented + from . import _elementwise + + return getattr(_elementwise, func_name)(other, self) + + def inplace(self, other): + from . import _elementwise + + return self._inplace(getattr(_elementwise, func_name)(self, other)) + + forward.__name__ = f"__{name}__" + reflected.__name__ = f"__r{name}__" + inplace.__name__ = f"__i{name}__" + return forward, reflected, inplace + + +for _name, _func in [ + ("add", "add"), + ("sub", "subtract"), + ("mul", "multiply"), + ("truediv", "divide"), + ("floordiv", "floor_divide"), + ("mod", "remainder"), + ("pow", "pow"), + ("and", "bitwise_and"), + ("or", "bitwise_or"), + ("xor", "bitwise_xor"), + ("lshift", "bitwise_left_shift"), + ("rshift", "bitwise_right_shift"), +]: + _f, _r, _i = _binary_methods(_name, _func) + setattr(Array, f"__{_name}__", _f) + setattr(Array, f"__r{_name}__", _r) + setattr(Array, f"__i{_name}__", _i) + +for _name, _func in [ + ("eq", "equal"), + ("ne", "not_equal"), + ("lt", "less"), + ("le", "less_equal"), + ("gt", "greater"), + ("ge", "greater_equal"), +]: + _f, _r, _i = _binary_methods(_name, _func) + setattr(Array, f"__{_name}__", _f) +del _name, _func, _f, _r, _i diff --git a/graphblas/array_api/_constants.py b/graphblas/array_api/_constants.py new file mode 100644 index 000000000..18a5473f8 --- /dev/null +++ b/graphblas/array_api/_constants.py @@ -0,0 +1,11 @@ +"""Constants.""" + +import math + +__all__ = ["e", "inf", "nan", "newaxis", "pi"] + +e = math.e +inf = math.inf +nan = math.nan +newaxis = None +pi = math.pi diff --git a/graphblas/array_api/_creation.py b/graphblas/array_api/_creation.py new file mode 100644 index 000000000..a4a32a9e3 --- /dev/null +++ b/graphblas/array_api/_creation.py @@ -0,0 +1,428 @@ +"""Creation functions. + +``zeros``, ``ones``, ``full``, ``empty`` and their ``_like`` forms store nothing: their value is +the fill, so they cost O(1). Conversions from dense data store every position whose value is +not identical to the fill (0 of the dtype unless ``fill_value=`` is given): ``-0.0`` is stored +when the fill is ``0.0``, and NaN is stored unless the fill is NaN. +""" + +import builtins +import math + +import numpy as np + +from .. import binary +from ..core.tensor import Tensor +from . import _utils +from ._array import Array, check_device +from ._dtypes import as_dtype, check_supported +from ._utils import gb_dtype + +__all__ = [ + "arange", + "asarray", + "empty", + "empty_like", + "eye", + "from_coo", + "from_dlpack", + "full", + "full_like", + "linspace", + "meshgrid", + "ones", + "ones_like", + "tril", + "triu", + "zeros", + "zeros_like", +] + +_DEFAULT_FLOAT = np.dtype("float64") + + +def _shape(shape): + if isinstance(shape, (int, np.integer)): + shape = (shape,) + shape = tuple(int(n) for n in shape) + if any(n < 0 for n in shape): + raise ValueError(f"negative dimensions are not allowed: {shape}") + return shape + + +def _default_fill(dtype): + return _utils.scalar(0, dtype) + + +def _from_numpy(a, fill_value=None): + """An Array of the numpy array ``a`` storing what differs from the fill.""" + dt = check_supported(a.dtype, "an array") + fill = _default_fill(dt) if fill_value is None else _utils.scalar(fill_value, dt) + if a.ndim == 0: + return Array._new(_utils.empty_0d(dt), a[()]) + if a.size == 0: + return Array._new(Tensor(gb_dtype(dt), a.shape), fill) + keep = _utils.differs(a, fill) + if keep.all(): + T = Tensor.from_dense(a, dtype=gb_dtype(dt)) + else: + coords = np.nonzero(keep) + T = Tensor.from_coo(coords, a[coords], gb_dtype(dt), shape=a.shape) + return Array._new(T, fill) + + +def _from_graphblas(obj, dtype, fill_value): + """A Tensor, Matrix, Vector or Scalar as an Array whose missing entries hold the fill.""" + from .. import Matrix, Scalar, Vector + from ..core.matrix import TransposedMatrix + + if isinstance(obj, Scalar): + dt = _utils.np_dtype(obj.dtype) + value = obj.value if obj.nvals else fill_value if fill_value is not None else 0 + out = Array._new(_utils.empty_0d(dt), value) + else: + if isinstance(obj, Tensor): + T = obj.dup() + elif isinstance(obj, Vector): + T = Tensor.from_vector(obj) + elif isinstance(obj, (Matrix, TransposedMatrix)): + T = Tensor.from_matrix(obj) + else: # pragma: no cover (safety) + raise TypeError(type(obj)) + dt = check_supported(_utils.np_dtype(T.dtype), "an array") + fill = _default_fill(dt) if fill_value is None else fill_value + out = Array._new(T, fill) + if dtype is not None and as_dtype(dtype) != out.dtype: + from ._elementwise import cast + + out = cast(out, as_dtype(dtype)) + return out + + +def _from_sparse(obj, dtype, fill_value): + """A scipy.sparse array or matrix, or a pydata/sparse array.""" + if hasattr(obj, "tocoo") and hasattr(obj, "format") and not hasattr(obj, "fill_value"): + coo = obj.tocoo() + coords = tuple(np.asarray(c) for c in coo.coords) + data, shape, fill = np.asarray(coo.data), coo.shape, None + if fill_value is not None: + # scipy.sparse's missing entries are zeros + raise ValueError("a scipy.sparse array has fill value 0; convert it first") + if not getattr(coo, "has_canonical_format", True): + coo.sum_duplicates() + coords, data = tuple(np.asarray(c) for c in coo.coords), np.asarray(coo.data) + else: + if obj.format != "coo": + obj = obj.asformat("coo") + coords = tuple(np.asarray(c) for c in obj.coords) + data, shape, fill = np.asarray(obj.data), obj.shape, obj.fill_value + if fill_value is not None and not _utils.same_value(fill_value, fill): + raise ValueError(f"fill_value={fill_value!r} differs from the sparse array's {fill!r}") + return from_coo(coords, data, shape, fill_value=fill, dtype=dtype) + + +def _is_sparse(obj): + module = type(obj).__module__ + return module.startswith(("scipy.sparse", "sparse")) + + +def asarray(obj, /, *, dtype=None, device=None, copy=None, fill_value=None): + """Convert ``obj`` to an Array. + + ``obj`` is an Array, a Python scalar or (nested) sequence, an object numpy converts (a + buffer, an ndarray), or (extensions) a ``graphblas`` Tensor, Matrix, Vector or Scalar, a + scipy.sparse array or a pydata/sparse array, whose missing entries hold the fill. + ``fill_value`` (an extension) is the value of unstored positions: 0 of the dtype by + default, or the sparse array's own fill value. + """ + check_device(device) + if isinstance(obj, Array): + if fill_value is not None and not _utils.same_value(fill_value, obj._fill): + raise ValueError( + "asarray of an Array keeps its fill value; xp.refill(x, fill_value) gives one " + "with another (the same values)" + ) + if dtype is not None and as_dtype(dtype) != obj.dtype: + if copy is False: + raise ValueError("changing the data type needs a copy, and copy=False") + from ._elementwise import cast + + return cast(obj, as_dtype(dtype)) + return obj._copy() if copy else obj + if copy is False: + raise ValueError("asarray stores its own copy of anything but an Array; copy=False") + from .. import Matrix, Scalar, Vector + from ..core.matrix import TransposedMatrix + + if isinstance(obj, (Tensor, Matrix, Vector, Scalar, TransposedMatrix)): + return _from_graphblas(obj, dtype, fill_value) + if _is_sparse(obj) and (hasattr(obj, "tocoo") or hasattr(obj, "coords")): + return _from_sparse(obj, dtype, fill_value) + dt = None if dtype is None else as_dtype(dtype) + if isinstance(obj, (list, tuple)): + stacked = _stack_arrays(obj) + if stacked is not None: + return asarray(stacked, dtype=dtype, fill_value=fill_value) + obj = _unwrap_nested(obj) + if dt is None and isinstance(obj, builtins.bool): + dt = np.dtype(bool) + elif dt is None and isinstance(obj, int) and not isinstance(obj, np.generic): + dt = np.dtype("int64") + a = np.asarray(obj, dtype=dt) + if not a.dtype.isnative: + a = a.astype(a.dtype.newbyteorder("=")) + if dt is None and a.dtype.kind == "f" and a.dtype.itemsize < 4: + a = a.astype(np.float32) + return _from_numpy(a, fill_value) + + +def _stack_arrays(obj): + """A sequence of n-d Arrays of one shape (or such sequences), stacked sparsely; None for + anything else. + """ + from ._manipulation import stack + + if not obj: + return None + parts = [] + for o in obj: + if isinstance(o, (list, tuple)): + o = _stack_arrays(o) + if not isinstance(o, Array) or o.ndim == 0: + return None + parts.append(o) + if len({p.shape for p in parts}) != 1: + return None + return stack(parts) + + +def _unwrap_nested(obj): + """A (nested) sequence with its 0-d Arrays replaced by their values (numpy would densify + every Array through ``__array__``, which works too, but slower). + """ + if isinstance(obj, Array): + return obj._fill if obj.ndim == 0 else obj.todense() + if isinstance(obj, (list, tuple)): + return [_unwrap_nested(o) for o in obj] + return obj + + +def from_coo(coords, data, shape=None, *, fill_value=None, dtype=None): + """An Array storing ``data`` at ``coords`` (one index array per axis, or an (ndim, nnz) + array), and ``fill_value`` (0 by default) elsewhere. ``data`` may be a scalar. Values at a + repeated coordinate are summed (logical or for bool), as pydata/sparse and scipy.sparse + do. + """ + data = np.asarray(data) + dt = as_dtype(dtype) if dtype is not None else check_supported(data.dtype, "from_coo data") + if dtype is not None: + data = data.astype(dt, copy=False) + if isinstance(coords, np.ndarray) and coords.ndim == 2: + coords = tuple(coords) + coords = tuple(np.asarray(c) for c in coords) + if shape is None: + shape = tuple(int(c.max()) + 1 if c.size else 0 for c in coords) + shape = _shape(shape) + fill = _default_fill(dt) if fill_value is None else _utils.scalar(fill_value, dt) + if len(shape) == 0: + value = data.reshape(-1)[0] if data.size else fill + return Array._new(_utils.empty_0d(dt), value) + if coords and coords[0].size: + gdt = gb_dtype(dt) + plus = binary.plus[gdt] # (lor for bool) + if data.ndim: + T = Tensor.from_coo(coords, data, gdt, shape=shape, dup_op=plus) + else: + # One value everywhere (iso, which keeps no duplicates); when coordinates + # repeat, each sums its copies: built again from the value repeated + T = Tensor.from_coo(coords, data, gdt, shape=shape) + if T.nvals < coords[0].size and dt.kind != "b": + data = np.broadcast_to(data, coords[0].shape) + T = Tensor.from_coo(coords, data, gdt, shape=shape, dup_op=plus) + else: + T = Tensor(gb_dtype(dt), shape) + return Array._new(T, fill) + + +def _full(shape, value, dtype): + shape = _shape(shape) + return Array._new(Tensor(gb_dtype(dtype), shape), value) + + +def _fill_dtype(fill_value): + if isinstance(fill_value, Array): + return fill_value.dtype + if isinstance(fill_value, np.generic): + return check_supported(fill_value.dtype, "full") + if isinstance(fill_value, builtins.bool): + return np.dtype(bool) + if isinstance(fill_value, int): + return np.dtype("int64") + if isinstance(fill_value, float): + return _DEFAULT_FLOAT + if isinstance(fill_value, complex): + return np.dtype("complex128") + raise TypeError(f"fill_value must be a Python scalar, not {type(fill_value).__name__}") + + +def full(shape, fill_value, *, dtype=None, device=None): + check_device(device) + dt = _fill_dtype(fill_value) if dtype is None else as_dtype(dtype) + if isinstance(fill_value, Array): + fill_value = fill_value._value("full's fill_value") + if isinstance(fill_value, int) and dt.kind in "iu": + info = np.iinfo(dt) + if not info.min <= fill_value <= info.max: + raise OverflowError(f"fill_value {fill_value} is out of range for {dt}") + return _full(shape, fill_value, dt) + + +def zeros(shape, *, dtype=None, device=None): + check_device(device) + dt = _DEFAULT_FLOAT if dtype is None else as_dtype(dtype) + return _full(shape, 0, dt) + + +def ones(shape, *, dtype=None, device=None): + check_device(device) + dt = _DEFAULT_FLOAT if dtype is None else as_dtype(dtype) + return _full(shape, 1, dt) + + +def empty(shape, *, dtype=None, device=None): + """Zeros: an empty array stores nothing, and its fill is 0 of the dtype.""" + return zeros(shape, dtype=dtype, device=device) + + +def _like(x, dtype, device): + if not isinstance(x, Array): + raise TypeError(f"needs an Array, not {type(x).__name__}") + check_device(device) + return x.shape, x.dtype if dtype is None else as_dtype(dtype) + + +def full_like(x, /, fill_value, *, dtype=None, device=None): + shape, dt = _like(x, dtype, device) + if isinstance(fill_value, Array): + fill_value = fill_value._value("full_like's fill_value") + return _full(shape, fill_value, dt) + + +def zeros_like(x, /, *, dtype=None, device=None): + shape, dt = _like(x, dtype, device) + return _full(shape, 0, dt) + + +def ones_like(x, /, *, dtype=None, device=None): + shape, dt = _like(x, dtype, device) + return _full(shape, 1, dt) + + +def empty_like(x, /, *, dtype=None, device=None): + return zeros_like(x, dtype=dtype, device=device) + + +def eye(n_rows, n_cols=None, /, *, k=0, dtype=None, device=None): + check_device(device) + dt = _DEFAULT_FLOAT if dtype is None else as_dtype(dtype) + n_cols = n_rows if n_cols is None else n_cols + shape = _shape((n_rows, n_cols)) + i = np.arange(builtins.max(0, -k), builtins.max(0, builtins.min(n_rows, n_cols - k))) + if i.size == 0: + return _full(shape, 0, dt) + T = Tensor.from_coo((i, i + k), _utils.scalar(1, dt), gb_dtype(dt), shape=shape) + return Array._new(T, 0) + + +def arange(start, /, stop=None, step=1, *, dtype=None, device=None): + check_device(device) + if stop is None: + start, stop = 0, start + dt = None if dtype is None else as_dtype(dtype) + if step == 0: + raise ValueError("arange's step must not be 0") + n = builtins.max(0, math.ceil((stop - start) / step)) + _utils.check_dense(n, "arange") + if dt is None: + if builtins.all(isinstance(v, int) for v in (start, stop, step)): + dt = np.dtype("int64") + else: + dt = _DEFAULT_FLOAT + a = np.arange(start, stop, step, dtype=dt) + return _from_numpy(a[:n] if a.size > n else a) + + +def linspace(start, stop, /, num, *, dtype=None, device=None, endpoint=True): + check_device(device) + _utils.check_dense(num, "linspace") + if dtype is None: + is_complex = builtins.any(isinstance(v, complex) for v in (start, stop)) + dt = np.dtype("complex128") if is_complex else _DEFAULT_FLOAT + else: + dt = as_dtype(dtype) + a = np.linspace(start, stop, num, endpoint=endpoint, dtype=dt) + return _from_numpy(a) + + +def from_dlpack(x, /, *, device=None, copy=None): + check_device(device) + if isinstance(x, Array): + return x._copy() if copy else x + if copy is False: + raise ValueError("from_dlpack stores its own copy of anything but an Array; copy=False") + return _from_numpy(np.from_dlpack(x)) + + +def meshgrid(*arrays, indexing="xy"): + from ._manipulation import _broadcast_to, reshape + + if indexing not in {"xy", "ij"}: + raise ValueError(f"indexing must be 'xy' or 'ij', not {indexing!r}") + for x in arrays: + if not isinstance(x, Array) or x.ndim != 1: + raise ValueError("meshgrid needs 1-d arrays") + if len({x.dtype for x in arrays}) > 1: + raise ValueError("meshgrid needs arrays of one data type") + n = len(arrays) + shape = [x.shape[0] for x in arrays] + if indexing == "xy" and n >= 2: + shape[0], shape[1] = shape[1], shape[0] + out = [] + for i, x in enumerate(arrays): + axis = i + if indexing == "xy" and n >= 2 and i < 2: + axis = 1 - i + s = [1] * n + s[axis] = x.shape[0] + out.append(_broadcast_to(reshape(x, tuple(s)), tuple(shape))) + return tuple(out) + + +def _tri(x, k, lower): + if not isinstance(x, Array): + raise TypeError(f"needs an Array, not {type(x).__name__}") + if x.ndim < 2: + raise ValueError("tril and triu need at least 2 dimensions") + zero = _default_fill(x.dtype) + if not _utils.same_value(x._fill, zero): + # The kept triangle holds the fill where nothing is stored: store it there first + _utils.check_dense(x.size, "tril or triu of an array whose fill is not 0") + x = _from_numpy(x.todense()) + T = x._T() + if T.nvals == 0: + return _full(x.shape, 0, x.dtype) + coords, values = T.to_coo() + i, j = coords[-2], coords[-1] + keep = (j - i <= k) if lower else (j - i >= k) + if not keep.any(): + return _full(x.shape, 0, x.dtype) + coords = tuple(c[keep] for c in coords) + return Array._new(Tensor.from_coo(coords, values[keep], T.dtype, shape=T.shape), zero) + + +def tril(x, /, *, k=0): + return _tri(x, k, lower=True) + + +def triu(x, /, *, k=0): + return _tri(x, k, lower=False) diff --git a/graphblas/array_api/_dtypes.py b/graphblas/array_api/_dtypes.py new file mode 100644 index 000000000..176121836 --- /dev/null +++ b/graphblas/array_api/_dtypes.py @@ -0,0 +1,171 @@ +"""Data types, type promotion and the data type functions. + +The namespace's dtypes are numpy dtype instances (``xp.float64 is np.dtype("float64")``), and +an Array's ``dtype`` is one of them, so they compare and hash alike. Type promotion is numpy's, +with Python scalars weak (NEP 50). +""" + +import builtins +from typing import NamedTuple + +import numpy as np + +__all__ = ["astype", "can_cast", "finfo", "iinfo", "isdtype", "result_type"] + +_NAMES = ( + "bool", + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", + "float32", + "float64", + "complex64", + "complex128", +) +DTYPES = {name: np.dtype(name) for name in _NAMES} +_SUPPORTED = frozenset(DTYPES.values()) + +_KINDS = { + "bool": frozenset({DTYPES["bool"]}), + "signed integer": frozenset(DTYPES[n] for n in ("int8", "int16", "int32", "int64")), + "unsigned integer": frozenset(DTYPES[n] for n in ("uint8", "uint16", "uint32", "uint64")), + "real floating": frozenset({DTYPES["float32"], DTYPES["float64"]}), + "complex floating": frozenset({DTYPES["complex64"], DTYPES["complex128"]}), +} +_KINDS["integral"] = _KINDS["signed integer"] | _KINDS["unsigned integer"] +_KINDS["numeric"] = _KINDS["integral"] | _KINDS["real floating"] | _KINDS["complex floating"] + + +def as_dtype(dtype): + """``dtype`` (a dtype of this namespace, or anything numpy takes for one) as a supported + numpy dtype; TypeError otherwise. + """ + if dtype is None: + raise TypeError("dtype must not be None here") + try: + dt = np.dtype(dtype) + except TypeError: + raise TypeError(f"{dtype!r} is not a data type") from None + if dt not in _SUPPORTED: + raise TypeError(f"data type {dt} is not supported by graphblas.array_api") + return dt + + +def check_supported(dt, what): + if dt not in _SUPPORTED: + raise TypeError(f"{what} would have data type {dt}, which graphblas.array_api lacks") + return dt + + +def _dtype_of(x): + from ._array import Array + + if isinstance(x, Array): + return x.dtype + return as_dtype(x) + + +def isdtype(dtype, kind): + if isinstance(dtype, type) and issubclass(dtype, np.generic): + dtype = np.dtype(dtype) + if not isinstance(dtype, np.dtype) or dtype not in _SUPPORTED: + raise TypeError(f"isdtype needs a data type of this namespace, not {dtype!r}") + if isinstance(kind, tuple): + return builtins.any(isdtype(dtype, k) for k in kind) + if isinstance(kind, str): + try: + return dtype in _KINDS[kind] + except KeyError: + raise ValueError(f"unknown data type kind {kind!r}") from None + return dtype == as_dtype(kind) + + +def result_type(*arrays_and_dtypes): + from ._array import Array + + if not arrays_and_dtypes: + raise ValueError("result_type needs at least one array or data type") + args = [] + has_typed = False + for x in arrays_and_dtypes: + if isinstance(x, Array): + args.append(x.dtype) + has_typed = True + elif isinstance(x, (builtins.bool, int, float, complex)) and not isinstance(x, np.generic): + args.append(x) # weak, as numpy takes a Python scalar + else: + args.append(as_dtype(x)) + has_typed = True + if not has_typed: + raise ValueError("result_type needs at least one array or data type, not only scalars") + # numpy promotes many arguments pairwise in an order of its own: fold left to right, as the + # Array API's promotion table does, so that the result does not depend on the order + rt = None + for arg in args: + if isinstance(arg, np.dtype): + rt = arg if rt is None else np.result_type(rt, arg) + for arg in args: + if not isinstance(arg, np.dtype): + rt = np.result_type(rt, arg) + return check_supported(rt, "result_type") + + +def can_cast(from_, to, /): + return builtins.bool(np.can_cast(_dtype_of(from_), as_dtype(to), casting="safe")) + + +class finfo_object(NamedTuple): + bits: int + eps: float + max: float + min: float + smallest_normal: float + dtype: np.dtype + + +class iinfo_object(NamedTuple): + bits: int + max: int + min: int + dtype: np.dtype + + +def finfo(type, /): # noqa: A002 (the Array API's name) + dt = _dtype_of(type) + if dt.kind not in "fc": + raise TypeError(f"finfo needs a floating-point data type, not {dt}") + info = np.finfo(dt) + return finfo_object( + builtins.int(info.bits), + float(info.eps), + float(info.max), + float(info.min), + float(info.smallest_normal), + np.dtype(info.dtype), + ) + + +def iinfo(type, /): # noqa: A002 (the Array API's name) + dt = _dtype_of(type) + if dt.kind not in "iu": + raise TypeError(f"iinfo needs an integer data type, not {dt}") + info = np.iinfo(dt) + return iinfo_object(builtins.int(info.bits), builtins.int(info.max), builtins.int(info.min), dt) + + +def astype(x, dtype, /, *, copy=True, device=None): + from ._array import Array, check_device + from ._elementwise import cast + + check_device(device) + if not isinstance(x, Array): + raise TypeError(f"astype needs an Array, not {type(x)}") + dt = as_dtype(dtype) + if dt == x.dtype: + return x._copy() if copy else x + return cast(x, dt) diff --git a/graphblas/array_api/_elementwise.py b/graphblas/array_api/_elementwise.py new file mode 100644 index 000000000..b52fd44ef --- /dev/null +++ b/graphblas/array_api/_elementwise.py @@ -0,0 +1,1036 @@ +"""Element-wise functions: numpy's answer on every position (DESIGN D3b: NaN, inf and signed +zeros exact, transcendental functions within libm's rounding, the sign of a zero at an unstored +position unspecified). + +A function of one array maps the stored entries and the fill alike. A function of two arrays +takes the union of their entries, an unstored position of one operand taking its fill there, +and the new fill is the function of the two fills. Python scalars are weak (NEP 50). + +Where a GraphBLAS operator gives numpy's answer exactly for the dtypes at hand, the entries +are computed in GraphBLAS; otherwise numpy maps the stored values (the pattern stays as it is, +see ``_utils.map_values``). numpy also computes every fill, so an entry equal to the fill maps +to the new fill. +""" + +import builtins +import math + +import numpy as np + +from .. import binary, monoid, select, unary +from ..core.scalar import Scalar +from ..core.tensor import Tensor +from ..dtypes import BOOL, INT64 +from . import _lazy, _utils +from ._array import Array +from ._dtypes import check_supported +from ._utils import gb_dtype, gb_scalar + +__all__ = [ + "abs", + "acos", + "acosh", + "add", + "asin", + "asinh", + "atan", + "atan2", + "atanh", + "bitwise_and", + "bitwise_invert", + "bitwise_left_shift", + "bitwise_or", + "bitwise_right_shift", + "bitwise_xor", + "ceil", + "clip", + "conj", + "copysign", + "cos", + "cosh", + "divide", + "equal", + "exp", + "expm1", + "floor", + "floor_divide", + "greater", + "greater_equal", + "hypot", + "imag", + "isfinite", + "isinf", + "isnan", + "less", + "less_equal", + "log", + "log1p", + "log2", + "log10", + "logaddexp", + "logical_and", + "logical_not", + "logical_or", + "logical_xor", + "maximum", + "minimum", + "multiply", + "negative", + "nextafter", + "not_equal", + "positive", + "pow", + "real", + "reciprocal", + "remainder", + "round", + "sign", + "signbit", + "sin", + "sinh", + "sqrt", + "square", + "subtract", + "tan", + "tanh", + "trunc", +] + +_FLOAT = "f" +_INEXACT = "fc" +_INTEGRAL = "iu" +_REAL = "biuf" + + +# Operands +def _as_operand(x, name): + """An Array; a Python scalar stays as it is (weak); a numpy scalar or array becomes an + Array (strong, as numpy treats it). + """ + if isinstance(x, Array): + return x + if isinstance(x, (np.generic, np.ndarray)): + from ._creation import asarray + + return asarray(x) + if isinstance(x, (builtins.bool, int, float, complex)): + return x + raise TypeError(f"{name} needs arrays or Python scalars, not {type(x).__name__}") + + +def _probe(x): + """What numpy's type resolution sees: an empty array of an Array's dtype, or the weak + Python scalar itself. + """ + if isinstance(x, Array): + return np.empty(0, x.dtype) + return x + + +def _fill_operand(x): + """The value of ``x`` at an unstored position, typed as numpy would type ``x``.""" + if isinstance(x, Array): + return np.asarray(x._fill) + return x + + +def _fill_of(ufunc, rdt, operands): + """``ufunc`` of the operands' fills: the result's fill. + + An Array that stores every position never shows its fill, so where the fills raise + (integers to a negative power) and one operand is such an Array, the result's fill is + never seen either, and 0 stands in for it. + """ + # numpy picks its loop by its operands' shapes (clip breaks ties by its 0-d path + # otherwise): a fill stands in for an array as a 1-element array, a 0-d Array stays 0-d, + # and a Python scalar stays weak + args = [ + ( + (np.reshape(x._fill, (1,)) if x.ndim else np.asarray(x._fill)) + if isinstance(x, Array) + else x + ) + for x in operands + ] + try: + return np.asarray(ufunc(*args)).reshape(-1)[0] + except (ValueError, ArithmeticError): + if any(isinstance(x, Array) and x.ndim and x._T().nvals == x.size for x in operands): + return _utils.scalar(0, rdt) + raise + + +def _result(T, fill): + if T.dtype == BOOL and T.ndim and T.nvals: + if T.nvals == T.size: + T, fill = _normalize_bool(T, fill) + elif fill and 2 * T.nvals > T.size: + T, fill = _majority_false(T) + return Array._new(T, fill) + + +def _majority_false(T): + """A bool Tensor with fill True that stores mostly False entries, more than half its + positions (``equal(deg, 0)`` of a ``deg`` that leaves a few zero degrees unstored): + the same values with fill False, storing only the True positions (its unstored ones and + its stored True entries). O(size), which is under twice its entries. Else as it is. + """ + ntrue = T.reduce_scalar(monoid.plus[INT64]).new().value + if 2 * (T.nvals - ntrue) <= T.size: + return T, True + R = Tensor(BOOL, T.shape) + R(~T.S)[...] << True + if ntrue: + R(T.select(select.valueeq, True).new().S)[...] << True + return R, False + + +def _normalize_bool(T, fill): + """A bool Tensor that stores every position never shows its fill, so it takes fill False + (the fill contractions and masks work with; a fill True made a mostly-True mask times a + sparse matrix dense). Its False entries are dropped when at most a quarter of the + entries are True (``equal(deg, 0)`` of a fully stored ``deg`` then stores only its True + entries). Otherwise they stay; dropping more costs about as much as the comparison and + saves a later consumer little. + + A sample of the leading entries says whether dropping is likely: a mostly-True result pays + only for the sample (a count over all 2e7 values cost as much as the comparison at 1 + thread). Then one select of the True values decides, and is the result when few. + """ + n = T.nvals + rows = -(-_SAMPLE * T.shape[0] // n) # (leading rows holding about _SAMPLE entries) + if rows < T.shape[0]: + sample = T[(slice(0, rows),) + (slice(None),) * (T.ndim - 1)].new() + if 3 * sample.reduce_scalar(monoid.plus[INT64]).new().value > sample.nvals: + return T, False # (over a third True in the sample: dropping is unlikely) + S = T.select(select.valueeq, True).new() + return (S if 4 * S.nvals <= n else T), False + + +_SAMPLE = 2**16 + + +def _zero_d(dtype, fill): + return Array._new(_utils.empty_0d(dtype), fill) + + +# Unary functions +_ALL_FACTS = ("finite", "nan", "zero") + + +def _unary(name, ufunc, x, gb=None, *, const=None, keeps=()): + """``ufunc`` of every position of ``x``. + + ``gb(dtype)`` gives a GraphBLAS UnaryOp equal to ``ufunc`` on that dtype, a callable + taking the Tensor, or None; ``const(dtype)`` gives the value of every position when the + result does not depend on the entries (``isnan`` of an integer array), else None. + ``keeps``: the facts about the stored values (``_utils.fact``) that hold for the result + when they hold for ``x``. + """ + x = _as_operand(x, name) + if not isinstance(x, Array): + from ._creation import asarray + + x = asarray(x) + op = gb(x.dtype) if gb is not None else None + with np.errstate(all="ignore"): + rdt = ufunc(np.empty(0, x.dtype)).dtype + if getattr(op, "opclass", None) == "UnaryOp": + # The fill by the op that maps the entries, so that an entry equal to the fill maps + # to the new fill (libm's rounding may differ from numpy's by an ulp) + s = Scalar.from_value(x._fill, gb_dtype(x.dtype)) + fill = s.apply(op).new(dtype=gb_dtype(rdt)).value + else: + fill = _fill_of(ufunc, rdt, (x,)) + check_supported(rdt, f"{name} of {x.dtype}") + T = x._T() + if T.ndim == 0: + return _zero_d(rdt, fill) + if const is not None and (value := const(x.dtype)) is not None: + return _result(_utils.empty_tensor(rdt, T.shape), value) + if T.nvals == 0: + return _result(_utils.empty_tensor(rdt, T.shape), fill) + if op is None: + R = _utils.map_values(ufunc, [T], rdt) + elif op == "copy": + R = T.dup(dtype=gb_dtype(rdt)) + elif callable(op) and not hasattr(op, "opclass"): + R = op(T) + else: + R = T.apply(op).new(dtype=gb_dtype(rdt)) + return _utils.keep_facts(x, _result(R, fill), keeps) + + +def _typed(op, kinds): + """``gb`` for _unary: ``op`` typed for dtypes of these kinds, else numpy.""" + + def get(dtype): + return op[gb_dtype(dtype)] if dtype.kind in kinds else None + + return get + + +def _copy_if(kinds): + def get(dtype): + return "copy" if dtype.kind in kinds else None + + return get + + +def _first_of(*getters): + def get(dtype): + for g in getters: + if (op := g(dtype)) is not None: + return op + return None + + return get + + +def abs(x, /): # noqa: A001 (the Array API's name) + # (GraphBLAS's integer abs of the most negative value is undefined in C; numpy wraps) + gb = _first_of(_typed(unary.abs, _FLOAT), _copy_if("ub")) + # (|z| of a finite complex z may overflow, and |inf + nan j| is inf: only "zero" holds) + keeps = ("zero",) if isinstance(x, Array) and x.dtype.kind == "c" else _ALL_FACTS + return _unary("abs", np.abs, x, gb, keeps=keeps) + + +def negative(x, /): + return _unary("negative", np.negative, x, _typed(unary.ainv, "iufc"), keeps=_ALL_FACTS) + + +def positive(x, /): + return _unary("positive", np.positive, x, _copy_if("iufc"), keeps=_ALL_FACTS) + + +# (a rounding keeps a value finite, or NaN; it may make a zero) +_ROUNDING_KEEPS = ("finite", "nan") + + +def sqrt(x, /): + # IEEE square root is correctly rounded, in C as in numpy + return _unary("sqrt", np.sqrt, x, _typed(unary.sqrt, _FLOAT)) + + +def floor(x, /): + gb = _first_of(_typed(unary.floor, _FLOAT), _copy_if("biu")) + return _unary("floor", np.floor, x, gb, keeps=_ROUNDING_KEEPS) + + +def ceil(x, /): + gb = _first_of(_typed(unary.ceil, _FLOAT), _copy_if("biu")) + return _unary("ceil", np.ceil, x, gb, keeps=_ROUNDING_KEEPS) + + +def trunc(x, /): + gb = _first_of(_typed(unary.trunc, _FLOAT), _copy_if("biu")) + return _unary("trunc", np.trunc, x, gb, keeps=_ROUNDING_KEEPS) + + +def round(x, /): # noqa: A001 (the Array API's name) + # numpy rounds half to even; C's round (GraphBLAS's) rounds half away from zero + return _unary("round", np.round, x, _copy_if("biu"), keeps=_ROUNDING_KEEPS) + + +def _const_unless(kinds, value): + def get(dtype): + return None if dtype.kind in kinds else value + + return get + + +def isnan(x, /): + return _unary( + "isnan", np.isnan, x, _typed(unary.isnan, _INEXACT), const=_const_unless("fc", False) + ) + + +def isinf(x, /): + return _unary( + "isinf", np.isinf, x, _typed(unary.isinf, _INEXACT), const=_const_unless("fc", False) + ) + + +def isfinite(x, /): + return _unary( + "isfinite", + np.isfinite, + x, + _typed(unary.isfinite, _INEXACT), + const=_const_unless("fc", True), + ) + + +def logical_not(x, /): + return _unary("logical_not", np.logical_not, x, _typed(unary.lnot, "b")) + + +def bitwise_invert(x, /): + gb = _first_of(_typed(unary.lnot, "b"), _typed(unary.bnot, _INTEGRAL)) + return _unary("bitwise_invert", np.invert, x, gb) + + +def square(x, /): + def gb(dtype): + if dtype.kind not in "iuf": + return None + # x * x on the stored entries, as numpy's square computes it + return lambda T: T.ewise_mult(T, binary.times[gb_dtype(dtype)]).new() + + return _unary("square", np.square, x, gb) + + +def reciprocal(x, /): + return _unary("reciprocal", np.reciprocal, x, _typed(unary.minv, _FLOAT)) + + +def conj(x, /): + gb = _first_of(_typed(unary.conj, "c"), _copy_if(_REAL)) + return _unary("conj", np.conj, x, gb, keeps=_ALL_FACTS) + + +def real(x, /): + return _unary("real", np.real, x, _first_of(_typed(unary.creal, "c"), _copy_if(_REAL))) + + +def imag(x, /): + def const(dtype): + return None if dtype.kind == "c" else 0 + + return _unary("imag", np.imag, x, _typed(unary.cimag, "c"), const=const) + + +def sign(x, /): + return _unary("sign", np.sign, x) + + +def signbit(x, /): + return _unary("signbit", np.signbit, x) + + +def _libm_unary(name, ufunc, op): + """``ufunc`` computed by the SuiteSparse unary op ``op`` (C's libm, in parallel) for real + floating dtypes, by numpy otherwise: within libm's rounding of numpy's answer, special + values (NaN, inf, signed zeros) exact (DESIGN D3b). + + Complex values stay on numpy: C libraries treat the many special cases of complex + arguments differently from platform to platform, and SuiteSparse's complex expm1 and + log1p are formulas of its own, inexact near 0. Integer inputs stay too (numpy's result + dtype there, such as float16 for int8, is its own). + """ + + def func(x, /): + return _unary(name, ufunc, x, _typed(op, _FLOAT)) + + func.__name__ = name + func.__qualname__ = name + return func + + +acos = _libm_unary("acos", np.arccos, unary.acos) +acosh = _libm_unary("acosh", np.arccosh, unary.acosh) +asin = _libm_unary("asin", np.arcsin, unary.asin) +asinh = _libm_unary("asinh", np.arcsinh, unary.asinh) +atan = _libm_unary("atan", np.arctan, unary.atan) +atanh = _libm_unary("atanh", np.arctanh, unary.atanh) +cos = _libm_unary("cos", np.cos, unary.cos) +cosh = _libm_unary("cosh", np.cosh, unary.cosh) +exp = _libm_unary("exp", np.exp, unary.exp) +log = _libm_unary("log", np.log, unary.log) +log2 = _libm_unary("log2", np.log2, unary.log2) +log10 = _libm_unary("log10", np.log10, unary.log10) +sin = _libm_unary("sin", np.sin, unary.sin) +sinh = _libm_unary("sinh", np.sinh, unary.sinh) +tan = _libm_unary("tan", np.tan, unary.tan) +tanh = _libm_unary("tanh", np.tanh, unary.tanh) +expm1 = _libm_unary("expm1", np.expm1, unary.expm1) +log1p = _libm_unary("log1p", np.log1p, unary.log1p) + + +# Binary functions +_MANTISSA = {"f": {4: 24, 8: 53}, "c": {8: 24, 16: 53}} + + +def _exact_cast(src, dst): + """Whether every value of dtype ``src`` converts to ``dst`` without change.""" + if src == dst or src.kind == "b": + return True + if src.kind in _INTEGRAL: + if dst.kind in _INTEGRAL: + return np.can_cast(src, dst, casting="safe") + if dst.kind in _INEXACT and dst != np.dtype(bool): + return src.itemsize * 8 - (src.kind == "i") <= _MANTISSA[dst.kind][dst.itemsize] + return False + if src.kind == "f" and dst.kind in _INEXACT: + return _MANTISSA["f"][src.itemsize] <= _MANTISSA[dst.kind][dst.itemsize] + if src.kind == "c" and dst.kind == "c": + return src.itemsize <= dst.itemsize + return False + + +def _python_scalar_fits(value, dtype): + """Whether numpy converts the Python scalar ``value`` to ``dtype`` without changing it + (out-of-range Python ints get special treatment in numpy's comparisons). + """ + if dtype.kind in _INTEGRAL and isinstance(value, int): + info = np.iinfo(dtype) + return info.min <= value <= info.max + return True + + +class _Binary: + """How to compute one binary function. + + ``kind`` is "arith" (the op works in the result dtype), "cmp" (in the promoted dtype of + the operands, with a bool result) or "logical" (in bool). ``gb(dtype)`` gives the + GraphBLAS operator that equals ``ufunc`` on operands of the dtype it works in, or None. + ``gb_values(a, b, dtype)``, in place of ``gb``, may look at the operands' values too. + ``annihilator(a, b, fill_a, fill_b, dtype)`` says whether the union equals the + intersection (``ewise_mult``): whether each operand's fill makes the result the result's + fill whatever the other operand holds. + """ + + def __init__(self, name, ufunc, kind="arith", gb=None, annihilator=None, gb_values=None): + self.name = name + self.ufunc = ufunc + self.kind = kind + self.gb = gb + self.annihilator = annihilator + self.gb_values = gb_values + + +def _binary(spec, x1, x2): + a = _as_operand(x1, spec.name) + b = _as_operand(x2, spec.name) + if spec.name == "multiply" and (out := _lazy.masked_multiply(a, b)) is not None: + return out # (an SDDMM) + ufunc = spec.ufunc + if not isinstance(a, Array) and not isinstance(b, Array): + from ._creation import asarray + + a = asarray(a) + with np.errstate(all="ignore"): + rdt = ufunc(_probe(a), _probe(b)).dtype + fill = _fill_of(ufunc, rdt, (a, b)) + check_supported(rdt, f"{spec.name} of {_dtype_name(a)} and {_dtype_name(b)}") + a_arr = isinstance(a, Array) and a.ndim > 0 + b_arr = isinstance(b, Array) and b.ndim > 0 + if not a_arr and not b_arr: + return _zero_d(rdt, fill) + if spec.kind == "arith": + ct = rdt + elif spec.kind == "cmp": + # (a Python int out of range of an integer dtype keeps that dtype here; _gb_op then + # leaves the comparison to numpy, which compares such an int exactly) + ct = np.result_type(_probe(a), _probe(b)) + else: + ct = np.dtype(bool) + op = _gb_op(spec, a, b, ct) + if a_arr and b_arr: + return _binary_arrays(spec, a, b, op, ct, rdt, fill) + if a_arr: + T, other, left = a._T(), b, False + else: + T, other, left = b._T(), a, True + if T.nvals == 0: + return _result(_utils.empty_tensor(rdt, T.shape), fill) + if op is not None: + s = gb_scalar(_fill_operand(other), ct) + expr = T.apply(op, left=s) if left else T.apply(op, right=s) + return _result(expr.new(dtype=gb_dtype(rdt)), fill) + s = _fill_operand(other) + if left: + R = _utils.map_values(lambda v: ufunc(s, v), [T], rdt) + else: + R = _utils.map_values(lambda v: ufunc(v, s), [T], rdt) + return _result(R, fill) + + +def _dtype_name(x): + return str(x.dtype) if isinstance(x, Array) else type(x).__name__ + + +def _gb_op(spec, a, b, ct): + """The typed GraphBLAS op for ``spec`` working in ``ct``, when it gives numpy's answer; + else None. + """ + if spec.gb_values is not None: + return spec.gb_values(a, b, ct) + if spec.gb is None: + return None + for x in (a, b): + if isinstance(x, Array): + if spec.kind != "arith" and not _exact_cast(x.dtype, ct): + return None + elif not _python_scalar_fits(x, ct): + return None + return spec.gb(ct) + + +def _binary_arrays(spec, a, b, op, ct, rdt, fill): + Ta, Tb = a._T(), b._T() + try: + shape = Ta.shape if Ta.shape == Tb.shape else np.broadcast_shapes(Ta.shape, Tb.shape) + except ValueError: + raise ValueError( + f"{spec.name}: shapes {Ta.shape} and {Tb.shape} do not broadcast" + ) from None + if Ta.nvals == 0 and Tb.nvals == 0: + return _result(_utils.empty_tensor(rdt, shape), fill) + if op is not None: + fa = _utils.scalar(a._fill, ct) + fb = _utils.scalar(b._fill, ct) + if spec.annihilator is not None: + # An operand that stores every position never shows its fill: take it as 0 + zero = _zero(ct) + fa_seen = zero if _stores_all(a) else fa + fb_seen = zero if _stores_all(b) else fb + if spec.annihilator(a, b, fa_seen, fb_seen, ct): + # Where either operand holds its fill the result is the fill: the intersection + _utils.mult_expansion_check(Ta, Tb, shape) + R = Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)) + with np.errstate(all="ignore"): + return _result(R, spec.ufunc(fa_seen, fb_seen)) + if (e := _absorbing(spec, a, b, fa, fb, ct)) is not None: + # Each operand's fill absorbs the op: where either holds its fill the result is + # that fill, so the intersection, which never replicates a broadcast operand + _utils.mult_expansion_check(Ta, Tb, shape) + R = Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)) + return _result(R, _utils.scalar(e, rdt)) + if spec.name in {"multiply", "divide"} and ct.kind == "f": + out = _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape) + if out is not None: + return out + _utils.expansion_check(Ta, Tb, shape) + expr = Ta.ewise_union(Tb, op, gb_scalar(fa, ct), gb_scalar(fb, ct), allow_expansion=True) + return _result(expr.new(dtype=gb_dtype(rdt)), fill) + # numpy over the values: both operands' values on the union of entries, aligned + _utils.expansion_check(Ta, Tb, shape) + da, db = a.dtype, b.dtype + A = Ta.ewise_union( + Tb, binary.first[gb_dtype(da)], gb_scalar(a._fill, da), gb_scalar(0, da), + allow_expansion=True, + ).new() # fmt: skip + B = Ta.ewise_union( + Tb, binary.second[gb_dtype(db)], gb_scalar(0, db), gb_scalar(b._fill, db), + allow_expansion=True, + ).new() # fmt: skip + return _result(_utils.map_values(spec.ufunc, [A, B], rdt), fill) + + +def _absorbing_values(name, ct): + """The values ``e`` with ``op(e, y) == op(y, e) == e`` for the binary ``name`` working in + ``ct``, for every y but those ``_absorbing`` excludes (NaN, the opposite infinity). + """ + if name in {"add", "maximum", "logical_or"} and ct.kind == "b": + return (True,) + if name == "minimum" and ct.kind == "b": + return (False,) + if ct.kind == "f": + return {"add": (np.inf, -np.inf), "maximum": (np.inf,), "minimum": (-np.inf,)}.get(name, ()) + if ct.kind in "iu": + info = np.iinfo(ct) + return {"maximum": (info.max,), "minimum": (info.min,)}.get(name, ()) + return () + + +def _absorbing(spec, a, b, fa, fb, ct): + """The value both operands' fills take when it absorbs the op (inf for add and maximum, + -inf for add and minimum, True for logical_or, the integer extremes), so that the result + is the intersection of their entries with that fill; else None. An operand that stores + every position never shows its fill, so it counts as having that value, but not both. + Floats: no operand may hold NaN (numpy propagates it) or, for add, the opposite infinity + (inf + -inf is NaN): cached facts. + """ + full_a, full_b = _stores_all(a), _stores_all(b) + if full_a and full_b: + return None # (the union is the intersection already) + for e in _absorbing_values(spec.name, ct): + e = _utils.scalar(e, ct) + if not (full_a or _utils.same_value(fa, e)) or not (full_b or _utils.same_value(fb, e)): + continue + if ct.kind == "f": + opposite = "neginf" if e > 0 else "posinf" + for x in (a, b): + if _utils.fact(x, "nan") or (spec.name == "add" and _utils.fact(x, opposite)): + return None + return e + return None + + +def _power(x1, x2): + """The power by numpy, but where numpy takes a square root for an exponent of 0.5 that + repeats (a scalar, or an exponent of stride 0): the Array API's pow(-inf, 0.5) is +inf and + pow(-0.0, 0.5) is +0.0, where sqrt gives NaN and -0.0 (review finding L11). + """ + out = np.power(x1, x2) + if out.dtype.kind == "f": + x1 = np.asarray(x1) + fix = (np.asarray(x2) == 0.5) & ((x1 == -np.inf) | (x1 == 0)) + if fix.any(): + out = np.where(fix, np.where(x1 == 0, 0.0, np.inf), out).astype(out.dtype) + return out + + +def _typed_for(op_by_kind): + """``gb`` for _Binary: ``op_by_kind[dtype.kind]`` typed for the dtype, else None.""" + + def get(dtype): + op = op_by_kind.get(dtype.kind) + return None if op is None else op[gb_dtype(dtype)] + + return get + + +def _zero(dtype): + return _utils.scalar(0, dtype) + + +def _stores_all(x): + """Whether ``x`` stores every position, so that its fill is never seen.""" + return x._T().nvals == x.size + + +def _both_fills_zero(a, b, fa, fb, ct): + zero = _zero(ct) + return _utils.same_value(fa, zero) and _utils.same_value(fb, zero) + + +def _multiply_intersects(a, b, fa, fb, ct): + """Whether ``a * b`` stores only the intersection of entries: both fills are 0 and, for + floats, no stored value is inf or NaN (whose product with 0.0 is NaN). + + A ``-0.0`` product with an unstored zero reads as the fill, 0.0: the sign of a zero at an + unstored position is unspecified (DESIGN D3a, D3b), and keeping it would make a row + scaling ``d[:, None] * A`` with signed ``d`` store every position. + """ + if not _both_fills_zero(a, b, fa, fb, ct): + return False + if ct.kind != "f": + return True + return _utils.fact(a, "finite") and _utils.fact(b, "finite") + + +def _one_sided(spec, a, b, fa, fb, op, ct, rdt, shape): + """``a * b`` or ``a / b`` on the entries of one operand only, when that operand's fill + is 0 and every value of the other makes 0 times (or divided by) it a zero: a row scaling + ``A * d[:, None]`` or ``A / d`` stores A's entries. The other operand's fill is computed + at the first operand's entries (a broadcast operand smaller than them is stored in full + first, so that it is not replicated). None when that does not hold. + """ + sides = [(a, b)] if spec.name == "divide" else [(a, b), (b, a)] + for X, Y in sides: + fx = fa if X is a else fb + fy = fb if X is a else fa + if _stores_all(X) or X._T().shape != shape or fx != 0: + continue + if _stores_all(Y): + fy = _utils.scalar(1, ct) # (never seen) + if not _zero_times(spec, Y, fy): + continue + Tx, Ty = X._T(), Y._T() + with np.errstate(all="ignore"): + fill = spec.ufunc(fx, fy) if X is a else spec.ufunc(fy, fx) + if Ty.shape != shape and not _stores_all(Y) and Tx.nvals < Y.size: + # Y broadcasts, shows its fill, and is larger than X's entries: its value at each + # of X's entries, gathered from Y itself at the projected coordinates, O(entries + # of X). (Storing Y in full cost O(size of Y), and the union replicated Y over its + # broadcast axes: exp(v)[None, :] times X of shape (1000, 1e6) with 4 entries made + # v 1000 times.) A Y no larger is stored in full below: the broadcast ewise_mult + # then needs no coordinates (a row scaling of 1e7 entries: 5 ms, against 143) + Ty = _gather(Tx, Y, shape) + elif not _stores_all(Y) and (Y.size <= Tx.nvals + 1024 or Ty.shape != shape): + # (O(size of Y): every position of Y stored, so the product is an intersection) + other = 0 if Y._fill != 0 else 1 + Ty = _utils.with_fill(Y, other)._T() + Ta, Tb = (Tx, Ty) if X is a else (Ty, Tx) + if Ty.nvals == Y.size: + _utils.mult_expansion_check(Ta, Tb, shape) + return _result(Ta.ewise_mult(Tb, op).new(dtype=gb_dtype(rdt)), fill) + _utils.expansion_check(Ta, Tb, shape) + R = Tensor(gb_dtype(rdt), shape) + fa_s, fb_s = gb_scalar(fa, ct), gb_scalar(fb, ct) + R(Tx.S) << Ta.ewise_union(Tb, op, fa_s, fb_s, allow_expansion=True) + return _result(R, fill) + return None + + +def _gather(Tx, Y, shape): + """A Tensor with ``Tx``'s entries (of ``shape``) holding the values of ``Y`` (an Array + that broadcasts to ``shape``) there: its stored value, or its fill. + """ + coords = Tx.to_coo(values=False)[0] + Ty = Y._T() + lead = len(shape) - Y.ndim + # (a broadcast axis of Y has size 1: index 0 there) + index = tuple( + coords[lead + i] if n != 1 else np.zeros(len(coords[0]), np.int64) + for i, n in enumerate(Ty.shape) + ) + values = Ty.vindex[index].to_dense(fill_value=Y._fill) + return Tensor.from_coo(coords, values, Ty.dtype, shape=shape) + + +def _zero_times(spec, y, fy): + """Whether 0 * v (or 0 / v) is a zero for every value v of ``y`` (its fill ``fy`` too): + finite values (for divide, also no 0). NaN never qualifies; the sign of the zero does not + matter (an unstored position reads the fill). + """ + if spec.name == "divide": + ok = fy != 0 and not np.isnan(fy) + return ok and not _utils.fact(y, "zero") and not _utils.fact(y, "nan") + return bool(np.isfinite(fy)) and _utils.fact(y, "finite") + + +def _max_min_gb(op): + """GraphBLAS's max and min of floats are C's fmax and fmin, which skip NaN; numpy's + propagate it. They agree when neither operand holds NaN (a cached fact), fills included. + Either may take either zero of 0.0 and -0.0: the Array API leaves that order unspecified. + """ + + def get(a, b, ct): + if ct.kind in "biu": + return op[gb_dtype(ct)] + if ct.kind != "f": + return None + for x in (a, b): + if isinstance(x, Array): + if np.isnan(x._fill) or _utils.fact(x, "nan"): + return None + elif isinstance(x, float) and np.isnan(x): + return None + return op[gb_dtype(ct)] + + return get + + +_SPECS = { + "add": _Binary( + "add", np.add, gb=_typed_for({"b": binary.lor, **dict.fromkeys("iufc", binary.plus)}) + ), + "subtract": _Binary( + "subtract", np.subtract, gb=_typed_for(dict.fromkeys("iufc", binary.minus)) + ), + # Complex products and quotients are computed by numpy: C's formulas round otherwise. + # A float zero fill annihilates only nonnegative finite values: 0.0 * -1.0 is -0.0, and + # 0.0 * inf is NaN. + "multiply": _Binary( + "multiply", + np.multiply, + gb=_typed_for({"b": binary.land, **dict.fromkeys("iuf", binary.times)}), + annihilator=_multiply_intersects, + ), + "divide": _Binary("divide", np.divide, gb=_typed_for({"f": binary.truediv})), + "equal": _Binary("equal", np.equal, "cmp", gb=_typed_for(dict.fromkeys("biufc", binary.eq))), + "not_equal": _Binary( + "not_equal", np.not_equal, "cmp", gb=_typed_for(dict.fromkeys("biufc", binary.ne)) + ), + "less": _Binary("less", np.less, "cmp", gb=_typed_for(dict.fromkeys("biuf", binary.lt))), + "less_equal": _Binary( + "less_equal", np.less_equal, "cmp", gb=_typed_for(dict.fromkeys("biuf", binary.le)) + ), + "greater": _Binary( + "greater", np.greater, "cmp", gb=_typed_for(dict.fromkeys("biuf", binary.gt)) + ), + "greater_equal": _Binary( + "greater_equal", np.greater_equal, "cmp", gb=_typed_for(dict.fromkeys("biuf", binary.ge)) + ), + "logical_and": _Binary( + "logical_and", + np.logical_and, + "logical", + gb=lambda dt: binary.land[gb_dtype(dt)], + annihilator=_both_fills_zero, + ), + "logical_or": _Binary( + "logical_or", np.logical_or, "logical", gb=lambda dt: binary.lor[gb_dtype(dt)] + ), + "logical_xor": _Binary( + "logical_xor", np.logical_xor, "logical", gb=lambda dt: binary.lxor[gb_dtype(dt)] + ), + "bitwise_and": _Binary( + "bitwise_and", + np.bitwise_and, + gb=_typed_for({"b": binary.land, "i": binary.band, "u": binary.band}), + annihilator=_both_fills_zero, + ), + "bitwise_or": _Binary( + "bitwise_or", + np.bitwise_or, + gb=_typed_for({"b": binary.lor, "i": binary.bor, "u": binary.bor}), + ), + "bitwise_xor": _Binary( + "bitwise_xor", + np.bitwise_xor, + gb=_typed_for({"b": binary.lxor, "i": binary.bxor, "u": binary.bxor}), + ), + "maximum": _Binary("maximum", np.maximum, gb_values=_max_min_gb(binary.max)), + "minimum": _Binary("minimum", np.minimum, gb_values=_max_min_gb(binary.min)), + "copysign": _Binary("copysign", np.copysign, gb=_typed_for({"f": binary.copysign})), + # numpy's own definitions: floored division and remainder (with its results for division + # by zero), integer powers, shifts past the width, and libm-free special values + "floor_divide": _Binary("floor_divide", np.floor_divide), + "remainder": _Binary("remainder", np.remainder), + "pow": _Binary("pow", _power), + "atan2": _Binary("atan2", np.arctan2), + "hypot": _Binary("hypot", np.hypot), + "logaddexp": _Binary("logaddexp", np.logaddexp), + "nextafter": _Binary("nextafter", np.nextafter), + "bitwise_left_shift": _Binary("bitwise_left_shift", np.left_shift), + "bitwise_right_shift": _Binary("bitwise_right_shift", np.right_shift), +} + + +def binary_numpy(name, func, x1, x2): + """``func`` (a numpy function of two arrays, which broadcasts) of two operands with the + fill rules of a binary function, computed by numpy over the values (for internal use). + """ + return _binary(_Binary(name, func), x1, x2) + + +def _make_binary(name): + spec = _SPECS[name] + + def func(x1, x2, /): + return _binary(spec, x1, x2) + + func.__name__ = name + func.__qualname__ = name + return func + + +add = _make_binary("add") +subtract = _make_binary("subtract") +multiply = _make_binary("multiply") +divide = _make_binary("divide") +equal = _make_binary("equal") +not_equal = _make_binary("not_equal") +less = _make_binary("less") +less_equal = _make_binary("less_equal") +greater = _make_binary("greater") +greater_equal = _make_binary("greater_equal") +logical_and = _make_binary("logical_and") +logical_or = _make_binary("logical_or") +logical_xor = _make_binary("logical_xor") +bitwise_and = _make_binary("bitwise_and") +bitwise_or = _make_binary("bitwise_or") +bitwise_xor = _make_binary("bitwise_xor") +maximum = _make_binary("maximum") +minimum = _make_binary("minimum") +copysign = _make_binary("copysign") +floor_divide = _make_binary("floor_divide") +remainder = _make_binary("remainder") +pow = _make_binary("pow") # noqa: A001 (the Array API's name) +atan2 = _make_binary("atan2") +hypot = _make_binary("hypot") +logaddexp = _make_binary("logaddexp") +nextafter = _make_binary("nextafter") +bitwise_left_shift = _make_binary("bitwise_left_shift") +bitwise_right_shift = _make_binary("bitwise_right_shift") + + +def clip(x, /, min=None, max=None): + """``x`` clipped as numpy's ``clip`` does (NaN propagates; a tie gives the bound), in + the dtype of ``x``, over the union of the entries of ``x`` and of array bounds. + """ + if not isinstance(x, Array): + raise TypeError(f"clip needs an Array, not {type(x).__name__}") + lo = _clip_bound(min, x) + hi = _clip_bound(max, x) + if lo is None and hi is None: + return x._copy() + if all(not isinstance(b, Array) or b.ndim == 0 for b in (lo, hi)): + lo_v, hi_v = (b._fill if isinstance(b, Array) else b for b in (lo, hi)) + out = _unary("clip", lambda v: np.clip(v, lo_v, hi_v), x) + else: + out = _clip_arrays(x, lo, hi) + return out if out.dtype == x.dtype else cast(out, x.dtype) + + +def _clip_arrays(x, lo, hi): + """Clip with array bounds as numpy does, over the union of the entries (aligned values); + the result has the shape the three broadcast to. + """ + operands = (x, lo, hi) + shape = np.broadcast_shapes(*[b.shape for b in operands if isinstance(b, Array)]) + is_array = [isinstance(b, Array) and b.ndim > 0 for b in operands] + + def value(b): # (a value where nothing is stored, typed as numpy types it) + return np.asarray(b._fill) if isinstance(b, Array) else b + + def probe(b): + return np.empty(0, b.dtype) if isinstance(b, Array) else b + + def value1(b): # (an array's fill as a 1-element array: numpy's 0-d clip breaks ties + # otherwise) + if isinstance(b, Array): + return np.reshape(b._fill, (1,)) if b.ndim else np.asarray(b._fill) + return b + + with np.errstate(all="ignore"): + rdt = np.clip(probe(x), probe(lo), probe(hi)).dtype + fill = np.clip(value1(x), value1(lo), value1(hi))[0] + if math.prod(shape) == 0: + return _result(_utils.empty_tensor(rdt, shape), fill) + + def func(*arrays): + it = iter(arrays) + v, lo_v, hi_v = ( + next(it) if arr else value(b) for b, arr in zip(operands, is_array, strict=True) + ) + return np.clip(v, lo_v, hi_v) + + tensors = _utils.aligned([b for b, arr in zip(operands, is_array, strict=True) if arr], shape) + return _result(_utils.map_values(func, tensors, rdt), fill) + + +def _clip_bound(bound, x): + """A Python integer bound clamped to the range of an integer ``x`` (the Array API's clip + keeps the dtype of ``x``, and numpy refuses an integer out of its range). + """ + if isinstance(bound, (np.generic, np.ndarray)): + return _as_operand(bound, "clip") + if isinstance(bound, (builtins.bool, int, float)): + if x.dtype.kind in _INTEGRAL and isinstance(bound, float): + return bound + if x.dtype.kind in _INTEGRAL and not _python_scalar_fits(bound, x.dtype): + info = np.iinfo(x.dtype) + return builtins.min(builtins.max(bound, int(info.min)), int(info.max)) + return bound + + +# Casting (astype) +def _gb_cast_ok(src, dst): + """Whether GraphBLAS's typecast from ``src`` to ``dst`` gives numpy's ``astype``. + + C converts integers and floats as numpy does; float to integer differs where the value + does not fit (NaN, inf, out of range: numpy's result is undefined there, GraphBLAS + saturates), and complex to real is numpy's own rule. + """ + if src.kind == "c": + return dst.kind in "cb" + if src.kind == "f": + return dst.kind in "fcb" + return True + + +def _astype(values, dt): + # numpy's astype from complex to real keeps the real part (and warns that it does) + if values.dtype.kind == "c" and dt.kind not in "cb": + values = values.real + return values.astype(dt) + + +def cast(x, dtype): + """``x.astype(dtype)`` by numpy's rules, entries and fill alike.""" + dt = np.dtype(dtype) + with np.errstate(all="ignore"): + fill = _astype(np.asarray(x._fill), dt) + T = x._T() + if T.ndim == 0: + return _zero_d(dt, fill) + if T.nvals == 0: + return _result(_utils.empty_tensor(dt, T.shape), fill) + if _gb_cast_ok(x.dtype, dt): + return _result(T.dup(dtype=gb_dtype(dt)), fill) + return _result(_utils.map_values(lambda v: _astype(v, dt), [T], dt), fill) diff --git a/graphblas/array_api/_indexing.py b/graphblas/array_api/_indexing.py new file mode 100644 index 000000000..7d7b51bac --- /dev/null +++ b/graphblas/array_api/_indexing.py @@ -0,0 +1,511 @@ +"""Indexing: ``x[key]``, ``x[key] = value``, ``take`` and ``take_along_axis``. + +Keys follow numpy: integers, slices, ``...`` and ``None`` (basic indexing, extracted in +GraphBLAS), and integer and boolean arrays (advanced indexing, gathered from the stored +entries, so a position the array does not store is the fill in the result too). The Array API +standard specifies an index of integers and integer arrays only (vectorized indexing), or one +boolean array alone; other combinations follow numpy. + +Costs: basic indexing is O(entries in the region). An integer index array is read densely +(its size is the result's); a boolean index array costs O(its entries) when its fill is +False, O(its size) when its fill is True. Assigning a scalar other than the fill, or an array +whose fill is not ``x``'s, to a region stores every position of the region; assigning by +several index arrays writes every selected position (O(selection)). +""" + +import builtins +import math +import operator + +import numpy as np + +from ..core.tensor import Tensor +from . import _lazy, _utils +from ._array import Array +from ._utils import gb_dtype + +__all__ = ["take", "take_along_axis"] + + +# Keys +class _Adv: + """An integer index array (advanced indexing) for one axis, nonnegative.""" + + __slots__ = ("index",) + + def __init__(self, index): + self.index = index + + +class _NewAxis: + """``None``, or a 0-d boolean index: a new axis of length 1 (0 for False).""" + + __slots__ = ("length",) + + def __init__(self, length=1): + self.length = length + + +class _Gap: + """A ``...`` that covers no axis. It takes no axis, but separates the advanced indices on + either side of it, as in numpy: ``x[:, 1, ..., [0, 2]]`` of a 3-d ``x`` puts the + broadcast advanced axis first. + """ + + __slots__ = () + + +_GAP = _Gap() + + +def _is_int(k): + return isinstance(k, (int, np.integer)) and not isinstance(k, (bool, np.bool_)) + + +def _index_array(k): + """An index of array type: a boolean Array as it is, else a numpy array; None when ``k`` + is not an array. + """ + if isinstance(k, Array): + if k.dtype.kind == "b": + return k + if k.dtype.kind not in "iu": + raise IndexError( + f"arrays used as indices must be of integer or boolean type: {k.dtype}" + ) + return k.todense() + if isinstance(k, (list, np.ndarray)): + arr = np.asarray(k) + if arr.dtype.kind not in "biu": + if arr.size == 0: + return arr.astype(np.int64) + raise IndexError( + f"arrays used as indices must be of integer or boolean type: {arr.dtype}" + ) + return arr + return None + + +def _mask_coords(mask, shape): + """The coordinates, in C order, where the boolean ``mask`` (an Array or ndarray) is True; + ``shape`` is that of the axes it covers. + """ + if mask.shape != tuple(shape): + # (the standard lets a dimension of the mask be 0 instead of the array's) + if any(m not in (n, 0) for m, n in zip(mask.shape, shape, strict=True)): + raise IndexError(f"boolean index of shape {mask.shape} does not match {shape}") + return tuple(np.empty(0, np.int64) for _ in mask.shape) + if isinstance(mask, np.ndarray): + return np.nonzero(mask) + if not mask._fill: + T = mask._T() + if T.nvals == 0: + return tuple(np.empty(0, np.int64) for _ in mask.shape) + # (stored entries may be False, equal to the fill) + return T.select("==", True).new().to_coo(values=False)[0] + _utils.check_dense(mask.size, "a boolean index whose fill is True") + return np.nonzero(mask.todense()) + + +def _normalize(x, key): + """``key`` as a list of items: an int, slice or _Adv for each axis of ``x``, and _NewAxis + items for new axes. A boolean array of M dimensions becomes M _Adv items. + """ + if not isinstance(key, tuple): + key = (key,) + items = [] + for k in key: + if k is None: + items.append(_NewAxis()) + elif k is Ellipsis or isinstance(k, slice): + items.append(k) + elif isinstance(k, (bool, np.bool_)): + items.append(_NewAxis(1 if k else 0)) + elif _is_int(k): + items.append(int(k)) + elif isinstance(k, Array) and k.ndim == 0: + if k.dtype.kind == "b": + items.append(_NewAxis(1 if k._fill else 0)) + elif k.dtype.kind in "iu": + items.append(operator.index(k)) + else: + raise IndexError(f"a 0-d index must be an integer or boolean, not {k.dtype}") + else: + arr = _index_array(k) + if arr is None: + raise IndexError(f"invalid index {k!r}") + items.append(arr) + if sum(k is Ellipsis for k in items if not isinstance(k, (np.ndarray, Array))) > 1: + raise IndexError("an index can only have a single ellipsis ('...')") + + def width(k): + if isinstance(k, (np.ndarray, Array)) and k.dtype.kind == "b": + return k.ndim + return 0 if k is Ellipsis or isinstance(k, _NewAxis) else 1 + + consumed = sum(width(k) for k in items) + if consumed > x.ndim: + raise IndexError(f"too many indices: {consumed} for an array of {x.ndim} dimensions") + i = next((i for i, k in enumerate(items) if k is Ellipsis), None) + if i is None: + items.extend([slice(None)] * (x.ndim - consumed)) + else: + items[i : i + 1] = [slice(None)] * (x.ndim - consumed) or [_GAP] + out = [] + axis = 0 + for k in items: + if isinstance(k, (_NewAxis, _Gap)): + out.append(k) + continue + if isinstance(k, (np.ndarray, Array)) and k.dtype.kind == "b": + if k.ndim == 0: # (a 0-d boolean numpy array) + out.append(_NewAxis(1 if k[()] else 0)) + continue + coords = _mask_coords(k, x.shape[axis : axis + k.ndim]) + out.extend(_Adv(np.asarray(c, dtype=np.int64)) for c in coords) + axis += k.ndim + continue + n = x.shape[axis] + if isinstance(k, np.ndarray): + idx = k.astype(np.int64, copy=False) + if idx.size and (idx.min() < -n or idx.max() >= n): + raise IndexError(f"index out of bounds for axis {axis} with size {n}") + out.append(_Adv(np.where(idx < 0, idx + n, idx))) + elif isinstance(k, int): + if not -n <= k < n: + raise IndexError(f"index {k} is out of bounds for axis {axis} with size {n}") + out.append(k % n) + else: + out.append(k) + axis += 1 + return out + + +def _has_adv(items): + return any(isinstance(k, _Adv) for k in items) + + +def _adjacent(items): + """Whether the advanced items (index arrays, with the ints among them) are next to each + other; numpy then keeps their axes in place, else puts them first. + """ + positions = [i for i, k in enumerate(items) if isinstance(k, (_Adv, int))] + return positions == list(range(positions[0], positions[-1] + 1)) + + +def _finish(out, items, nadv): + """``out`` holds the advanced axes (``nadv`` of them) first, then the sliced axes in + order: move the advanced ones where numpy puts them and insert the new axes. + """ + from ._manipulation import expand_dims, permute_dims + + comps = [] + placed = True + if nadv or _has_adv(items): + placed = not _adjacent(items) + if placed: + comps.append("adv") + for k in items: + if isinstance(k, (_NewAxis, slice)): + comps.append(k) + elif isinstance(k, _Gap): + continue + elif not placed: + comps.append("adv") + placed = True + order = [] + new = [] + rest = nadv + for c in comps: + if isinstance(c, str): + order.extend(range(nadv)) + elif isinstance(c, slice): + order.append(rest) + rest += 1 + else: + new.append((len(order) + len(new), c.length)) + if order != list(range(len(order))): + out = permute_dims(out, tuple(order)) + if new: + out = expand_dims(out, tuple(p for p, _ in new)) + if any(length == 0 for _, length in new): + empty = {p for p, length in new if length == 0} + out = out[tuple(slice(0, 0) if a in empty else slice(None) for a in range(out.ndim))] + return out + + +def _basic_shape(shape, axes_key): + return tuple( + len(range(*k.indices(n))) + for k, n in zip(axes_key, shape, strict=True) + if isinstance(k, slice) + ) + + +# getitem +def getitem(x, key): + items = _normalize(x, key) + axes_key = [k for k in items if not isinstance(k, (_NewAxis, _Gap))] + if _has_adv(items): + out, nadv = _advanced_extract(x, axes_key) + else: + # (a pending product slices its operands instead) + if (out := _lazy.sliced(x, axes_key)) is None: + out = _basic_extract(x, tuple(axes_key)) + nadv = 0 + return _finish(out, items, nadv) + + +def _basic_extract(x, axes_key): + T = x._T() + if x.ndim == 0 or all(k == slice(None) for k in axes_key): + return x._copy() + shape = _basic_shape(T.shape, axes_key) + if T.nvals == 0 or 0 in shape: + return Array._new(Tensor(T.dtype, shape), x._fill) + return Array._new(T[axes_key].new(), x._fill) + + +def _advanced_extract(x, axes_key): + """``x`` indexed by ``axes_key`` (one item per axis) with the advanced axes first, then the + sliced ones in order, and the number of advanced axes. + """ + sliced = tuple(k if isinstance(k, slice) else slice(None) for k in axes_key) + y = _basic_extract(x, sliced) if any(k != slice(None) for k in sliced) else x + adv_axes = [a for a, k in enumerate(axes_key) if not isinstance(k, slice)] + rest = [a for a, k in enumerate(axes_key) if isinstance(k, slice)] + indices = np.broadcast_arrays( + *[ + k.index if isinstance(k, _Adv) else np.asarray(k, np.int64) + for k in axes_key + if not isinstance(k, slice) + ] + ) + S2 = indices[0].shape + adv_shape = tuple(y.shape[a] for a in adv_axes) + rest_shape = tuple(y.shape[a] for a in rest) + T = y._T() + if T.nvals == 0 or math.prod(S2) == 0 or math.prod(rest_shape) == 0: + return Array._new(Tensor(T.dtype, S2 + rest_shape), x._fill), len(S2) + if not rest: + # Vectorized (pointwise) indexing of every axis + return Array._new(T.vindex[tuple(indices)], x._fill), len(S2) + Tp = T.permute_dims(tuple(adv_axes) + tuple(rest)) + if len(adv_axes) > 1: + flat = np.ravel_multi_index(tuple(indices), adv_shape) + Tp = Tp.reshape((math.prod(adv_shape), *rest_shape)) + else: + flat = indices[0] + return Array._new(Tp.take(flat.astype(np.int64, copy=False), axis=0), x._fill), len(S2) + + +# setitem +def setitem(x, key, value): + if not isinstance(x, Array): # pragma: no cover (safety) + raise TypeError(type(x)) + value = _setitem_value(x, value) + items = _normalize(x, key) + axes_key = [k for k in items if not isinstance(k, (_NewAxis, _Gap))] + if any(isinstance(k, _NewAxis) and k.length == 0 for k in items): + return # (a False index selects nothing) + if isinstance(value, Array) and not _has_adv(items): + value = _without_new_axes(value, items) + if value.ndim == 0: + value = _scalar_value(value._fill, x.dtype) + if x.ndim == 0: + if isinstance(value, Array): + raise ValueError(f"cannot assign an array of shape {value.shape} to a 0-d array") + x._fill = value + else: + adv = [k for k in axes_key if isinstance(k, _Adv)] + if ( + len(adv) == 1 + and adv[0].index.ndim == 1 + and not any(isinstance(k, int) for k in axes_key) + ): + # One 1-d index array among slices: outer (GraphBLAS) and vectorized indexing agree + _region_setitem( + x, tuple(k.index if isinstance(k, _Adv) else k for k in axes_key), value + ) + elif adv: + _scatter_setitem(x, [k for k in items if not isinstance(k, _NewAxis)], value) + else: + _region_setitem(x, tuple(axes_key), value) + x._facts = None + + +def _without_new_axes(value, items): + """``value`` for a basic key with new axes (``None``): the selection numpy broadcasts it + against has those axes, of length 1, which the region assigned to does not; drop the + value's axes that line up with them. + """ + kinds = [k for k in items if isinstance(k, (_NewAxis, slice))] + offset = len(kinds) - value.ndim + drop = [p - offset for p, k in enumerate(kinds) if isinstance(k, _NewAxis) and p - offset >= 0] + if not drop: + return value + if any(value.shape[q] != 1 for q in drop): + raise ValueError( + f"cannot assign an array of shape {value.shape} to a selection of " + f"{len(kinds)} dimensions whose new axes have length 1" + ) + from ._manipulation import squeeze + + return squeeze(value, axis=tuple(drop)) + + +def _setitem_value(x, value): + """The value cast to ``x``'s dtype: a numpy scalar, or an Array of at least 1 dimension.""" + from ._creation import asarray + from ._elementwise import cast + + dt = x.dtype + if isinstance(value, (builtins.bool, int, float, complex, np.generic)): + return _scalar_value(value, dt) + if not isinstance(value, Array): + value = asarray(value) + if value.ndim == 0: + return _scalar_value(value._fill, dt) + return value if value.dtype == dt else cast(value, dt) + + +def _scalar_value(value, dt): + if isinstance(value, (complex, np.complexfloating)) and dt.kind != "c": + raise TypeError(f"cannot assign a complex value to an array of {dt}") + with np.errstate(all="ignore"): + return np.array(value, dtype=dt)[()] + + +def _region_setitem(x, key, value): + """Assign to the region an outer index selects (GraphBLAS assign).""" + T = x._T() + region = tuple( + len(range(*k.indices(n))) if isinstance(k, slice) else len(k) + for k, n in zip(key, x.shape, strict=True) + if not isinstance(k, int) + ) + if not isinstance(value, Array): + if math.prod(region) == 0: + return + if _utils.same_value(value, x._fill): + if T.nvals: + del T[key] + else: + T[key] = _utils.gb_scalar(value, x.dtype) + return + try: + ok = np.broadcast_shapes(value.shape, region) == region + except ValueError: + ok = False + if not ok: + raise ValueError( + f"cannot assign an array of shape {value.shape} to a region of shape {region}" + ) + if math.prod(region) == 0: + return + # The region takes the value's structure: an unstored position of the value deletes + T[key] = _utils.with_fill(value, x._fill)._T() + + +def _scatter_setitem(x, axes_key, value): + """Assign by several index arrays: write every selected position.""" + # numpy lays out the selection; index broadcast views of the coordinates with the key (a + # ``...`` that covers no axis stays: it decides where the advanced axes go) + key = tuple( + k.index if isinstance(k, _Adv) else Ellipsis if isinstance(k, _Gap) else k for k in axes_key + ) + coords = [] + for a, n in enumerate(x.shape): + grid = np.arange(n, dtype=np.int64).reshape([n if b == a else 1 for b in range(x.ndim)]) + coords.append(np.asarray(np.broadcast_to(grid, x.shape)[key])) + shape = coords[0].shape + size = math.prod(shape) + if size == 0: + return + _utils.check_dense(size, "assigning to the selected positions") + if isinstance(value, Array): + try: + values = np.broadcast_to(value.todense(), shape) + except ValueError: + raise ValueError( + f"cannot assign an array of shape {value.shape} to a selection of shape {shape}" + ) from None + else: + values = np.broadcast_to(np.asarray(value, dtype=x.dtype), shape) + coords = [c.reshape(-1) for c in coords] + values = values.reshape(-1) + # The last of repeated positions wins, as in numpy + flat = np.ravel_multi_index(coords, x.shape) + _, last = np.unique(flat[::-1], return_index=True) + keep = np.sort(flat.size - 1 - last) + coords = [c[keep] for c in coords] + values = values[keep] + stored = _utils.differs(values, x._fill) + T = x._T() + gbdt = gb_dtype(x.dtype) + if stored.any(): + U = Tensor.from_coo([c[stored] for c in coords], values[stored], gbdt, shape=x.shape) + T(U.S) << U + if not stored.all() and T.nvals: + D = Tensor.from_coo([c[~stored] for c in coords], True, shape=x.shape) + T(D.S) << Tensor(gbdt, x.shape) + + +# take, take_along_axis +def _index_operand(indices, name): + if isinstance(indices, Array): + if indices.dtype.kind not in "iu": + raise TypeError(f"{name} needs integer indices, not {indices.dtype}") + return indices.todense().astype(np.int64, copy=False) + arr = np.asarray(indices) + if arr.dtype.kind not in "iu" and arr.size: + raise TypeError(f"{name} needs integer indices, not {arr.dtype}") + return arr.astype(np.int64, copy=False) + + +def _check_bounds(idx, n, axis): + if idx.size and (idx.min() < -n or idx.max() >= n): + raise IndexError(f"index out of bounds for axis {axis} with size {n}") + + +def take(x, indices, /, *, axis=None): + """``x``'s elements at ``indices`` along ``axis`` (the flattened array if None).""" + from ._manipulation import _normalize_axis + + if not isinstance(x, Array): + raise TypeError(f"take needs an Array, not {type(x).__name__}") + idx = _index_operand(indices, "take") + if x.ndim == 0: + raise ValueError("take needs an array of at least one dimension") + if axis is None: + if x.ndim > 1: + raise ValueError("take needs an axis for an array of more than one dimension") + axis = 0 + axis = _normalize_axis(axis, x.ndim) + _check_bounds(idx, x.shape[axis], axis) + T = x._T() + shape = x.shape[:axis] + idx.shape + x.shape[axis + 1 :] + if T.nvals == 0 or math.prod(shape) == 0: + return Array._new(Tensor(T.dtype, shape), x._fill) + return Array._new(T.take(idx, axis=axis), x._fill) + + +def take_along_axis(x, indices, /, *, axis=-1): + from ._manipulation import _normalize_axis + + if not isinstance(x, Array): + raise TypeError(f"take_along_axis needs an Array, not {type(x).__name__}") + idx = _index_operand(indices, "take_along_axis") + if idx.ndim != x.ndim: + raise ValueError("take_along_axis needs indices with as many dimensions as x") + axis = _normalize_axis(axis, x.ndim) + _check_bounds(idx, x.shape[axis], axis) + others = [s for a, s in enumerate(x.shape) if a != axis] + iothers = [s for a, s in enumerate(idx.shape) if a != axis] + bshape = list(np.broadcast_shapes(tuple(others), tuple(iothers))) + bshape.insert(axis, idx.shape[axis]) + shape = tuple(bshape) + T = x._T() + if T.nvals == 0 or math.prod(shape) == 0: + return Array._new(Tensor(T.dtype, shape), x._fill) + return Array._new(T.take_along_axis(idx, axis=axis), x._fill) diff --git a/graphblas/array_api/_info.py b/graphblas/array_api/_info.py new file mode 100644 index 000000000..d155d409f --- /dev/null +++ b/graphblas/array_api/_info.py @@ -0,0 +1,53 @@ +"""``__array_namespace_info__``: what the namespace supports.""" + +from ._array import CPU_DEVICE +from ._dtypes import _KINDS, DTYPES + +__all__ = ["__array_namespace_info__"] + + +class Info: + __module__ = "graphblas.array_api" + + def capabilities(self): + return { + "boolean indexing": True, + "data-dependent shapes": True, + "max dimensions": None, + } + + def default_device(self): + return CPU_DEVICE + + def devices(self): + return (CPU_DEVICE,) + + def default_dtypes(self, *, device=None): + _check(device) + return { + "real floating": DTYPES["float64"], + "complex floating": DTYPES["complex128"], + "integral": DTYPES["int64"], + "indexing": DTYPES["int64"], + } + + def dtypes(self, *, device=None, kind=None): + _check(device) + if kind is None: + return dict(DTYPES) + kinds = kind if isinstance(kind, tuple) else (kind,) + selected = set() + for k in kinds: + if k not in _KINDS: + raise ValueError(f"unknown data type kind {k!r}") + selected |= _KINDS[k] + return {name: dt for name, dt in DTYPES.items() if dt in selected} + + +def _check(device): + if device is not None and device != CPU_DEVICE: + raise ValueError(f"unknown device {device!r}") + + +def __array_namespace_info__(): + return Info() diff --git a/graphblas/array_api/_lazy.py b/graphblas/array_api/_lazy.py new file mode 100644 index 000000000..a9c5afd41 --- /dev/null +++ b/graphblas/array_api/_lazy.py @@ -0,0 +1,265 @@ +"""Deferred contractions (DESIGN D9, scoped to one level of deferral). + +matmul, tensordot, vecdot and einsum return an Array whose Tensor is not computed yet when +they skip the operands' unstored positions (their fills are the semiring's identity) and the +result has a dimension. Such a pending Array holds copy-on-write snapshots (``Tensor.dup``, +O(1)) of its operands' Tensors, so changing an operand afterwards does not change it (the +operand copies its data then, and the snapshots keep the old data alive until the result is +computed), and computes its Tensor on the first read (``Array._T``), once. + +Two consumers use a pending product P without computing it whole: + +- ``multiply(S, P)`` or ``multiply(P, S)``, with S of P's shape and both fills 0, computes P + only at S's stored positions (a masked product, then the multiply: an SDDMM). Elsewhere S + is 0, and 0 times P's value there is 0 when that value is finite. That holds when the + operands are finite and no sum of products can overflow: the number of terms of each sum + times the largest magnitude of each operand stays below the dtype's largest value (one + GraphBLAS pass per operand for each). Otherwise P is computed whole (a sparse product, as + without deferral), then multiplied. +- ``P[key]`` with a basic index (ints and slices) slices the operands along the indexed axes + of the result, then contracts (pending again). + +Everything else reads the Tensor, which computes it. Values do not depend on whether a result +was deferred; ``enabled = False`` turns deferral off, for debugging. +""" + +import math + +import numpy as np + +from .. import binary, monoid, select, semiring, unary +from ..core.tensor import Tensor +from ..dtypes import FC64, FP64, lookup_dtype +from . import _utils +from ._array import Array +from ._utils import gb_dtype, np_dtype + +# Private kill switch: False computes every contraction when it is called +enabled = True + + +class Contraction: + """A contraction not computed yet. + + ``build(*tensors)`` gives its TensorExpression; ``tensors`` are snapshots of the operands' + Tensors. ``spec``, ``(terms, output, semiring)`` of an einsum that computes it, lets a + slice of the result slice the operands instead (None: no such einsum). ``finite`` holds + each operand's finiteness: True, False, or None when not checked yet. ``nterms``: the + number of terms each result position sums (the product of the contracted sizes). + ``product``: whether it is one product of its two operands, which takes the mask. + """ + + __slots__ = "tensors", "build", "gbdt", "spec", "finite", "nterms", "product" + + def __init__(self, tensors, build, gbdt, spec, finite, nterms, product=False): + self.tensors = tensors + self.build = build + self.gbdt = gbdt + self.spec = spec + self.finite = finite + self.nterms = nterms + self.product = product + + def compute(self, mask=None, **opts): + return self.build(*self.tensors).new(dtype=self.gbdt, mask=mask, **opts) + + def all_finite(self): + """Whether every value of the product is finite: every operand value is finite, and + ``nterms`` times the largest magnitude of each operand is below half the largest + value of the result's dtype (the half covers the rounding of the sums). + """ + for i, ok in enumerate(self.finite): + if ok is None: + ok = self.finite[i] = _all_finite(self.tensors[i]) + if not ok: + return False + dt = np_dtype(self.gbdt) + if dt.kind not in "fc": + return True # (integers wrap, as numpy's do: 0 times a wrapped sum is 0) + bound = float(self.nterms) + for T in self.tensors: + bound *= _max_abs(T) # (an overflow here is inf, and fails the test) + return bound < float(np.finfo(dt).max) / 2 + + +def _all_finite(T): + return np_dtype(T.dtype).kind not in "fc" or _utils.tensor_fact(T, "finite") + + +def _max_abs(T): + """The largest magnitude of the stored values of ``T``, as a float (NaN skipped).""" + if T.nvals == 0: + return 0.0 + if np_dtype(T.dtype).kind == "c": + A = T.apply(unary.abs[FC64]).new() + return float(A.reduce_scalar(monoid.max[FP64]).new().value) + # (two reductions in place: a copy of |values| cost 40 ms of 155 at n = 1e6, k = 64) + hi = T.reduce_scalar(monoid.max[FP64]).new().value + lo = T.reduce_scalar(monoid.min[FP64]).new().value + return max(abs(hi), abs(lo)) + + +def _known_finite(x): + """Whether ``x``'s values are finite, if ``x`` already knows (from its cached facts); + else None. + """ + if x.dtype.kind not in "fc": + return True + facts = x._facts + return None if facts is None else facts.get("finite") + + +def contraction(arrays, build, shape, dtype, fill, spec=None, nterms=1, product=False, typed=None): + """The Array of ``build(*tensors)`` over the arrays' Tensors (``build(*tensors, op=)`` + with another semiring), with fill ``fill``: pending when deferral is on and the result + has a dimension, else computed now. ``nterms``: the number of terms each result position + sums; ``product``: see ``Contraction``; ``typed``: the semiring. + + When an operand holds inf or NaN, the operands' stored entries equal to their fill go + first, so that such a value meets an explicit fill as it meets an unstored one, with no + term (DESIGN D12a). When a float min or max monoid would drop a stored NaN term, the + result is computed now and NaN set where such a term is (D12b). + """ + if any(x.dtype.kind in "fc" and not _utils.fact(x, "finite") for x in arrays): + from ._manipulation import prune + + arrays = [prune(x) for x in arrays] + tensors = [x._T() for x in arrays] + gbdt = gb_dtype(dtype) + nan = ( + typed is not None + and typed.monoid.name in {"min", "max"} + and np_dtype(typed.return_type).kind == "f" + and any(x.dtype.kind == "f" and _utils.fact(x, "nan") for x in arrays) + ) + if nan: + return Array._new(_with_nan_terms(build, tensors, gbdt), fill) + if not enabled or not shape: + return Array._new(build(*tensors).new(dtype=gbdt), fill) + snapshots = [T.dup() for T in tensors] + finite = [_known_finite(x) for x in arrays] + node = Contraction(snapshots, build, gbdt, spec, finite, nterms, product) + return Array._deferred(node, shape, dtype, fill) + + +def _with_nan_terms(build, tensors, gbdt): + """``build(*tensors)`` with NaN wherever a term holds a stored NaN: SuiteSparse's float + min and max are fmin and fmax, which skip NaN, where numpy's propagate it. Those positions + are the product of one operand's NaN entries with the others' patterns (any_pair). + """ + R = build(*tensors).new(dtype=gbdt) + bool_ = lookup_dtype(bool) + patterns = [T.apply(unary.one[bool_]).new() for T in tensors] + N = None + for i, T in enumerate(tensors): + # (every operand is float here: its fill is the identity of min or max, an infinity) + nan = T.apply(unary.isnan).new().select(select.valueeq, True).new() + if nan.nvals == 0: + continue + ops = [nan if j == i else P for j, P in enumerate(patterns)] + Ni = build(*ops, op=semiring.any_pair[bool_]).new() + N = Ni if N is None else N.ewise_add(Ni, binary.lor).new() + if N is not None and N.nvals: + R(N.S) << _utils.gb_scalar(np.nan, np_dtype(gbdt)) + return R + + +def masked_multiply(a, b): + """``multiply(a, b)`` when one of them is a pending product P and the other, S, lets P be + computed at S's stored positions only (see the module's docstring); else None. + """ + if not isinstance(a, Array) or not isinstance(b, Array): + return None + if b._pending is not None: + S, P = a, b + elif a._pending is not None: + S, P = b, a + else: + return None + if S is P or S.shape != P.shape or bool(S._fill) or bool(P._fill): + return None # (P * P: computed once, as a whole) + node = P._pending + if not node.all_finite(): + return None # (P is computed whole, a sparse product as without deferral) + TS = S._T() + hint = _dot_hint(node, TS.nvals, P.shape) + PS = Array._new(node.compute(mask=TS.S, **hint), P._fill) + from ._elementwise import multiply + + return multiply(PS, S) if P is a else multiply(S, PS) + + +def _dot_hint(node, nmask, shape=()): + """``axb_method="dot"`` for a product under a mask of ``nmask`` entries when dot products + at the mask's entries cost well below a saxpy over the whole product, else nothing. + + SuiteSparse takes a masked saxpy for A @ B with B held by row, all of A @ B's work: an + SDDMM with full A (n x 32) and B (32 x n), n = 4000 and 4e4 mask entries, took 0.6 s, + against 0.6 ms as dot products (SuiteSparse transposes B for them). + """ + if not node.product: + return {} + Ta, Tb = node.tensors + k = max(node.nterms, 1) + # (a batched product pairs entries within each batch only) + if not shape: + nbatch = 1 + elif node.spec is not None: + (ta, tb), output = node.spec[:2] + nbatch = math.prod(shape[output.index(c)] for c in output if c in ta and c in tb) + else: + nbatch = math.prod(shape[:-2]) # (a matmul whose batch axes broadcast) + saxpy = Ta.nvals * Tb.nvals / (k * max(nbatch, 1)) + dot = nmask * k * (Ta.nvals / max(Ta.size, 1) + Tb.nvals / max(Tb.size, 1)) + return {"axb_method": "dot"} if 4 * dot < saxpy else {} + + +def sliced(x, axes_key): + """``x[axes_key]`` (an int or a slice per axis) of a pending ``x`` that an einsum + computes: its operands sliced along the indexed axes, contracted (pending again). None + when that does not apply. + """ + node = x._pending + if node is None or node.spec is None or all(k == slice(None) for k in axes_key): + return None + terms, output, typed = node.spec + dropped = {c for c, k in zip(output, axes_key, strict=True) if not isinstance(k, slice)} + new_terms = ["".join(c for c in term if c not in dropped) for term in terms] + out = "".join(c for c in output if c not in dropped) + if not out or not all(new_terms): + return None # (a 0-d result or operand: computed whole, then indexed) + shape = tuple( + len(range(*k.indices(n))) + for k, n in zip(axes_key, x.shape, strict=True) + if isinstance(k, slice) + ) + keys = dict(zip(output, axes_key, strict=True)) + tensors = [] + for term, T in zip(terms, node.tensors, strict=True): + key = tuple(keys.get(c, slice(None)) for c in term) + if any(k != slice(None) for k in key): + T = T[key].new() + tensors.append(T) + if 0 in shape or any(T.nvals == 0 for T in tensors): + return Array._new(Tensor(node.gbdt, shape), x._fill) + from ..core.tensor_einsum import einsum + + subscripts = ",".join(new_terms) + "->" + out + + def build(*tensors, op=typed): + return einsum(subscripts, *tensors, op=op) + + # (a slice of a finite operand is finite) + spec = (tuple(new_terms), out, typed) + product = _one_product(new_terms, out) + sub = Contraction(tensors, build, node.gbdt, spec, list(node.finite), node.nterms, product) + return Array._deferred(sub, shape, x.dtype, x._fill) + + +def _one_product(terms, output): + """Whether an einsum of ``terms`` into ``output`` is one product of two operands, whose + result needs no permutation or sum after it (so a mask goes into the product). + """ + from ..core.tensor_einsum import _direct_axes + + return len(terms) == 2 and _direct_axes(*terms, output) is not None diff --git a/graphblas/array_api/_linalg.py b/graphblas/array_api/_linalg.py new file mode 100644 index 000000000..8e5adfa00 --- /dev/null +++ b/graphblas/array_api/_linalg.py @@ -0,0 +1,614 @@ +"""Linear algebra in the main namespace: matmul, matrix_transpose, tensordot and vecdot, and +the extension einsum. + +Contractions treat a fill of zero as a structural zero (DESIGN D6), as scipy.sparse and +pydata/sparse do: a product with an unstored position is skipped, so 0 * inf there is not NaN, +unlike numpy. Sums may round differently from numpy's (the order of the terms differs). An +operand that stores every position never shows its fill, so a nonzero fill there is taken as +0; another nonzero fill makes the contraction go through numpy densely, under the densify +limit on its operands and its result (DESIGN D12c). + +Semirings (an extension, DESIGN D11): ``semiring=`` takes a ``graphblas`` Semiring or its +name ("min_plus", "max_min", ...), whose monoid takes the place of the sum and whose binary +operator that of the product. A contraction skips the unstored positions, which is right when +every operand's fill is the identity ``e`` of the semiring's monoid and ``e`` annihilates the +multiply from both sides, ``mult(e, y) == mult(y, e) == e``: inf for min_plus, -inf for +max_plus, False for lor_land, 0 for plus_times. The result's fill is then ``e``; other fills +raise ValueError. As in D6 the skipping is structural: an infinity meeting an unstored +position makes no NaN (inf + -inf in min_plus). The annihilator is checked on a set of finite +probe values of each operand type, once per typed semiring. The default (None) is plus_times, +or lor_land for bool, under the rules of the first paragraph. + +A contraction that skips the unstored positions is deferred: its result computes on first use, +or only where a consumer needs it (``multiply`` by a zero-filled array, a slice; see +``_lazy``). +""" + +import math +import weakref + +import numpy as np + +from .. import semiring as semirings +from ..core.operator.semiring import Semiring, TypedBuiltinSemiring, TypedUserSemiring +from ..core.scalar import Scalar +from ..core.tensor import Tensor +from ..core.vector import Vector +from . import _lazy, _utils +from ._array import Array +from ._dtypes import check_supported +from ._utils import gb_dtype, np_dtype + +__all__ = ["compute", "einsum", "matmul", "matrix_transpose", "tensordot", "vecdot"] + + +def _check(x, name): + if not isinstance(x, Array): + raise TypeError(f"{name} needs Arrays, not {type(x).__name__}") + + +def _operand(x): + """A numpy array as an Array (``ndarray @ Array`` comes here through ``__rmatmul__``).""" + if isinstance(x, np.ndarray): + from ._creation import asarray + + return asarray(x) + return x + + +def compute(x, /): + """``x``, computed now (an extension). matmul, tensordot, vecdot and einsum return an Array + that computes its value on first use, or only where a consumer needs it (DESIGN D9); + ``compute`` computes it at once (to time it, or before handing it on). Any other Array + is returned as it is. + """ + _check(x, "compute") + x._T() + return x + + +def _semiring(dt): + return (semirings.lor_land if dt.kind == "b" else semirings.plus_times)[gb_dtype(dt)] + + +def _resolve(semiring, dt, name): + """``(typed semiring, whether it is the default)`` for the ``semiring=`` argument, typed + for ``dt``, the operands' promoted dtype. + """ + default = _semiring(dt) + if semiring is None: + return default, True + if isinstance(semiring, str): + try: + sr = semirings.from_string(semiring) + except ValueError: + raise ValueError(f"{name}: unknown semiring {semiring!r}") from None + elif isinstance(semiring, (Semiring, TypedBuiltinSemiring, TypedUserSemiring)): + sr = semiring + else: + raise TypeError( + f"{name}'s semiring must be a graphblas Semiring or its name, not " + f"{type(semiring).__name__}" + ) + if isinstance(sr, Semiring): + if dt.kind == "b" and sr.name == "plus_times": + return default, True # (numpy's sum of bools is their lor) + try: + sr = sr[gb_dtype(dt)] + except KeyError: + raise TypeError(f"{name}: the semiring {sr.name} does not work with {dt}") from None + same = sr is default or ( + sr.name == default.name and sr.type == default.type and sr.type2 == default.type2 + ) + return sr, same + + +def _probes(dt): + """Values of the numpy dtype ``dt`` to check an annihilator on: both signs, zeros and the + extremes, all finite (an infinity meets an unstored position only structurally, as in D6). + """ + if dt.kind == "b": + return np.array([False, True]) + if dt.kind in "iu": + info = np.iinfo(dt) + values = [0, 1, 2, 7, info.max, info.max // 3] + if dt.kind == "i": + values += [-1, -5, info.min, info.min // 3] + return np.array(values, dtype=dt) + big = np.finfo(dt).max + values = [0.0, -0.0, 1.0, -1.0, 0.5, -2.5, 3e9, big, -big, np.finfo(dt).tiny] + if dt.kind == "c": + values += [1j, -1j, 1 + 2j, -3.5 + 0.5j, complex(0, big)] + return np.array(values, dtype=dt) + + +_ANNIHILATES = weakref.WeakKeyDictionary() + + +def _annihilates(typed): + """Whether the identity of the typed semiring's monoid annihilates its multiply from both + sides on the probe values (cached per typed semiring object: two user semirings may share + a name). + """ + key = typed + if (ok := _ANNIHILATES.get(key)) is None: + ok = _ANNIHILATES[key] = _check_annihilates(typed) + return ok + + +def _check_annihilates(typed): + if typed.is_positional: + return False + mult = typed.binaryop + e = typed.monoid.identity + for left in (True, False): + e_type, y_type = (typed.type, typed.type2) if left else (typed.type2, typed.type) + ys = _probes(np_dtype(y_type)) + v = Vector.from_coo(np.arange(ys.size), ys, dtype=y_type, size=ys.size) + s = Scalar.from_value(_utils.scalar(e, np_dtype(e_type)), dtype=e_type) + r = (v.apply(mult, left=s) if left else v.apply(mult, right=s)).new() + if r.nvals != ys.size or not bool(np.all(r.to_coo()[1] == e)): + return False + return True + + +def _operand_name(name, i): + return f"operand {i}" if name == "einsum" else f"x{i + 1}" + + +def _spell(value): + """``value`` as the namespace spells it (``xp.inf`` for inf).""" + value = np.asarray(value).item() + if isinstance(value, float) and np.isinf(value): + return "xp.inf" if value > 0 else "-xp.inf" + return repr(value) + + +def _sound(typed, arrays, name, *, multiplies=True): + """The result dtype and fill of a contraction of ``arrays`` with a typed semiring other + than the default, and the arrays, after checking that it may skip their unstored + positions (ValueError if not). An array that stores every position never shows its fill, + so it is relabeled with the identity when its dtype holds it (O(1)). ``multiplies``: + whether values meet the multiply (einsum of one operand only sums). + """ + e = typed.monoid.identity + out_dt = check_supported(np_dtype(typed.return_type), name) + if multiplies and not _annihilates(typed): + mult = typed.binaryop.name + raise ValueError( + f"{name}: the semiring {typed.name} on {np_dtype(typed.type)} cannot skip unstored " + f"positions: the identity {e!r} of its monoid {typed.monoid.name} does not " + f"annihilate {mult} ({mult}(e, y) and {mult}(y, e) must be e for every y), so no " + "fill works. Use it on Tensors (x.tensor), which skip missing entries." + ) + out = [] + for i, x in enumerate(arrays): + in_dt = np_dtype(typed.type if i == 0 else typed.type2) + fill = _utils.scalar(x._fill, in_dt) + if fill != e and (relabeled := _relabel_full(x, e)) is not None: + x, fill = relabeled, e + out.append(x) + if fill != e: # (a NaN fill too) + spelled = _spell(e) + raise ValueError( + f"{name} with the semiring {typed.name} needs operands whose fill is {e!r}, " + f"the identity of {typed.monoid.name}" + + (f", which annihilates {typed.binaryop.name}" if multiplies else "") + + f"; {_operand_name(name, i)} has the fill {x._fill.item()!r}. Make the array " + f"with fill_value={spelled}, or use xp.refill(x, {spelled}), which keeps its " + "values (it stores the positions that held the old fill)." + ) + return out_dt, _utils.scalar(e, out_dt), out + + +def _relabel_full(x, e): + """``x`` with fill ``e`` when it stores every position (its fill is never seen) and its + dtype holds ``e`` exactly; else None. + """ + if x._T().nvals != x.size: + return None + with np.errstate(all="ignore"): + fill = np.asarray(e).astype(x.dtype)[()] # (inf as an integer: no longer inf) + if not _utils.same_value(np.asarray(fill).astype(np.asarray(e).dtype)[()], e): + return None + return Array._new(x._T(), fill) + + +def _is_zero(fill): + return not bool(fill) + + +def _zero_fills(*arrays): + """The arrays with fill 0, or None when one cannot have it. An array that stores every + position never shows its fill, so its fill is relabeled 0 (O(1)). One with a finite fill + whose size fits under the densify limit (a vector or matrix of ones by a sparse matrix) is + stored in full first, O(its size): an unstored zero of another operand times its finite + values is 0 either way, so the answer does not depend on what the others store. Any other + nonzero fill gives None: the contraction is dense (a non-finite fill times another's + unstored zero is NaN there, as numpy has it). + """ + out = [] + for x in arrays: + if not _is_zero(x._fill): + T = x._T() + if T.nvals != x.size: + small = x.size <= _utils.densify_limit() + if x.ndim == 0 or not small or not np.isfinite(x._fill): + return None + # (its fill stored at every position it does not store, the new fill 0) + T = _utils.with_fill(x, _utils.scalar(0, x.dtype))._T() + x = Array._new(T, _utils.scalar(0, x.dtype)) + out.append(x) + return out + + +def _dense(func, shape, *arrays): + """``func`` of the arrays densely by numpy, under the densify limit on the operands and on + the result of ``shape``. + """ + from ._creation import _from_numpy + + what = "a contraction with a nonzero fill that does not store every position" + for x in arrays: + _utils.check_dense(x.size, what) + _utils.check_dense(math.prod(shape), what) + out = np.asarray(func(*[x.todense() for x in arrays])) + return _from_numpy(out) + + +def _rdt(name, func, x1, x2): + probe = func( + np.ones((1,) * max(x1.ndim, 1), x1.dtype), np.ones((1,) * max(x2.ndim, 1), x2.dtype) + ) + return check_supported(np.asarray(probe).dtype, name) + + +def matmul(x1, x2, /, *, semiring=None): + """The matrix product (numpy's matmul rules). ``semiring`` (an extension): a graphblas + Semiring or its name, under the module's rules for fills. + """ + x1, x2 = _operand(x1), _operand(x2) + _check(x1, "matmul") + _check(x2, "matmul") + if x1.ndim == 0 or x2.ndim == 0: + raise ValueError("matmul needs arrays of at least one dimension") + K = x1.shape[-1] + if x2.shape[-2 if x2.ndim > 1 else 0] != K: + raise ValueError(f"matmul: shapes {x1.shape} and {x2.shape} do not align") + b1 = x1.shape[:-2] + b2 = x2.shape[:-2] + try: + np.broadcast_shapes(b1, b2) + except ValueError: + raise ValueError(f"matmul: shapes {x1.shape} and {x2.shape} do not broadcast") from None + rdt = _rdt("matmul", np.matmul, x1, x2) + typed, default = _resolve(semiring, rdt, "matmul") + if not default: + out_dt, fill, (x1, x2) = _sound(typed, (x1, x2), "matmul") + return _product(x1, x2, typed, out_dt, fill) + if (ops := _zero_fills(x1, x2)) is None: + return _dense(np.matmul, _matmul_shape(x1, x2), x1, x2) + return _product(*ops, typed, rdt, _utils.scalar(0, rdt)) + + +def _matmul_shape(a, b): + shape = np.broadcast_shapes(a.shape[:-2], b.shape[:-2]) + return shape + (a.shape[-2:-1] if a.ndim > 1 else ()) + (b.shape[-1:] if b.ndim > 1 else ()) + + +def _product(a, b, typed, out_dt, fill): + """``a @ b`` with the typed semiring, skipping unstored positions (deferred, see + ``_lazy``); the result's fill is ``fill``. + """ + Ta, Tb = a._T(), b._T() + shape = _matmul_shape(a, b) + if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: + return Array._new(Tensor(gb_dtype(out_dt), shape), fill) + + def build(Ta, Tb, op=typed): + return Ta.matmul(Tb, op) + + spec = _matmul_terms(a.shape, b.shape) + spec = None if spec is None else (*spec, typed) + return _lazy.contraction((a, b), build, shape, out_dt, fill, spec, a.shape[-1], True, typed) + + +# Index letters for the einsum that computes a contraction (for slicing a pending result) +_LETTERS = "abcdefghlmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" + + +def _matmul_terms(sa, sb): + """``(terms, output)`` of the einsum equal to matmul of shapes ``sa`` and ``sb``, or None + when a batch axis of size 1 broadcasts (einsum has no broadcasting). + """ + ba, bb = sa[:-2], sb[:-2] + nb = max(len(ba), len(bb)) + if nb > len(_LETTERS): + return None # pragma: no cover (safety) + batch = _LETTERS[:nb] + for b in (ba, bb): + for n, m in zip(b[::-1], np.broadcast_shapes(ba, bb)[::-1], strict=False): + if n != m: + return None + ta = batch[nb - len(ba) :] + ("ik" if len(sa) > 1 else "k") + tb = batch[nb - len(bb) :] + ("kj" if len(sb) > 1 else "k") + out = batch + ("i" if len(sa) > 1 else "") + ("j" if len(sb) > 1 else "") + return (ta, tb), out + + +def _tensordot_terms(na, nb, a_axes, b_axes): + """``(terms, output)`` of the einsum equal to tensordot of arrays of ``na`` and ``nb`` + axes contracting ``a_axes`` with ``b_axes``. + """ + if na + nb > len(_LETTERS): + return None # pragma: no cover (safety) + ta = _LETTERS[:na] + tb = list(_LETTERS[na : na + nb]) + for i, j in zip(a_axes, b_axes, strict=True): + tb[j] = ta[i] + tb = "".join(tb) + out = "".join(c for c in ta if c not in tb) + "".join(c for c in tb if c not in ta) + return (ta, tb), out + + +def matrix_transpose(x, /): + _check(x, "matrix_transpose") + return x.mT + + +def _tensordot_axes(x1, x2, axes): + if isinstance(axes, int): + if axes < 0: + raise ValueError("tensordot's axes must be nonnegative") + if axes > min(x1.ndim, x2.ndim): + raise ValueError("tensordot: more axes than the arrays have") + return tuple(range(x1.ndim - axes, x1.ndim)), tuple(range(axes)) + a, b = axes + a = tuple(_axis(i, x1.ndim) for i in (a if isinstance(a, (list, tuple)) else (a,))) + b = tuple(_axis(i, x2.ndim) for i in (b if isinstance(b, (list, tuple)) else (b,))) + if len(a) != len(b): + raise ValueError("tensordot needs as many axes of each array") + return a, b + + +def _axis(i, ndim): + i = int(i) + if not -ndim <= i < ndim: + raise IndexError(f"tensordot: axis {i} is out of bounds for an array of {ndim} dimensions") + return i % ndim + + +def tensordot(x1, x2, /, *, axes=2, semiring=None): + """The tensor dot product (numpy's tensordot). ``semiring`` (an extension): a graphblas + Semiring or its name, under the module's rules for fills. + """ + _check(x1, "tensordot") + _check(x2, "tensordot") + a_axes, b_axes = _tensordot_axes(x1, x2, axes) + for a, b in zip(a_axes, b_axes, strict=True): + if x1.shape[a] != x2.shape[b]: + raise ValueError(f"tensordot: axis {a} of x1 and axis {b} of x2 differ in size") + rdt = check_supported(np.result_type(x1.dtype, x2.dtype), "tensordot") + typed, default = _resolve(semiring, rdt, "tensordot") + free1 = tuple(n for i, n in enumerate(x1.shape) if i not in a_axes) + free2 = tuple(n for i, n in enumerate(x2.shape) if i not in b_axes) + shape = free1 + free2 + if not default: + out_dt, fill, (x1, x2) = _sound(typed, (x1, x2), "tensordot") + return _tensordot(x1, x2, a_axes, b_axes, shape, typed, out_dt, fill) + if (ops := _zero_fills(x1, x2)) is None: + func = lambda a, b: np.tensordot(a, b, axes=(a_axes, b_axes)) # noqa: E731 + return _dense(func, shape, x1, x2) + return _tensordot(*ops, a_axes, b_axes, shape, typed, rdt, _utils.scalar(0, rdt)) + + +def _tensordot(a, b, a_axes, b_axes, shape, typed, out_dt, fill): + Ta, Tb = a._T(), b._T() + if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: + return Array._new(Tensor(gb_dtype(out_dt), shape), fill) + axes = (list(a_axes), list(b_axes)) + + def build(Ta, Tb, op=typed): + return Ta.tensordot(Tb, axes, op) + + spec = _tensordot_terms(a.ndim, b.ndim, a_axes, b_axes) + spec = None if spec is None else (*spec, typed) + nterms = math.prod(a.shape[i] for i in a_axes) + return _lazy.contraction((a, b), build, shape, out_dt, fill, spec, nterms, True, typed) + + +def vecdot(x1, x2, /, *, axis=-1, semiring=None): + """Sum over ``axis`` of ``conj(x1) * x2`` (broadcast), with zero fills structural. + + ``semiring`` (an extension): a graphblas Semiring or its name, whose monoid reduces + ``mult(x1, x2)`` over ``axis``, under the module's rules for fills. The default semiring + (plus_times, or lor_land for bool), given or not, conjugates ``x1``; others do not. + """ + from ._elementwise import conj + + _check(x1, "vecdot") + _check(x2, "vecdot") + ndim = min(x1.ndim, x2.ndim) + if ndim == 0: + raise ValueError("vecdot needs arrays of at least one dimension") + if not -ndim <= axis < 0: + if 0 <= axis < ndim and x1.ndim == x2.ndim: + axis -= x1.ndim + else: + raise ValueError(f"vecdot's axis must be in [-{ndim}, -1], not {axis}") + if x1.shape[axis] != x2.shape[axis]: + raise ValueError("vecdot: the axis differs in size") + shape = np.broadcast_shapes(x1.shape, x2.shape) + rdt = check_supported(np.result_type(x1.dtype, x2.dtype), "vecdot") + typed, default = _resolve(semiring, rdt, "vecdot") + out_shape = tuple(n for i, n in enumerate(shape) if i != len(shape) + axis) + if not default: + out_dt, fill, (x1, x2) = _sound(typed, (x1, x2), "vecdot") + return _vecdot(x1, x2, axis, shape, out_shape, typed, out_dt, fill) + if (ops := _zero_fills(x1, x2)) is None: + return _dense(lambda a, b: np.vecdot(a, b, axis=axis), out_shape, x1, x2) + a0, b0 = ops + if a0.dtype.kind == "c": + a0 = conj(a0) + return _vecdot(a0, b0, axis, shape, out_shape, typed, rdt, _utils.scalar(0, rdt)) + + +def _vecdot(a, b, axis, shape, out_shape, typed, out_dt, fill): + """The semiring's monoid over ``axis`` of ``mult(a, b)`` (broadcast to ``shape``).""" + Ta, Tb = a._T(), b._T() + if Ta.nvals == 0 or Tb.nvals == 0 or math.prod(shape) == 0: + return Array._new(Tensor(gb_dtype(out_dt), out_shape), fill) + if _both_broadcast(Ta.shape, Tb.shape, shape): + # (an outer product: ewise_mult replicates one operand) + _utils.expansion_check(Ta, Tb, shape) + axis += len(shape) + + def build(Ta, Tb, op=typed): + return Ta.ewise_mult(Tb, op.binaryop).new().reduce(op.monoid, axis=axis) + + return _lazy.contraction( + (a, b), build, out_shape, out_dt, fill, nterms=shape[axis], typed=typed + ) + + +def _both_broadcast(sa, sb, shape): + from ..core.tensor_ewise import _broadcast_factor + + return _broadcast_factor(sa, shape) > 1 and _broadcast_factor(sb, shape) > 1 + + +def einsum(subscripts, /, *operands, semiring=None): + """Einstein summation (an extension), with numpy's subscripts: ``einsum("ij,jk->ik", a, + b)`` is ``a @ b``. Each index appears at most once per operand (no diagonals), and + ``...`` is not supported. + + ``semiring``: a graphblas Semiring or its name; the default is plus_times, or lor_land + for bool. Its multiply combines the operands of each term, and its monoid sums the terms + over the indices not in the output. Fills follow the module's rules; with the default + semiring an operand whose fill is not zero makes the computation go through numpy + densely (under the densify limit). + """ + from ..core.tensor_einsum import _einsum_plan + from ..core.tensor_einsum import einsum as tensor_einsum + + if not isinstance(subscripts, str): + raise TypeError(f"einsum's subscripts must be a str, not {type(subscripts).__name__}") + if not operands: + raise ValueError("einsum needs at least one operand") + ops = [_operand(x) for x in operands] + for x in ops: + _check(x, "einsum") + rdt = check_supported(np.result_type(*(x.dtype for x in ops)), "einsum") + typed, default = _resolve(semiring, rdt, "einsum") + shapes = tuple(x.shape for x in ops) + shape = _einsum_plan(subscripts, shapes, "greedy")[2] + scalars = [x._fill for x in ops if x.ndim == 0] + if default and any(map(bool, scalars)) and all(np.isfinite(scalars)): + # A finite 0-d operand is a scalar factor of every term (and of the sum, for + # plus_times); a non-finite one goes dense, as numpy has it (inf times the unstored + # zeros is NaN) + return _einsum_scalars(subscripts, ops) + if default: + if (zeroed := _zero_fills(*ops)) is None: + return _dense(lambda *a: np.einsum(subscripts, *a), shape, *ops) + ops = zeroed + out_dt, fill = rdt, _utils.scalar(0, rdt) + else: + out_dt, fill, ops = _sound(typed, ops, "einsum", multiplies=len(ops) > 1) + if math.prod(shape) == 0 or any(x._T().nvals == 0 for x in ops): + return Array._new(Tensor(gb_dtype(out_dt), shape), fill) + + def build(*tensors, op=typed): + return tensor_einsum(subscripts, *tensors, op=op) + + terms, output = _einsum_plan(subscripts, shapes, "greedy")[:2] + if len(ops) == 3: + masked = _masked_three(terms, output, ops, typed, out_dt, fill) + if masked is not None: + return masked + sizes = { + c: n for term, s in zip(terms, shapes, strict=True) for c, n in zip(term, s, strict=True) + } + nterms = math.prod(n for c, n in sizes.items() if c not in output) + product = _lazy._one_product(terms, output) + return _lazy.contraction( + ops, build, shape, out_dt, fill, (terms, output, typed), nterms, product, typed + ) + + +# (monoid, multiply) pairs where the multiply distributes over the monoid, so that a term's +# third factor can multiply the sum of the other two's products: s * (a + b) = s*a + s*b +_DISTRIBUTES = { + ("plus", "times"), + ("lor", "land"), + ("min", "plus"), + ("max", "plus"), + ("min", "max"), + ("max", "min"), +} + + +def _masked_three(terms, output, ops, typed, out_dt, fill): + """An einsum of three operands in which one operand's indices are exactly those the + other two leave after contracting their shared ones (triangle counting's + "ij,jk,ik->", an SDDMM's "ik,ij,jk->ik"), computed as that pair's product at the third + operand's stored positions only, times the third, then summed into the output; else + None. Every term of the full einsum is a term of this one: a position the third operand + does not store gives no term either way (DESIGN D6). Sums differ only in their order + (rounding, and where a sum overflows). An operand holding inf or NaN prunes the explicit + fills first, as every contraction does (D12a); a NaN under a float min or max monoid + takes the full einsum, which propagates it (D12b). With plus_times the third operand + must be finite: inf * (x + y) is not inf * x + inf * y when x and y differ in sign. + """ + if (typed.monoid.name, typed.binaryop.name) not in _DISTRIBUTES or not all(x.ndim for x in ops): + return None + finite = [x.dtype.kind not in "fc" or _utils.fact(x, "finite") for x in ops] + if not all(finite): + if typed.monoid.name in {"min", "max"} and any(_utils.fact(x, "nan") for x in ops): + return None + if any(x.dtype.kind == "c" for x in ops): + return None # (complex infinities do not distribute) + from ._manipulation import prune + + ops = [prune(x) for x in ops] + for m, tm in enumerate(terms): + a, b = (t for i, t in enumerate(terms) if i != m) + contracted = (set(a) & set(b)) - set(tm) + # (the pair's other indices, batch ones included, are exactly the third's) + if not contracted or set(a) | set(b) != set(tm) | contracted or not set(output) <= set(tm): + continue + if typed.monoid.name == "plus" and not finite[m]: + continue + A, B = (x for i, x in enumerate(ops) if i != m) + P = einsum(f"{a},{b}->{tm}", A, B, semiring=typed) + TS = ops[m]._T() + if (node := P._pending) is not None: + TP = node.compute(mask=TS.S, **_lazy._dot_hint(node, TS.nvals, P.shape)) + else: + TP = P._T() # (computed already: a NaN or deferral off) + R = Array._new(TS.ewise_mult(TP, typed.binaryop).new(dtype=gb_dtype(out_dt)), fill) + return R if output == tm else einsum(f"{tm}->{output}", R, semiring=typed) + return None + + +def _einsum_scalars(subscripts, ops): + """Einsum with the 0-d operands taken out as a product of scalars: the einsum of the + others, times that product. + """ + from ._elementwise import multiply + + inputs, arrow, output = subscripts.replace(" ", "").partition("->") + terms = inputs.split(",") + keep = [i for i, x in enumerate(ops) if x.ndim > 0] + factor = None + for x in ops: + if x.ndim == 0: + factor = x if factor is None else multiply(factor, x) + if not keep: + return _dense(lambda *a: np.einsum(subscripts, *a), (), *ops) + rest = ",".join(terms[i] for i in keep) + arrow + output + if not arrow: + # (implicit output: the letters that appear once, in order, as numpy picks them) + counts = {c: "".join(terms).count(c) for c in "".join(terms)} + rest += "->" + "".join(sorted(c for c, k in counts.items() if k == 1)) + return multiply(einsum(rest, *(ops[i] for i in keep)), factor) diff --git a/graphblas/array_api/_manipulation.py b/graphblas/array_api/_manipulation.py new file mode 100644 index 000000000..db163c21c --- /dev/null +++ b/graphblas/array_api/_manipulation.py @@ -0,0 +1,394 @@ +"""Manipulation functions: shape changes share the stored data (copy-on-write).""" + +import math + +import numpy as np + +from .. import binary, select +from ..core.tensor import Tensor +from . import _utils +from ._array import Array +from ._utils import gb_dtype + +__all__ = [ + "broadcast_arrays", + "broadcast_shapes", + "broadcast_to", + "concat", + "expand_dims", + "flip", + "moveaxis", + "permute_dims", + "prune", + "refill", + "repeat", + "reshape", + "roll", + "squeeze", + "stack", + "tile", + "unstack", +] + + +def _check_array(x, name): + if not isinstance(x, Array): + raise TypeError(f"{name} needs an Array, not {type(x).__name__}") + + +def _normalize_axis(axis, ndim): + axis = int(axis) + if not -ndim <= axis < ndim: + raise IndexError(f"axis {axis} is out of bounds for {ndim} dimensions") + return axis % ndim + + +def _new_shape(shape, size): + if isinstance(shape, (int, np.integer)): + shape = (shape,) + shape = tuple(int(n) for n in shape) + unknown = [i for i, n in enumerate(shape) if n == -1] + if len(unknown) > 1: + raise ValueError("reshape takes at most one -1") + if any(n < -1 for n in shape): + raise ValueError(f"invalid shape {shape}") + if unknown: + known = math.prod(n for n in shape if n != -1) + if known == 0 or size % known: + raise ValueError(f"cannot reshape an array of size {size} into shape {shape}") + i = unknown[0] + shape = (*shape[:i], size // known, *shape[i + 1 :]) + if math.prod(shape) != size: + raise ValueError(f"cannot reshape an array of size {size} into shape {shape}") + return shape + + +def reshape(x, /, shape, *, copy=None): + _check_array(x, "reshape") + shape = _new_shape(shape, x.size) + T = x._T() + if x.ndim == 0 or T.nvals == 0: + return Array._new(Tensor(T.dtype, shape), x._fill) + if shape == (): + # (one position: its entry, as a 0-d Array's value) + return Array._new(_utils.empty_0d(T.dtype), T.to_coo()[1][0]) + if shape == T.shape: + return x._copy() if copy else Array._new(T.dup(), x._fill) + return Array._new(T.reshape(shape, copy=copy), x._fill) + + +def permute_dims(x, /, axes): + _check_array(x, "permute_dims") + axes = tuple(_normalize_axis(a, x.ndim) for a in axes) + if sorted(axes) != list(range(x.ndim)): + raise ValueError(f"axes {axes} must be a permutation of range({x.ndim})") + if x.ndim == 0: + return x._copy() + return Array._new(x._T().permute_dims(axes), x._fill) + + +def expand_dims(x, /, axis=0): + _check_array(x, "expand_dims") + axes = axis if isinstance(axis, tuple) else (axis,) + ndim = x.ndim + len(axes) + norm = sorted(_normalize_axis(a, ndim) for a in axes) + if len(set(norm)) != len(norm): + raise ValueError(f"repeated axis in {axis}") + T = x._T() + if x.ndim == 0: + return Array._new(Tensor(T.dtype, (1,) * ndim), x._fill) + return Array._new(T.expand_dims(tuple(norm)), x._fill) + + +def squeeze(x, /, axis): + _check_array(x, "squeeze") + axes = axis if isinstance(axis, tuple) else (axis,) + norm = [_normalize_axis(a, x.ndim) for a in axes] + if len(set(norm)) != len(norm): + raise ValueError(f"repeated axis in {axis}") + for a in norm: + if x.shape[a] != 1: + raise ValueError(f"cannot squeeze axis {a} of size {x.shape[a]}") + if not norm: + return x._copy() + T = x._T() + shape = tuple(n for i, n in enumerate(x.shape) if i not in norm) + if not shape: + return reshape(x, ()) + return Array._new(T.squeeze(tuple(norm)), x._fill) + + +def broadcast_to(x, /, shape): + """``x`` broadcast to ``shape``: its entries replicated over the broadcast axes (the + result's entries, under the densify limit for the replication). + """ + _check_array(x, "broadcast_to") + return _broadcast_to(x, shape) + + +def broadcast_shapes(*shapes): + return tuple(np.broadcast_shapes(*shapes)) + + +def broadcast_arrays(*arrays): + for x in arrays: + _check_array(x, "broadcast_arrays") + shape = np.broadcast_shapes(*[x.shape for x in arrays]) + return tuple(_broadcast_to(x, shape) for x in arrays) + + +def _broadcast_to(x, shape): + shape = tuple(int(n) for n in shape) + if np.broadcast_shapes(x.shape, shape) != shape: + raise ValueError(f"cannot broadcast shape {x.shape} to {shape}") + if x.shape == shape: + return x._copy() + T = x._T() + if x.ndim == 0 or T.nvals == 0: + return Array._new(Tensor(T.dtype, shape), x._fill) + E = Tensor(T.dtype, shape) + _utils.expansion_check(T, E, shape) + first = binary.first[gb_dtype(x.dtype)] + return Array._new(T.ewise_add(E, first, allow_expansion=True).new(), x._fill) + + +def concat(arrays, /, *, axis=0): + """Join arrays along ``axis`` (None: flattened first), O(entries), in parallel. + + The result's fill is the one held by the most positions; a part with another fill stores + its fill explicitly (O(that part's size), under the densify limit). + """ + from ._dtypes import result_type + from ._elementwise import cast + + arrays = list(arrays) + if not arrays: + raise ValueError("concat needs at least one array") + for a in arrays: + _check_array(a, "concat") + if axis is None: + arrays = [reshape(a, (-1,)) for a in arrays] + axis = 0 + ndim = arrays[0].ndim + if ndim == 0 or any(a.ndim != ndim for a in arrays): + raise ValueError("concat needs arrays of one number of dimensions, at least 1") + axis = _normalize_axis(axis, ndim) + shape = list(arrays[0].shape) + for a in arrays[1:]: + if any(a.shape[i] != n for i, n in enumerate(shape) if i != axis): + raise ValueError(f"shapes {arrays[0].shape} and {a.shape} do not concatenate") + shape[axis] = sum(a.shape[axis] for a in arrays) + shape = tuple(shape) + dtype = result_type(*arrays) + arrays = [a if a.dtype == dtype else cast(a, dtype) for a in arrays] + fill = _common_fill(arrays) + # (one GxB_Matrix_concat of the stored Matrices; a part stored in another layout moves) + T = Tensor._concat([_utils.with_fill(a, fill)._T() for a in arrays], axis) + return Array._new(T, fill) + + +def _common_fill(arrays): + """The fill held by the most positions of ``arrays`` (the first such, on ties).""" + totals = [] + for a in arrays: + for i, (fill, _) in enumerate(totals): + if _utils.same_value(fill, a._fill): + totals[i] = (fill, totals[i][1] + a.size) + break + else: + totals.append((a._fill, a.size)) + best = max(n for _, n in totals) + return next(fill for fill, n in totals if n == best) + + +def stack(arrays, /, *, axis=0): + arrays = list(arrays) + if not arrays: + raise ValueError("stack needs at least one array") + for a in arrays: + _check_array(a, "stack") + if len({a.shape for a in arrays}) > 1: + raise ValueError(f"stack needs arrays of one shape, not {[a.shape for a in arrays]}") + axis = _normalize_axis(axis, arrays[0].ndim + 1) + return concat([expand_dims(a, axis=axis) for a in arrays], axis=axis) + + +def unstack(x, /, *, axis=0): + _check_array(x, "unstack") + if x.ndim == 0: + raise ValueError("unstack needs an array of at least one dimension") + axis = _normalize_axis(axis, x.ndim) + before = (slice(None),) * axis + return tuple(x[(*before, i)] for i in range(x.shape[axis])) + + +def _axes(axis, ndim): + if axis is None: + return tuple(range(ndim)) + axes = axis if isinstance(axis, tuple) else (axis,) + norm = tuple(_normalize_axis(a, ndim) for a in axes) + if len(set(norm)) != len(norm): + raise ValueError(f"repeated axis in {axis}") + return norm + + +def flip(x, /, *, axis=None): + _check_array(x, "flip") + axes = _axes(axis, x.ndim) + if x.ndim == 0: + return x._copy() + return x[tuple(slice(None, None, -1) if a in axes else slice(None) for a in range(x.ndim))] + + +def moveaxis(x, source, destination, /): + _check_array(x, "moveaxis") + src = source if isinstance(source, tuple) else (source,) + dst = destination if isinstance(destination, tuple) else (destination,) + if len(src) != len(dst): + raise ValueError("moveaxis needs as many destinations as sources") + src = _axes(src, x.ndim) + dst = _axes(dst, x.ndim) + order = [a for a in range(x.ndim) if a not in src] + for d, s in sorted(zip(dst, src, strict=True)): + order.insert(d, s) + return permute_dims(x, tuple(order)) + + +def _entries(x): + """The stored entries of ``x`` as (coordinates, values), in stored order.""" + T = x._T() + if T.nvals == 0: + return [np.empty(0, np.int64) for _ in range(x.ndim)], np.empty(0, x.dtype) + idx, values = T.to_coo(sort=False) + return list(idx), values + + +def _from_entries(coords, values, shape, x): + """An Array of ``x``'s dtype and fill storing ``values`` at ``coords`` in ``shape``.""" + shape = tuple(shape) + if len(values) == 0: + return Array._new(Tensor(x._T().dtype, shape), x._fill) + T = Tensor.from_coo(tuple(coords), values, x._T().dtype, shape=shape) + return Array._new(T, x._fill) + + +def repeat(x, repeats, /, *, axis=None): + """Each element repeated along ``axis``, O(entries of the result).""" + _check_array(x, "repeat") + if axis is None: + x = reshape(x, (-1,)) + axis = 0 + axis = _normalize_axis(axis, x.ndim) + n = x.shape[axis] + if isinstance(repeats, Array): + if repeats.dtype.kind not in "iu" or repeats.ndim != 1: + raise TypeError("repeats must be an int or a 1-d integer array") + counts = repeats.todense().astype(np.int64) + else: + counts = np.asarray(int(repeats), dtype=np.int64) + counts = np.broadcast_to(counts, (n,)) + if (counts < 0).any(): + raise ValueError("repeats must be nonnegative") + starts = np.zeros(n + 1, np.int64) + np.cumsum(counts, out=starts[1:]) + shape = list(x.shape) + shape[axis] = int(starts[-1]) + coords, values = _entries(x) + k = counts[coords[axis]] + which = np.repeat(np.arange(len(values)), k) + within = np.arange(which.size) - np.repeat(np.cumsum(k) - k, k) + coords = [c[which] for c in coords] + coords[axis] = starts[:-1][coords[axis]] + within + return _from_entries(coords, values[which], shape, x) + + +def roll(x, /, shift, *, axis=None): + _check_array(x, "roll") + if axis is None: + return reshape(roll(reshape(x, (-1,)), shift, axis=0), x.shape) + shifts = shift if isinstance(shift, tuple) else (shift,) + axes = axis if isinstance(axis, tuple) else (axis,) + if len(shifts) == 1 and len(axes) > 1: + shifts = shifts * len(axes) + if len(shifts) != len(axes): + raise ValueError("roll needs one shift per axis, or one shift") + total = [0] * x.ndim + for s, a in zip(shifts, axes, strict=True): + total[_normalize_axis(a, x.ndim)] += int(s) + coords, values = _entries(x) + for a, s in enumerate(total): + if s and x.shape[a]: + coords[a] = (coords[a] + s) % x.shape[a] + return _from_entries(coords, values, x.shape, x) + + +def tile(x, repetitions, /): + """``x`` repeated ``repetitions`` times along each axis, O(entries of the result).""" + _check_array(x, "tile") + reps = tuple(int(r) for r in repetitions) + if any(r < 0 for r in reps): + raise ValueError("repetitions must be nonnegative") + ndim = max(len(reps), x.ndim) + reps = (1,) * (ndim - len(reps)) + reps + if x.ndim < ndim: + x = reshape(x, (1,) * (ndim - x.ndim) + x.shape) + shape = tuple(n * r for n, r in zip(x.shape, reps, strict=True)) + coords, values = _entries(x) + ntiles = math.prod(reps) + if ntiles == 0 or len(values) == 0: + return Array._new(Tensor(x._T().dtype, shape), x._fill) + tiles = np.unravel_index(np.arange(ntiles), reps) + out = [] + for c, t, n in zip(coords, tiles, x.shape, strict=True): + out.append((c[None, :] + (t * n)[:, None]).reshape(-1)) + return _from_entries(out, np.tile(values, ntiles), shape, x) + + +def refill(x, /, fill_value): + """``x`` with the fill value ``fill_value`` (an extension): the same value at every + position. Positions ``x`` does not store held its old fill, so they are stored now + (O(size), under the densify limit; O(1) when ``x`` stores every position), and stored + entries equal to ``fill_value`` go (``prune``). A 0-d array's value is its fill, so it + is returned as it is. ``fill_value`` must be a value of ``x``'s dtype. + """ + if not isinstance(x, Array): + raise TypeError(f"refill needs an Array, not {type(x).__name__}") + with np.errstate(all="ignore"): + fill = _utils.scalar(fill_value, x.dtype) + back = np.asarray(fill).astype(np.result_type(np.asarray(fill_value), x.dtype))[()] + if not _utils.same_value(back, np.asarray(fill_value)[()]): + raise ValueError(f"refill: {fill_value!r} is not a value of {x.dtype}") + if x.ndim == 0: + return x._copy() + T = x._T() + if _utils.same_value(fill, x._fill): + y = x + elif T.nvals == x.size: + y = Array._new(T, fill) + else: + y = _utils.with_fill(x, fill) + return prune(y) + + +def prune(x, /): + """``x`` without the stored entries that equal its fill value (an extension): the same + values, fewer stored entries. "Equal" is "identical": with fill 0.0 a stored -0.0 stays, + and with fill NaN the stored NaNs go. ``nnz`` counts stored entries, these included, + until they are pruned. + """ + if not isinstance(x, Array): + raise TypeError(f"prune needs an Array, not {type(x).__name__}") + T = x._T() + fill = x._fill + if T.nvals == 0: + return x._copy() + kind = x.dtype.kind + if kind in "biu" or (kind == "f" and fill != 0 and not np.isnan(fill)): + # (a value equal to a nonzero number is that number) + return Array._new(T.select(select.valuene, fill).new(), fill) + if fill == 0 and not _utils.fact(x, "zero"): + return x._copy() # (no stored zero of either sign) + keep = _utils.map_values(lambda v: _utils.differs(v, fill), [T], np.dtype(bool)) + R = Tensor(T.dtype, T.shape) + R(keep.V) << T + return Array._new(R, fill) diff --git a/graphblas/array_api/_searching.py b/graphblas/array_api/_searching.py new file mode 100644 index 000000000..33631c95d --- /dev/null +++ b/graphblas/array_api/_searching.py @@ -0,0 +1,383 @@ +"""Searching functions: argmax, argmin, count_nonzero, nonzero, searchsorted, where.""" + +import builtins +import math + +import numpy as np + +from ..core.tensor import Tensor +from ..core.vector import Vector +from . import _utils +from ._array import Array +from ._dtypes import check_supported +from ._utils import gb_dtype + +__all__ = ["argmax", "argmin", "count_nonzero", "nonzero", "searchsorted", "where"] + + +def _operand(x, name): + if isinstance(x, Array): + return x + if isinstance(x, (np.generic, np.ndarray)): + from ._creation import asarray + + return asarray(x) + if isinstance(x, (builtins.bool, int, float, complex)): + return x + raise TypeError(f"{name} needs arrays or Python scalars, not {type(x).__name__}") + + +def _probe(x): + return np.empty(0, x.dtype) if isinstance(x, Array) else x + + +def _value(x): + return np.asarray(x._fill) if isinstance(x, Array) else x + + +def where(condition, x1, x2, /): + """``x1`` where ``condition`` is true, else ``x2`` (broadcast); numpy's ``where``. + + Where the condition stores nothing it reads its fill, so the result there is the operand + that fill selects. When that operand needs no broadcast (or is a scalar, or stores + nothing), the result is that operand with the condition's stored positions written over: + O(entries of the condition and of that operand), the other operand read only at the + condition's entries (``xp.where(graph, labels, big)``, with ``labels`` broadcast along an + axis of ``graph``, stores ``graph``'s entries). Otherwise the entries are the union of + the three arrays' entries, with the broadcast operands' entries replicated. + """ + cond = _operand(condition, "where") + if not isinstance(cond, Array): + from ._creation import asarray + + cond = asarray(cond) + a = _operand(x1, "where") + b = _operand(x2, "where") + rdt = check_supported(np.result_type(_probe(a), _probe(b)), "where") + fill = np.where(np.asarray(cond._fill, dtype=bool), _value(a), _value(b)).astype(rdt)[()] + shapes = [x.shape for x in (cond, a, b) if isinstance(x, Array)] + shape = np.broadcast_shapes(*shapes) + arrays = [x for x in (cond, a, b) if isinstance(x, Array) and x.ndim] + if not shape or all(x._T().nvals == 0 for x in arrays): + return Array._new(Tensor(gb_dtype(rdt), shape), fill) + if (R := _where_at_condition(cond, a, b, shape, rdt)) is not None: + return Array._new(R, fill) + tensors = iter(_utils.aligned(arrays, shape)) + slots = [next(tensors) if isinstance(x, Array) and x.ndim else None for x in (cond, a, b)] + constants = [_value(x) for x in (cond, a, b)] + + def func(*values): + it = iter(values) + c, v1, v2 = ( + next(it) if t is not None else k for t, k in zip(slots, constants, strict=True) + ) + return np.where(np.asarray(c, dtype=bool), v1, v2).astype(rdt, copy=False) + + R = _utils.map_values(func, [t for t in slots if t is not None], rdt) + return Array._new(R, fill) + + +def _where_at_condition(cond, a, b, shape, rdt): + """``where(cond, a, b)``'s Tensor built from the operand ``cond``'s fill selects, with + ``cond``'s stored positions written over (see ``where``), or None when that operand + would have to be broadcast over stored entries, or ``cond`` itself broadcasts. + """ + from .. import binary + + Tc = cond._T() + if Tc.shape != shape: + return None + base = a if bool(cond._fill) else b + if isinstance(base, Array) and base.ndim and base.shape != shape and base._T().nvals: + return None + gdt = gb_dtype(rdt) + has_base = isinstance(base, Array) and base.ndim and base._T().nvals + other = b if bool(cond._fill) else a + if cond.dtype.kind != "c": + # All in GraphBLAS, in parallel. The condition's entries that select the base (those + # equal to its fill) hold the base's values, which the base supplies; only the others + # are computed, and the other operand is read there only (where(x > t, x, 0) + # stores x's entries > t; where(frontier, k, level) is a select and a union). + R = _selecting(Tc, cond, other, shape, rdt) + if has_base: + R = base._T().ewise_add(R, binary.second[gdt]).new(dtype=gdt) + return R + if not has_base: + # (a complex condition: through numpy) + coords, c = Tc.to_coo() + c = np.asarray(c, dtype=bool) + keep = c != bool(cond._fill) + coords = tuple(i[keep] for i in coords) + if isinstance(other, Array) and other.ndim: + V = _values_at(coords, other, shape) if len(coords[0]) else c[keep] + else: + V = np.full(len(coords[0]), _value(other)) + V = np.asarray(V).astype(rdt, copy=False) + return Tensor.from_coo(coords, V, gdt, shape=shape) if len(V) else Tensor(gdt, shape) + coords, c = Tc.to_coo() + c = np.asarray(c, dtype=bool) + values = [ + ( + (_values_at(coords, x, shape) if len(c) else c) + if isinstance(x, Array) and x.ndim + else _value(x) + ) + for x in (a, b) + ] + V = np.where(c, *values).astype(rdt, copy=False) + R = Tensor.from_coo(coords, V, gdt, shape=shape) if len(c) else Tensor(gdt, shape) + return base._T().ewise_add(R, binary.second[gdt]).new(dtype=gdt) + + +def _selecting(Tc, cond, other, shape, rdt): + """The entries of ``Tc`` (cond's Tensor, not complex) that select ``other`` (the truthy + ones, or the falsy ones under a truthy fill; NaN is truthy), holding ``other``'s values: + its fill, then its stored values written over. A broadcast ``other`` is read through a + diagonal scaling, without replicating it (xp.where(graph, labels, big)). + """ + from .. import binary, select + from ..dtypes import BOOL + + gdt = gb_dtype(rdt) + pick = select.valueeq if bool(cond._fill) else select.valuene + sel = Tc.select(pick, False if Tc.dtype == BOOL else 0).new() + second = binary.second[gdt] + To = other._T() if isinstance(other, Array) and other.ndim else None + if To is not None and other.shape != shape and other.size <= 4 * sel.nvals: + # (stored at every position, the broadcast operand makes the product a diagonal + # scaling, one pass; with gaps in the diagonal SuiteSparse multiplies in general, + # 17x slower at 1e7 entries) + if To.nvals < other.size: + To = _utils.with_fill(other, 0 if other._fill else 1)._T() + return sel.ewise_mult(To, second).new(dtype=gdt) + R = sel.apply(second, right=_utils.gb_scalar(_value(other), rdt)).new(dtype=gdt) + if To is not None and To.nvals: + R(accum=second) << sel.ewise_mult(To, second) + return R + + +def _values_at(coords, x, shape): + """The values of ``x`` (an Array that broadcasts to ``shape``) at ``coords`` of + ``shape``: its stored value, or its fill. O(len(coords)), read from x's own entries. + """ + lead = len(shape) - x.ndim + n = len(coords[0]) + # (a broadcast axis of x has size 1: index 0 there) + index = tuple( + coords[lead + i] if size != 1 else np.zeros(n, np.int64) for i, size in enumerate(x.shape) + ) + return x._T().vindex[index].to_dense(fill_value=x._fill) + + +def _check_array(x, name): + if not isinstance(x, Array): + raise TypeError(f"{name} needs an Array, not {type(x).__name__}") + + +def _arg(name, x, axis, keepdims): + """The first position of the extreme value of each fiber (argmax, argmin), NaN counting + as the extreme (numpy's rule), computed in GraphBLAS on the stored entries. + + With the fibers as the rows (or columns) of the stored Matrix: the extreme of each fiber + (a reduction), the entries equal to it (the Matrix scaled by the extremes as a diagonal, + under ``eq``), and the smallest position among those (a positional reduction); likewise + the first stored NaN. A fiber with an unstored position holds the fill there: when the + fill beats or ties the stored extreme, or is NaN, the fiber's first unstored position + takes part. That is 0 unless the fiber stores position 0, and otherwise is found from + that fiber's entries alone. Each step is a pass over the entries; nothing is sorted. + """ + from ._manipulation import _normalize_axis, expand_dims, reshape + + _check_array(x, name) + if x.dtype.kind == "c": + raise TypeError(f"{name} needs a real-valued array, not {x.dtype}") + if axis is None: + out = _arg(name, reshape(x, (-1,)), 0, False) + if keepdims: + out = reshape(out, (1,) * x.ndim) + return out + axis = _normalize_axis(axis, x.ndim) + L = x.shape[axis] + if L == 0: + raise ValueError(f"{name} of an empty axis") + out_shape = tuple(n for a, n in enumerate(x.shape) if a != axis) + T = x._T() + if T.nvals == 0 or math.prod(out_shape) == 0: + out = Array._new(Tensor(gb_dtype(np.int64), out_shape), 0) + return expand_dims(out, axis) if keepdims else out + R = _arg_fibers(name, x, T, axis) + if out_shape: + out = Array._new(Tensor.from_vector(R).reshape(out_shape), 0) + else: + out = Array._new(_utils.empty_0d(np.int64), np.int64(R.get(0, 0))) + return expand_dims(out, axis) if keepdims else out + + +def _arg_fibers(name, x, T, axis): + """A Vector over the fibers of ``T`` along ``axis`` (in C order) holding the first + position of each one's extreme (see ``_arg``), storing only the positions that are not 0. + """ + from .. import binary, monoid, select, semiring, unary + from ..dtypes import BOOL, INT64 + + L = T.shape[axis] + others = tuple(a for a in range(T.ndim) if a != axis) + # The fibers as the rows of M (byrow) or its columns: as stored when the layout allows + if T._layout == (others, (axis,)): + M, byrow = T._M, True + elif T._layout == ((axis,), others): + M, byrow = T._M, False + else: + M, byrow = T.to_matrix(rows=others), True + nfibers = M.nrows if byrow else M.ncols + if M.dtype == BOOL: + top = monoid.lor if name == "argmax" else monoid.land + else: + top = monoid.max if name == "argmax" else monoid.min + every = Vector(BOOL, L) + every << True # (iso: no values stored) + + def first(P, mask=None): + # (the smallest position along each fiber of P's entries) + v = Vector(INT64, nfibers) + if byrow: + v(mask=mask) << P.mxv(every, semiring.ss.min_firstj) + else: + v(mask=mask) << every.vxm(P, semiring.ss.min_secondi) + return v + + ext = (M.reduce_rowwise(top) if byrow else M.reduce_columnwise(top)).new() # (skips NaN) + # The entries equal to their fiber's extreme. With a value at every diagonal position, + # SuiteSparse runs the product as a row (column) scaling, one pass over M. + if nfibers <= 4 * M.nvals: + diag = Vector(ext.dtype, nfibers) + diag << 0 # (the value at a fiber with no entries scales nothing) + diag(accum=binary.second) << ext + else: + diag = ext + D = diag.diag() + eq = semiring.any_eq[M.dtype] + hits = (D.mxm(M, eq) if byrow else M.mxm(D, eq)).new() + R = first(hits.select(select.valueeq, True).new()) + del hits + # The first stored NaN, which beats every value + nan = None + if x.dtype.kind == "f" and _utils.fact(x, "nan"): + nan = first(M.apply(unary.isnan).new().select(select.valueeq, True).new()) + R(accum=binary.second) << nan + # The fibers where the fill takes part, if they have an unstored position: those whose + # stored extreme the fill beats or ties (and that store no NaN), or all, for a NaN fill + fill = x._fill + fill_nan = x.dtype.kind == "f" and np.isnan(fill) + if fill_nan: + need = ext + else: + beats = select.valuelt if name == "argmax" else select.valuegt + need = ext.select(beats, fill).new() + need(accum=binary.first) << ext.select(select.valueeq, fill) + if nan is not None: + keep = Vector(need.dtype, nfibers) + keep(~nan.S, replace=True) << need + need = keep + if need.nvals: + # A fiber whose first entry is past position 0 has its first unstored position at 0 + firsts = first(M, need.S) + R(~firsts.select(select.valuegt, 0).new().S, replace=True) << R # (there: 0) + at_0 = firsts.select(select.valueeq, 0).new() + if at_0.nvals: + _fill_at_0(name, M, byrow, R, ext, nan, fill, fill_nan, at_0, L) + return R.select(select.valuene, 0).new() + + +def _fill_at_0(name, M, byrow, R, ext, nan, fill, fill_nan, at_0, L): + """Update R at the fibers ``at_0`` (those that take the fill and store position 0) with + the fill's first position, read from those fibers' entries: the first position not equal + to its rank among the fiber's entries, or the number of entries when they fill a prefix + (none when that number is L: the fiber stores every position). + """ + fibers = at_0.to_coo(values=False)[0] + sub = M[fibers, :].new() if byrow else M[:, fibers].new().T.new() + r, c, _ = sub.to_coo() # (by fiber, positions sorted) + starts = np.flatnonzero(np.r_[True, r[1:] != r[:-1]]) + counts = np.diff(np.r_[starts, len(r)]) + rank = np.arange(len(r)) - np.repeat(starts, counts) + gap = np.minimum.reduceat(np.where(c != rank, rank, np.repeat(counts, counts)), starts) + best = R[fibers].new().to_dense(fill_value=L) + if fill_nan: + stored_nan = nan[fibers].new().to_dense(fill_value=L) if nan is not None else L + new = np.where(gap < L, np.minimum(gap, stored_nan), best) + else: + e = ext[fibers].new().to_dense() + wins = (fill > e) if name == "argmax" else (fill < e) + new = np.where(gap < L, np.where(wins, gap, np.minimum(gap, best)), best) + R[fibers] << Vector.from_coo(np.arange(len(fibers)), new.astype(np.int64), size=len(fibers)) + + +def argmax(x, /, *, axis=None, keepdims=False): + return _arg("argmax", x, axis, keepdims) + + +def argmin(x, /, *, axis=None, keepdims=False): + return _arg("argmin", x, axis, keepdims) + + +def count_nonzero(x, /, *, axis=None, keepdims=False): + from ._elementwise import not_equal + from ._statistical import sum as sum_ + + _check_array(x, "count_nonzero") + return sum_(not_equal(x, 0), axis=axis, dtype=np.dtype(np.int64), keepdims=keepdims) + + +def nonzero(x, /): + """The coordinates (C order) of the nonzero positions: O(entries) when the fill is zero, + O(size) otherwise (every unstored position is nonzero). + """ + from ._creation import asarray + from ._elementwise import not_equal + + _check_array(x, "nonzero") + if x.ndim == 0: + raise ValueError("nonzero needs an array of at least one dimension") + nz = not_equal(x, 0) + if nz._fill: + _utils.check_dense(x.size, "nonzero of an array whose fill is not zero") + coords = np.nonzero(nz.todense()) + else: + T = nz._T() + if T.nvals == 0: + coords = tuple(np.empty(0, np.int64) for _ in range(x.ndim)) + else: + coords = T.select("==", True).new().to_coo(values=False)[0] + return tuple(asarray(np.asarray(c, dtype=np.int64)) for c in coords) + + +def searchsorted(x1, x2, /, *, side="left", sorter=None): + """Insertion positions in the sorted 1-d ``x1``: the stored values are searched, and the + run of the fill counted where it sorts. O(entries of x1 + entries of x2). + """ + from ._creation import asarray + from ._indexing import take + + _check_array(x1, "searchsorted") + if x1.ndim != 1: + raise ValueError("searchsorted needs a 1-d x1") + if side not in {"left", "right"}: + raise ValueError(f"side must be 'left' or 'right', not {side!r}") + if sorter is not None: + x1 = take(x1, sorter) + T = x1._T() + stored = T.to_coo()[1] if T.nvals else np.empty(0, x1.dtype) + k = x1.size - T.nvals + fill = np.asarray([x1._fill]) + + def count(v): + n = np.searchsorted(stored, v, side=side) + return (n + k * np.searchsorted(fill, v, side=side)).astype(np.int64) + + if not isinstance(x2, Array): + x2 = asarray(x2) + out_fill = count(np.asarray(x2._fill)) + if x2.ndim == 0 or x2._T().nvals == 0: + return Array._new(Tensor(gb_dtype(np.int64), x2.shape), out_fill) + return Array._new(_utils.map_values(count, [x2._T()], np.dtype(np.int64)), out_fill) diff --git a/graphblas/array_api/_set.py b/graphblas/array_api/_set.py new file mode 100644 index 000000000..ef95a37b2 --- /dev/null +++ b/graphblas/array_api/_set.py @@ -0,0 +1,162 @@ +"""Set functions: isin, unique_all, unique_counts, unique_inverse, unique_values. + +The unique values of an array are those of its stored entries plus the fill, counted once per +unstored position: O(entries log entries). NaNs are distinct (each is its own value), so a +fill of NaN with unstored positions makes as many distinct values: that case goes through +numpy densely, under the densify limit. Values come sorted, as numpy's. +""" + +from typing import NamedTuple + +import numpy as np + +from ..core.tensor import Tensor +from . import _utils +from ._array import Array +from ._utils import gb_dtype + +__all__ = ["isin", "unique_all", "unique_counts", "unique_inverse", "unique_values"] + + +class UniqueAllResult(NamedTuple): + values: Array + indices: Array + inverse_indices: Array + counts: Array + + +class UniqueCountsResult(NamedTuple): + values: Array + counts: Array + + +class UniqueInverseResult(NamedTuple): + values: Array + inverse_indices: Array + + +def _vector(a): + from ._creation import _from_numpy + + return _from_numpy(np.asarray(a)) + + +def _unique(x, *, index=False, inverse=False, counts=False): + """The sorted unique values of ``x`` (as numpy's unique with equal_nan=False), and the + first flat index, the inverse (an Array of x's shape) and the counts when asked. + """ + if not isinstance(x, Array): + raise TypeError(f"needs an Array, not {type(x).__name__}") + T = x._T() + fill = x._fill + k = x.size - T.nvals + if k and x.dtype.kind in "fc" and np.isnan(fill): + _utils.check_dense(x.size, "unique of an array whose fill is NaN (each NaN is distinct)") + return _unique_dense(x, index, inverse, counts) + if T.nvals: + coords, values = T.to_coo() + flat = np.ravel_multi_index(coords, x.shape) if x.ndim else np.zeros(1, np.int64) + else: + coords = tuple(np.empty(0, np.int64) for _ in range(x.ndim)) + values, flat = np.empty(0, x.dtype), np.empty(0, np.int64) + if k: + # The fill once, at its first unstored position, with the other k - 1 added later + first = _first_unstored(flat, x.size) + at = np.searchsorted(flat, first) + values = np.insert(values, at, fill) + flat = np.insert(flat, at, first) + out = np.unique( + values, return_index=True, return_inverse=True, return_counts=True, equal_nan=False + ) + uvalues, uindex, uinverse, ucounts = out + uinverse = uinverse.reshape(-1) + result = {"values": _vector(uvalues)} + if k: + fill_slot = uinverse[at] + ucounts = ucounts.copy() + ucounts[fill_slot] += k - 1 + if index: + result["indices"] = _vector(flat[uindex].astype(np.int64)) + if counts: + result["counts"] = _vector(ucounts.astype(np.int64)) + if inverse: + inv_fill = int(uinverse[at]) if k else 0 + if k: + uinverse = np.delete(uinverse, at) + if x.ndim == 0: + inv = Array._new(_utils.empty_0d(np.int64), uinverse[0] if len(uinverse) else inv_fill) + elif T.nvals: + R = Tensor.from_coo( + coords, uinverse.astype(np.int64), gb_dtype(np.int64), shape=x.shape + ) + inv = Array._new(R, inv_fill) + else: + inv = Array._new(Tensor(gb_dtype(np.int64), x.shape), inv_fill) + result["inverse_indices"] = inv + return result + + +def _first_unstored(flat, size): + """The smallest flat position in ``range(size)`` not in the sorted ``flat``.""" + gaps = np.flatnonzero(flat != np.arange(len(flat))) + return int(gaps[0]) if len(gaps) else len(flat) + + +def _unique_dense(x, index, inverse, counts): + from ._creation import _from_numpy + + dense = x.todense() + uvalues, uindex, uinverse, ucounts = np.unique( + dense.reshape(-1), return_index=True, return_inverse=True, return_counts=True, + equal_nan=False, + ) # fmt: skip + result = {"values": _vector(uvalues)} + if index: + result["indices"] = _vector(uindex.astype(np.int64)) + if counts: + result["counts"] = _vector(ucounts.astype(np.int64)) + if inverse: + result["inverse_indices"] = _from_numpy(uinverse.reshape(x.shape).astype(np.int64)) + return result + + +def unique_values(x, /): + return _unique(x)["values"] + + +def unique_counts(x, /): + r = _unique(x, counts=True) + return UniqueCountsResult(r["values"], r["counts"]) + + +def unique_inverse(x, /): + r = _unique(x, inverse=True) + return UniqueInverseResult(r["values"], r["inverse_indices"]) + + +def unique_all(x, /): + r = _unique(x, index=True, inverse=True, counts=True) + return UniqueAllResult(r["values"], r["indices"], r["inverse_indices"], r["counts"]) + + +def isin(x1, x2, /, *, invert=False): + """Whether each element of ``x1`` is among ``x2``'s values (its stored values and its + fill): O(entries of both). + """ + from ._creation import asarray + + x1 = x1 if isinstance(x1, Array) else asarray(x1) + x2 = x2 if isinstance(x2, Array) else asarray(x2) + T2 = x2._T() + values = T2.to_coo(indices=False)[1] if T2.nvals else np.empty(0, x2.dtype) + if x2.size > T2.nvals: + values = np.append(values, x2._fill) + + def member(v): + return np.isin(v, values, invert=invert) + + fill = member(np.asarray(x1._fill)) + T1 = x1._T() + if x1.ndim == 0 or T1.nvals == 0: + return Array._new(Tensor(gb_dtype(bool), x1.shape), fill) + return Array._new(_utils.map_values(member, [T1], np.dtype(bool)), fill) diff --git a/graphblas/array_api/_sorting.py b/graphblas/array_api/_sorting.py new file mode 100644 index 000000000..d14489f49 --- /dev/null +++ b/graphblas/array_api/_sorting.py @@ -0,0 +1,123 @@ +"""Sorting functions: sort and argsort. + +``sort`` keeps the result as sparse as its input: along each fiber, the values below the fill +come first, then the run of values equal to the fill (the unstored positions among them), +then the values above it, NaN last; equal values keep their order (stable). Descending order +reverses the values but keeps equal values in their order, NaN first (as +``flip(sort(flip(x)))``). O(entries log entries). + +``argsort``'s result is a permutation of every position, dense: it goes through numpy, under +the densify limit. +""" + +import numpy as np + +from ..core.tensor import Tensor +from . import _utils +from ._array import Array + +__all__ = ["argsort", "sort"] + + +def _check(x, name): + if not isinstance(x, Array): + raise TypeError(f"{name} needs an Array, not {type(x).__name__}") + if x.ndim == 0: + raise ValueError(f"{name} needs an array of at least one dimension") + if x.dtype.kind == "c": + raise TypeError(f"{name} needs a real-valued array, not {x.dtype}") + + +def sort(x, /, *, axis=-1, descending=False, stable=True): + from ._manipulation import _normalize_axis, permute_dims + + _check(x, "sort") + axis = _normalize_axis(axis, x.ndim) + L = x.shape[axis] + T = x._T() + if T.nvals == 0 or L <= 1: + return x._copy() + others = tuple(a for a in range(x.ndim) if a != axis) + perm = (*others, axis) + coords, values = permute_dims(x, perm)._T().to_coo() # (by fiber, then position) + fshape = tuple(x.shape[a] for a in others) + fiber = np.ravel_multi_index(coords[:-1], fshape) if fshape else np.zeros(len(values), np.int64) + pos = coords[-1] + fill = x._fill + with np.errstate(invalid="ignore"): + if values.dtype.kind == "f": + nan_v, nan_f = np.isnan(values), bool(np.isnan(fill)) + equal = (values == fill) | (nan_v & nan_f) + below = (values < fill) | (~nan_v & nan_f) + else: + equal = values == fill + below = values < fill + above = ~(equal | below) + if descending: + below, above = above, below + # Rank the values (np.unique sorts NaN last and makes them one value) + _, rank = np.unique(values, return_inverse=True) + rank = rank.reshape(-1) + if descending: + rank = rank.max(initial=0) - rank + # Groups per fiber: 0 below the fill, 1 equal to it, 2 above it + group = np.where(below, 0, np.where(equal, 1, 2)) + order = np.lexsort((pos, rank, group, fiber)) + f_o, g_o, p_o = fiber[order], group[order], pos[order] + # Rank of each entry within its (fiber, group), and the sizes of the groups per fiber + key = f_o * 3 + g_o + starts = np.flatnonzero(np.r_[True, key[1:] != key[:-1]]) + sizes = np.diff(np.r_[starts, len(key)]) + within = np.arange(len(key)) - np.repeat(starts, sizes) + fibers, fiber_index = np.unique(f_o, return_inverse=True) + counts = np.zeros((len(fibers), 3), np.int64) + np.add.at(counts, (fiber_index, g_o), 1) + n_below = counts[fiber_index, 0] + n_equal = counts[fiber_index, 1] + unstored = L - counts.sum(axis=1)[fiber_index] + # An entry equal to the fill keeps its place among the run of the fill (stable): the + # positions before it in the run are every position before it but the other groups' + others_before = _count_before(f_o, p_o, g_o != 1) + new_pos = np.where( + g_o == 0, + within, + np.where(g_o == 1, n_below + p_o - others_before, n_below + n_equal + unstored + within), + ) + vals = values[order] + keep = (g_o != 1) | _utils.differs(vals, fill) + if not keep.any(): + return Array._new(Tensor(T.dtype, x.shape), fill) + out_coords = list(np.unravel_index(f_o[keep], fshape)) if fshape else [] + out_coords.append(new_pos[keep]) + R = Tensor.from_coo(out_coords, vals[keep], T.dtype, shape=tuple(x.shape[a] for a in perm)) + inverse = tuple(int(i) for i in np.argsort(perm)) + return permute_dims(Array._new(R, fill), inverse) + + +def _count_before(fiber, pos, flags): + """For each entry: how many flagged entries of its fiber have a smaller position.""" + order = np.lexsort((pos, fiber)) + f, fl = fiber[order], flags[order].astype(np.int64) + starts = np.flatnonzero(np.r_[True, f[1:] != f[:-1]]) + csum = np.cumsum(fl) - fl # (flagged before, over all fibers) + base = np.repeat(csum[starts], np.diff(np.r_[starts, len(f)])) + out = np.empty(len(fiber), np.int64) + out[order] = csum - base + return out + + +def argsort(x, /, *, axis=-1, descending=False, stable=True): + from ._creation import _from_numpy + from ._manipulation import _normalize_axis + + _check(x, "argsort") + axis = _normalize_axis(axis, x.ndim) + _utils.check_dense(x.size, "argsort (every position's index)") + dense = x.todense() + if descending: + # Descending and stable: equal values keep their order, as flip(sort(flip(x))) + flipped = np.flip(dense, axis=axis) + out = x.shape[axis] - 1 - np.flip(np.argsort(flipped, axis=axis, kind="stable"), axis=axis) + else: + out = np.argsort(dense, axis=axis, kind="stable") + return _from_numpy(out.astype(np.int64, copy=False)) diff --git a/graphblas/array_api/_statistical.py b/graphblas/array_api/_statistical.py new file mode 100644 index 000000000..111fe4593 --- /dev/null +++ b/graphblas/array_api/_statistical.py @@ -0,0 +1,591 @@ +"""Statistical functions, and the utility functions ``all``, ``any`` and ``diff``. + +A reduction over a fiber of ``L`` positions, ``k`` of them unstored, is the monoid over the +stored entries combined with the fill ``k`` times: ``plus`` adds ``k * fill``, ``times`` +multiplies by ``fill ** k``, and ``min``, ``max``, ``lor`` and ``land`` take the fill once +when ``k > 0``. The counts cost a second pass, skipped when the fill is the monoid's identity +or every position is stored. A fiber with no entries holds the result's fill. Floating-point +sums and products may round differently from numpy's (the order of the terms differs); +everything else is exact. +""" + +import builtins +import itertools +import math + +import numpy as np + +from .. import agg, binary, monoid, unary +from ..core.tensor import Tensor +from . import _utils +from ._array import Array +from ._dtypes import as_dtype, check_supported +from ._utils import gb_dtype + +__all__ = [ + "all", + "any", + "cumulative_prod", + "cumulative_sum", + "diff", + "max", + "mean", + "min", + "prod", + "std", + "sum", + "var", +] + + +def _check_array(x, name): + if not isinstance(x, Array): + raise TypeError(f"{name} needs an Array, not {type(x).__name__}") + + +def _axes(x, axis): + """The reduced axes, normalized and sorted.""" + if axis is None: + return tuple(range(x.ndim)) + axes = axis if isinstance(axis, tuple) else (axis,) + norm = [] + for a in axes: + a = int(a) + if not -x.ndim <= a < x.ndim: + raise IndexError(f"axis {a} is out of bounds for {x.ndim} dimensions") + norm.append(a % x.ndim) + if len(set(norm)) != len(norm): + raise ValueError(f"repeated axis in {axis}") + return tuple(sorted(norm)) + + +def _out_shape(shape, axes, keepdims): + if keepdims: + return tuple(1 if i in axes else n for i, n in enumerate(shape)) + return tuple(n for i, n in enumerate(shape) if i not in axes) + + +def _keepdims(out, x, axes, keepdims): + if keepdims and axes: + from ._manipulation import expand_dims + + return expand_dims(out, axes) + return out + + +def _reduce_tensor(T, op, axes): + axis = None if len(axes) == T.ndim else axes + return T.reduce(op, axis=axis).new() + + +def _combine(func, tensors, dtype, fill): + """An Array whose entries are ``func`` of the aligned values of ``tensors`` (one + pattern), and ``fill`` elsewhere. + """ + return Array._new(_utils.map_values(func, tensors, dtype), fill) + + +def _int_power(base, exponent, dtype): + """``base ** exponent`` wrapped to the integer ``dtype``, as repeated products wrap.""" + bits = dtype.itemsize * 8 + value = builtins.pow(int(base), int(exponent), 1 << bits) + if dtype.kind == "i" and value >= 1 << (bits - 1): + value -= 1 << bits + return np.asarray(value).astype(dtype)[()] + + +def _unstored(L, c, ct): + """``L - c`` (counts of unstored positions), in a dtype that does not overflow: uint64 + for integer sums (which wrap anyway), float64 otherwise (a fiber may hold 2**63 or more + positions). + """ + if ct.kind in "iub": + return np.uint64(L % (1 << 64)) - np.asarray(c).astype(np.uint64) + return np.float64(L) - c + + +def _scaled(k, fill, ct): + """``fill`` added ``k`` times (k >= 0), as repeated additions give it: wrapping for + integers, logical or for bool, and part by part for complex (0 * inf in one part must not + make NaN in the other). + """ + with np.errstate(all="ignore"): + if ct.kind in "iu": + f = np.asarray(fill).astype(np.int64 if ct.kind == "i" else np.uint64) + return (np.asarray(k, dtype=np.uint64) * f.astype(np.uint64)).astype(ct) + if ct.kind == "b": + return np.logical_and(np.asarray(k) > 0, fill) + if ct.kind == "c": + out = np.empty(np.shape(k), ct) + out.real = k * fill.real + out.imag = k * fill.imag + return out + return np.asarray(np.multiply(k, fill)).astype(ct) + + +class _Monoid: + def __init__(self, name, gb, npfunc, identity): + self.name = name + self.gb = gb + self.npfunc = npfunc + self.identity = identity # (dtype -> the identity) + + def repeat(self, fill, L, ct): + """The monoid over ``L`` copies of ``fill``.""" + if self.name == "sum": + k = np.uint64(L % (1 << 64)) if ct.kind in "iub" else np.float64(L) + return _scaled(k, fill, ct)[()] + if self.name == "prod": + with np.errstate(all="ignore"): + if ct.kind in "iu": + return _int_power(fill, L, ct) + if ct.kind == "b": + return fill + return np.power(fill, np.float64(L), dtype=ct) + return fill + + def combine(self, L, fill, ct): + """``func(r, count)``: a fiber's reduction of its stored entries ``r`` combined with + the fill at its ``L - count`` unstored positions. + """ + name = self.name + if name == "sum": + # (a fiber with no unstored position must not see 0 * inf) + return lambda r, c: np.where(c < L, r + _scaled(_unstored(L, c, ct), fill, ct), r) + if name == "prod": + if ct.kind in "iu": + return lambda r, c: r * _int_powers(fill, _unstored(L, c, ct), ct) + if ct.kind == "b": + return lambda r, c: np.where(c < L, r & fill, r) + return lambda r, c: np.where(c < L, r * np.power(fill, _unstored(L, c, ct)), r) + func = { + "min": np.minimum, + "max": np.maximum, + "any": np.logical_or, + "all": np.logical_and, + }[name] + return lambda r, c: np.where(c < L, func(r, fill), r) + + +def _int_powers(base, exponents, dtype): + """``base ** e`` for each count ``e`` (uint64), wrapped to the integer ``dtype``.""" + big = np.dtype("uint64" if dtype.kind == "u" else "int64") + with np.errstate(all="ignore"): + return np.power(big.type(base), exponents.astype(big)).astype(dtype) + + +def _min_identity(dt): + if dt.kind == "f": + return np.inf + if dt.kind == "b": + return True + return np.iinfo(dt).max + + +def _max_identity(dt): + if dt.kind == "f": + return -np.inf + if dt.kind == "b": + return False + return np.iinfo(dt).min + + +_MONOIDS = { + "sum": _Monoid("sum", monoid.plus, np.sum, lambda dt: 0), + "prod": _Monoid("prod", monoid.times, np.prod, lambda dt: 1), + "min": _Monoid("min", monoid.min, np.min, _min_identity), + "max": _Monoid("max", monoid.max, np.max, _max_identity), + "any": _Monoid("any", monoid.lor, np.any, lambda dt: False), + "all": _Monoid("all", monoid.land, np.all, lambda dt: True), +} + + +def _reduce(name, x, axis, keepdims, dtype=None): + out = _reduce_main(name, x, axis, keepdims, dtype) + if name == "sum" and out.dtype.kind in "fc": + # numpy's sum starts from +0.0: a sum of -0.0 is 0.0 (GraphBLAS starts from the first) + from ._elementwise import add + + out = add(out, 0.0) + return out + + +def _reduce_main(name, x, axis, keepdims, dtype): + _check_array(x, name) + spec = _MONOIDS[name] + if name in {"min", "max"} and x.dtype.kind == "c": + raise TypeError(f"{name} needs a real-valued array, not {x.dtype}") + if name in {"any", "all"}: + ct = np.dtype(bool) + elif name in {"sum", "prod"}: + if dtype is not None: + ct = as_dtype(dtype) + else: + ct = check_supported(spec.npfunc(np.zeros(1, x.dtype)).dtype, name) + else: + ct = x.dtype + axes = _axes(x, axis) + if x.ndim == 0: + kwargs = {"dtype": ct} if name in {"sum", "prod"} else {} + with np.errstate(all="ignore"): + value = spec.npfunc(np.asarray(x._fill), **kwargs) + return Array._new(_utils.empty_0d(ct), value) + if not axes: + return _reduce_no_axes(name, x, ct) + T = x._T() + shape = _out_shape(x.shape, axes, keepdims=False) + L = math.prod(x.shape[a] for a in axes) + with np.errstate(all="ignore"): + fill = _utils.scalar(x._fill, ct) + if L == 0: + if name in {"min", "max"}: + raise ValueError(f"{name} of an empty axis, which has no identity") + out = Array._new(Tensor(gb_dtype(ct), shape), spec.identity(ct)) + return _keepdims(out, x, axes, keepdims) + fill_out = spec.repeat(fill, L, ct) + if T.nvals == 0 or math.prod(shape) == 0: + return _keepdims(Array._new(Tensor(gb_dtype(ct), shape), fill_out), x, axes, keepdims) + nan_fix = name in {"min", "max"} and ct.kind == "f" and _has_nan(x) + full = T.nvals == x.size + if len(axes) == x.ndim: + # Over every axis: one Scalar, and the count is the number of entries + r = T.reduce_scalar(_gb_monoid(spec, ct)).new().value + if nan_fix: + value = np.nan + elif full or _utils.same_value(fill, spec.identity(ct)): + value = r + else: + with np.errstate(all="ignore"): + count = np.array([T.nvals], dtype=np.int64) + value = spec.combine(L, fill, ct)(np.array([r], dtype=ct), count)[0] + out = Array._new(_utils.empty_0d(ct), value) + return _keepdims(out, x, axes, keepdims) + R = _reduce_tensor(T, _gb_monoid(spec, ct), axes) + if not nan_fix and (full or _utils.same_value(fill, spec.identity(ct))): + out = Array._new(R, fill if full else fill_out) + else: + C = _reduce_tensor(T, agg.count, axes) + combine = spec.combine(L, fill, ct) + if nan_fix: + # GraphBLAS's min and max skip NaN; numpy's propagate it + N = _reduce_tensor(T.apply(unary.isnan).new(), monoid.lor[gb_dtype(bool)], axes) + + def func(r, c, n): + return np.where(n, np.nan, combine(r, c)) + + out = _combine(func, [R, C, N], ct, fill_out) + else: + out = _combine(combine, [R, C], ct, fill_out) + return _keepdims(out, x, axes, keepdims) + + +def _reduce_no_axes(name, x, ct): + """A reduction over no axes, as numpy computes it: the identity combined with each + element (so ``sum`` turns ``-0.0`` into ``0.0``), or a copy for ``min`` and ``max``. + """ + from ._elementwise import add, cast, multiply + + if name in {"sum", "prod"}: + identity = Array._new(_utils.empty_0d(ct), 0 if name == "sum" else 1) + out = (add if name == "sum" else multiply)(identity, x if x.dtype == ct else cast(x, ct)) + return out if out.dtype == ct else cast(out, ct) + return cast(x, ct) if x.dtype != ct else x._copy() + + +def _gb_monoid(spec, ct): + # (a sum of booleans in bool is their logical or, a product their logical and) + if ct.kind == "b" and spec.name in {"sum", "prod"}: + return (monoid.lor if spec.name == "sum" else monoid.land)[gb_dtype(ct)] + return spec.gb[gb_dtype(ct)] + + +def _has_nan(x): + return _utils.fact(x, "nan") + + +def sum(x, /, *, axis=None, dtype=None, keepdims=False): # noqa: A001 (the Array API's name) + return _reduce("sum", x, axis, keepdims, dtype) + + +def prod(x, /, *, axis=None, dtype=None, keepdims=False): + return _reduce("prod", x, axis, keepdims, dtype) + + +def min(x, /, *, axis=None, keepdims=False): + return _reduce("min", x, axis, keepdims) + + +def max(x, /, *, axis=None, keepdims=False): + return _reduce("max", x, axis, keepdims) + + +def any(x, /, *, axis=None, keepdims=False): # noqa: A001 (the Array API's name) + return _reduce("any", x, axis, keepdims) + + +def all(x, /, *, axis=None, keepdims=False): # noqa: A001 (the Array API's name) + return _reduce("all", x, axis, keepdims) + + +def _mean_dtype(x, name): + return check_supported(np.mean(np.zeros(1, x.dtype)).dtype, name) + + +def mean(x, /, *, axis=None, keepdims=False): + from ._elementwise import divide + + _check_array(x, "mean") + rdt = _mean_dtype(x, "mean") + L = math.prod(x.shape[a] for a in _axes(x, axis)) + # (an empty axis divides 0 by 0, as numpy does: nan, or nan+nanj for complex) + return divide(_reduce("sum", x, axis, keepdims, rdt), L) + + +def _squared_magnitude(x): + from ._elementwise import conj, multiply, real + + if x.dtype.kind == "c": + return real(multiply(x, conj(x))) + return multiply(x, x) + + +def var(x, /, *, axis=None, correction=0.0, keepdims=False): + """Two passes: the means, then the squared deviations from them (the stored entries' + in GraphBLAS, and the fill's once per fiber, times its count of unstored positions). + """ + from ._elementwise import add, binary_numpy, cast, divide, subtract + + _check_array(x, "var") + rdt = _mean_dtype(x, "var") + real_dt = np.dtype(rdt.char.lower()) if rdt.kind == "c" else rdt + axes = _axes(x, axis) + L = math.prod(x.shape[a] for a in axes) + if L == 0: + shape = _out_shape(x.shape, axes, keepdims) + return Array._new(Tensor(gb_dtype(real_dt), shape), np.nan) + denominator = builtins.max(L - correction, 0) + if not axes: + xr = cast(x, rdt) + out = divide(_squared_magnitude(subtract(xr, xr)), denominator) + return out if out.dtype == real_dt else cast(out, real_dt) + T = x._T() + m = mean(x, axis=axes) + if T.nvals: + C = _reduce_tensor(T, agg.count, axes) + D = _deviations(T, m, C, axes, rdt) + S2 = _utils.map_values(lambda d: np.real(d * np.conj(d)).astype(real_dt), [D], real_dt) + S2 = _reduce_tensor(S2, monoid.plus[gb_dtype(real_dt)], axes) + else: + S2 = Tensor(gb_dtype(real_dt), m.shape) + C = Tensor(gb_dtype(np.int64), m.shape) + k = subtract(L, Array._new(C, 0)) + dev = subtract(Array._new(_utils.empty_0d(rdt), x._fill), m) + # k unstored positions, each (fill - mean) ** 2 away; none for k == 0 (the fill may be inf) + missing = binary_numpy("var", _times_count, k, _squared_magnitude(dev)) + out = divide(add(Array._new(S2, 0), missing), denominator) + if out.dtype != real_dt: + out = cast(out, real_dt) + return _keepdims(out, x, axes, keepdims) + + +def _deviations(T, m, C, axes, rdt): + """Each entry of ``T`` minus its fiber's mean (``m``, ``C`` the entries per fiber).""" + minus = binary.minus[gb_dtype(rdt)] + if m.ndim == 0: + return T.apply(minus, right=_utils.gb_scalar(m._fill, rdt)).new() + Tm = m._T() + if Tm.nvals != C.nvals: # pragma: no cover (safety: store the mean of every fiber) + keep_shape = _out_shape(T.shape, axes, keepdims=True) + _utils.check_dense(math.prod(keep_shape), "var or std (the means)") + M = _utils.full_tensor(rdt, keep_shape, m._fill) + M(accum=binary.second) << Tm.expand_dims(axes) + return T.ewise_mult(M, minus).new() + # The means store exactly the fibers that hold entries (the pattern of the sums), and every + # entry lies in one of them: broadcast them as they are, with no dense copy + return T.ewise_mult(Tm.expand_dims(axes), minus).new() + + +def _times_count(k, d): + return np.where(k == 0, d.dtype.type(0), k * d) + + +def std(x, /, *, axis=None, correction=0.0, keepdims=False): + from ._elementwise import sqrt + + return sqrt(var(x, axis=axis, correction=correction, keepdims=keepdims)) + + +def _cumulative(name, npfunc, x, axis, dtype, include_initial): + _check_array(x, name) + if x.ndim == 0: + raise ValueError(f"{name} needs an array of at least one dimension") + if axis is None: + if x.ndim > 1: + raise ValueError(f"{name} of an array of more than one dimension needs an axis") + axis = 0 + axis = _axes(x, axis)[0] + if dtype is None: + reduce_name = "sum" if name == "cumulative_sum" else "prod" + ct = check_supported(_MONOIDS[reduce_name].npfunc(np.zeros(1, x.dtype)).dtype, name) + else: + ct = as_dtype(dtype) + with np.errstate(all="ignore"): + fill = np.asarray(x._fill).astype(ct) + identity = 0 if name == "cumulative_sum" else 1 + if x.size and fill == identity and not (include_initial and _negative_zero(fill)): + return _cumulative_sparse(npfunc, x, axis, ct, fill, include_initial) + _utils.check_dense(x.size + (x.size // builtins.max(x.shape[axis], 1)), name) + with np.errstate(all="ignore"): + out = npfunc(x.todense(), axis=axis, dtype=ct) + if include_initial: + shape = list(out.shape) + shape[axis] = 1 + initial = np.zeros(shape, ct) if name == "cumulative_sum" else np.ones(shape, ct) + out = np.concatenate([initial, out], axis=axis) + from ._creation import _from_numpy + + return _from_numpy(out) + + +def _negative_zero(fill): + """Whether a zero ``fill`` has a sign bit (an included initial 0 would then be stored).""" + parts = [fill.real, fill.imag] if np.iscomplexobj(fill) else [fill] + return builtins.any(bool(np.signbit(p)) for p in parts if np.asarray(p).dtype.kind == "f") + + +def _cumulative_sparse(npfunc, x, axis, ct, fill, include_initial): + """A cumulative sum (product) along ``axis`` when the fill is the identity, 0 (1): each + fiber is the fill up to its first entry, so only the positions from there on are stored, + O(sum over fibers of the positions from the first entry on). + + numpy accumulates one position at a time, so each run of unstored positions adds the fill + once more (0.0 + -0.0 is 0.0): each run becomes one element of the fiber's sequence, and + the sequences are accumulated by numpy itself, in order, for numpy's rounding. + """ + from ._manipulation import permute_dims + + T = x._T() + L = x.shape[axis] + others = tuple(a for a in range(x.ndim) if a != axis) + shape = tuple(x.shape[a] for a in others) + out_L = L + 1 if include_initial else L + perm = (*others, axis) + out_shape = (*shape, out_L) + if T.nvals == 0: + out = Array._new(Tensor(gb_dtype(ct), out_shape), fill) + return permute_dims(out, tuple(int(i) for i in np.argsort(perm))) + coords, values = permute_dims(x, perm)._T().to_coo() + with np.errstate(all="ignore"): + values = values.astype(ct) + fiber = np.ravel_multi_index(coords[:-1], shape) if shape else np.zeros(len(values), np.int64) + pos = coords[-1] + first = np.r_[True, fiber[1:] != fiber[:-1]] + last = np.r_[fiber[1:] != fiber[:-1], True] + prev = np.where(first, -1, np.r_[0, pos[:-1]]) + # The sequence per fiber: [fill if a run precedes the entry] + entry, ..., [fill if a run + # follows the last entry] + gap = pos > prev + 1 + tail = last & (pos < L - 1) + n_elem = 1 + gap + tail + seq_index = np.cumsum(n_elem) - n_elem # (where each entry's elements start) + seq = np.full(int(n_elem.sum()), fill, dtype=ct) + entry_at = seq_index + gap + seq[entry_at] = values + starts = np.flatnonzero(first) + acc = _grouped_accumulate(npfunc, seq, seq_index[starts], ct) + # Positions and values: each entry; the run before it (not before the first entry: there + # the values are the fill); the run after the last entry + out_pos, out_val, out_fib = [pos], [acc[entry_at]], [fiber] + inner = gap & ~first + if inner.any(): + lengths = (pos - prev - 1)[inner] + out_pos.append(_runs(prev[inner] + 1, lengths)) + out_val.append(np.repeat(acc[seq_index[inner]], lengths)) + out_fib.append(np.repeat(fiber[inner], lengths)) + if tail.any(): + lengths = (L - 1 - pos)[tail] + out_pos.append(_runs(pos[tail] + 1, lengths)) + out_val.append(np.repeat(acc[entry_at[tail] + 1], lengths)) + out_fib.append(np.repeat(fiber[tail], lengths)) + out_pos = np.concatenate(out_pos) + (1 if include_initial else 0) + _utils.check_dense(len(out_pos), "a cumulative sum or product (the positions after an entry)") + idx = list(np.unravel_index(np.concatenate(out_fib), shape)) if shape else [] + idx.append(out_pos) + R = Tensor.from_coo(idx, np.concatenate(out_val), gb_dtype(ct), shape=out_shape) + return permute_dims(Array._new(R, fill), tuple(int(i) for i in np.argsort(perm))) + + +def _runs(starts, lengths): + """Concatenated ranges ``start, start + 1, ..., start + length - 1``.""" + offsets = np.arange(int(lengths.sum())) - np.repeat(np.cumsum(lengths) - lengths, lengths) + return np.repeat(starts, lengths) + offsets + + +def _grouped_accumulate(npfunc, seq, group_starts, ct): + """The cumsum (cumprod) of each group of ``seq`` by numpy (groups start at ``group_starts``), + each in order, as numpy accumulates one fiber. + """ + bounds = np.r_[group_starts, len(seq)] + lengths = np.diff(bounds) + longest = int(lengths.max()) + out = np.empty_like(seq) + with np.errstate(all="ignore"): + if longest * len(lengths) <= 4 * len(seq) + 1024: + # One 2-d accumulation, the groups padded at their ends (padding comes after) + padded = np.zeros((len(lengths), longest), dtype=ct) + col = np.arange(len(seq)) - np.repeat(group_starts, lengths) + row = np.repeat(np.arange(len(lengths)), lengths) + padded[row, col] = seq + out[:] = npfunc(padded, axis=1, dtype=ct)[row, col] + else: + for a, b in itertools.pairwise(bounds): + out[a:b] = npfunc(seq[a:b], dtype=ct) + return out + + +def cumulative_sum(x, /, *, axis=None, dtype=None, include_initial=False): + return _cumulative("cumulative_sum", np.cumsum, x, axis, dtype, include_initial) + + +def cumulative_prod(x, /, *, axis=None, dtype=None, include_initial=False): + return _cumulative("cumulative_prod", np.cumprod, x, axis, dtype, include_initial) + + +def diff(x, /, *, axis=-1, n=1, prepend=None, append=None): + from ._elementwise import not_equal, subtract + from ._manipulation import concat + + _check_array(x, "diff") + if x.ndim == 0: + raise ValueError("diff needs an array of at least one dimension") + axis = _axes(x, axis)[0] + n = int(n) + if n < 0: + raise ValueError(f"diff's n must be nonnegative, not {n}") + if n == 0: + return x._copy() # (numpy's diff ignores prepend and append then) + parts = [_diff_part(p, x, axis) for p in (prepend, x, append) if p is not None] + if len(parts) > 1: + x = concat(parts, axis=axis) + # (numpy's diff of booleans is not_equal: subtracting booleans is an error) + op = not_equal if x.dtype.kind == "b" else subtract + before = (slice(None),) * axis + for _ in range(n): + if x.shape[axis] == 0: + break + x = op(x[(*before, slice(1, None))], x[(*before, slice(None, -1))]) + return x + + +def _diff_part(p, x, axis): + """A ``prepend`` or ``append`` of ``diff``: an Array, or a scalar broadcast to ``x``'s + shape with length 1 along ``axis`` (as numpy takes it). + """ + if isinstance(p, Array) and p.ndim: + return p + from ._creation import asarray + from ._manipulation import _broadcast_to + + shape = list(x.shape) + shape[axis] = 1 + return _broadcast_to(asarray(p), tuple(shape)) diff --git a/graphblas/array_api/_utils.py b/graphblas/array_api/_utils.py new file mode 100644 index 000000000..39a4827cb --- /dev/null +++ b/graphblas/array_api/_utils.py @@ -0,0 +1,263 @@ +"""Private helpers shared by the modules of ``graphblas.array_api``.""" + +import math + +import numpy as np + +from .. import binary, config, monoid, unary +from ..core.scalar import Scalar +from ..core.tensor import Tensor +from ..core.tensor_ewise import _broadcast_factor +from ..dtypes import DataType, lookup_dtype + + +class DensifyError(ValueError): + """An operation would allocate a dense array larger than the namespace allows.""" + + +def densify_limit(): + return config.get("array_api_densify_limit", 10**8) + + +def check_dense(n, what): + """Raise DensifyError when ``what`` would hold ``n`` elements densely past the limit.""" + if n > (limit := densify_limit()): + raise DensifyError( + f"{what} would hold {n} elements densely, more than the limit of {limit}. Raise " + "`graphblas.config` option array_api_densify_limit " + "(`with gb.config.set(array_api_densify_limit=...)`) to allow it." + ) + + +def gb_dtype(dtype): + return lookup_dtype(np.dtype(dtype)) + + +def np_dtype(gbdtype): + return np.dtype(gbdtype.np_type) + + +def scalar(value, dtype): + """``value`` as a numpy scalar of ``dtype`` (numpy's cast, warnings silenced).""" + with np.errstate(all="ignore"): + return np.asarray(value).astype(dtype, copy=False)[()] + + +_SCALARS = {} + + +def gb_scalar(value, dtype): + """A GraphBLAS Scalar of the numpy ``dtype`` holding ``value`` cast as numpy casts it. + + The Scalars are only read (as operands), so they are cached by their bytes: a new one + costs about as much as a small element-wise operation. + """ + value = scalar(value, dtype) + key = (value.dtype.str, value.tobytes()) + if (s := _SCALARS.get(key)) is None: + if len(_SCALARS) > 1024: + _SCALARS.clear() + s = _SCALARS[key] = Scalar.from_value(value, gb_dtype(dtype)) + return s + + +def same_value(a, b): + """Whether two scalars are identical as values: NaN matches NaN, -0.0 differs from 0.0.""" + a = np.asarray(a) + b = np.asarray(b) + if a.dtype.kind == "c" or b.dtype.kind == "c": + return same_value(a.real, b.real) and same_value(a.imag, b.imag) + if a.dtype.kind == "f" or b.dtype.kind == "f": + if np.isnan(a) or np.isnan(b): + return bool(np.isnan(a) and np.isnan(b)) + return bool(a == b and np.signbit(a) == np.signbit(b)) + return bool(np.equal(a, b)) + + +def differs(values, fill): + """Boolean array: where ``values`` is not identical to the scalar ``fill`` (invariant 4).""" + kind = values.dtype.kind + if kind == "c": + return differs(values.real, np.real(fill)) | differs(values.imag, np.imag(fill)) + if kind == "f": + if np.isnan(fill): + return ~np.isnan(values) + return (values != fill) | (np.signbit(values) != np.signbit(fill)) + return values != fill + + +def expansion_check(Ta, Tb, shape): + """Before a union of ``Ta`` and ``Tb`` broadcast to ``shape``: raise past the limit. + + The union replicates each broadcast operand's entries over the axes it broadcasts along. + """ + if Ta.shape == Tb.shape: + return + replicated = 0 + for T in (Ta, Tb): + factor = _broadcast_factor(T.shape, shape) + if factor > 1: + replicated += factor * T.nvals + check_dense(replicated, f"broadcasting shapes {Ta.shape} and {Tb.shape} to {shape}") + + +def mult_expansion_check(Ta, Tb, shape): + """Before an intersection (ewise_mult) of ``Ta`` and ``Tb`` broadcast to ``shape``: when + both broadcast, one is replicated (an outer product); raise past the limit. + """ + fa = _broadcast_factor(Ta.shape, shape) + fb = _broadcast_factor(Tb.shape, shape) + if fa > 1 and fb > 1: + replicated = min(fa * Ta.nvals, fb * Tb.nvals) + check_dense(replicated, f"an outer product of shapes {Ta.shape} and {Tb.shape}") + + +def empty_tensor(dtype, shape): + shape = tuple(shape) + return empty_0d(dtype) if not shape else Tensor(gb_dtype(dtype), shape) + + +_EMPTY_0D = {} + + +def empty_0d(dtype): + """An empty 0-d Tensor of ``dtype`` (a numpy dtype or a graphblas one): a 0-d Array + stores nothing, its value is its fill. An O(1) copy-on-write copy of one made once per + dtype: making one costs a Matrix and a format setting, about 6 us, half of a 0-d + operation (SAPS's time-stepping loops make one per operation). + """ + gdt = dtype if isinstance(dtype, DataType) else gb_dtype(dtype) + if (T := _EMPTY_0D.get(gdt)) is None: + T = _EMPTY_0D[gdt] = Tensor(gdt, ()) + return T.dup() + + +def full_tensor(dtype, shape, value): + """A Tensor holding ``value`` at every position (iso and full: O(1) memory).""" + T = Tensor(gb_dtype(dtype), tuple(shape)) + if math.prod(shape) > 0: + T << gb_scalar(value, dtype) + return T + + +def union_pattern(tensors, shape): + """A BOOL Tensor of ``shape`` with an entry wherever one of ``tensors``, broadcast to + ``shape``, has one. + """ + pair = binary.pair[lookup_dtype(bool)] + P = None + for T in tensors: + if P is None: + if T.shape == shape: + P = T.apply(pair, right=True).new() + continue + P = Tensor(lookup_dtype(bool), shape) + expansion_check(P, T, shape) + P = P.ewise_add(T, pair, allow_expansion=True).new() + return P + + +def aligned(arrays, shape): + """Tensors of ``shape`` with one pattern, the union of the arrays' entries (broadcast), + holding each array's values there: its entry, or its fill where it has none. + """ + tensors = [x._T() for x in arrays] + if len(tensors) == 1 and tensors[0].shape == shape: + return tensors + P = union_pattern(tensors, shape) + out = [] + for x, T in zip(arrays, tensors, strict=True): + dt = x.dtype + second = binary.second[gb_dtype(dt)] + V = P.ewise_union(T, second, gb_scalar(0, dt), gb_scalar(x._fill, dt), allow_expansion=True) + out.append(V.new()) + return out + + +def with_fill(x, fill): + """``x`` with the fill ``fill``: where it stores nothing, its old fill is stored (every + position it does not store is written: O(size), under the densify limit). + """ + from ._array import Array + + fill = scalar(fill, x.dtype) + if same_value(fill, x._fill): + return x + check_dense(x.size, "giving an array another fill value") + T = full_tensor(x.dtype, x.shape, x._fill) + Tx = x._T() + if Tx.nvals: + T(accum=binary.second[gb_dtype(x.dtype)]) << Tx + return Array._new(T, fill) + + +# Facts about an Array's stored values (not its fill), each found by one parallel GraphBLAS +# pass and cached on the Array; ``fact`` documents them. For each: the unary op that says +# whether a value has the property, the monoid that combines those answers, and the answer when +# the Array stores nothing. +# (A number in place of the unary op: whether a value equals it.) +_FACT_SCANS = { + "finite": (unary.isfinite, monoid.land, True), + "nan": (unary.isnan, monoid.lor, False), + "zero": (0, monoid.lor, False), + "posinf": (np.inf, monoid.lor, False), + "neginf": (-np.inf, monoid.lor, False), +} + + +def fact(x, name): + """A fact about the stored values of the Array ``x``, computed once (cached on ``x``): + "finite" (every value finite), "nan" (some value is NaN), "zero" (some value is 0 or + -0.0), "posinf" or "neginf" (some value is inf, or -inf). Integers and bools are finite, + with no NaN or infinity, without a look; each other fact is one + parallel GraphBLAS pass with no copy of the values (an apply to bool, then a reduction + that stops at the first deciding value). + """ + if name not in _FACT_SCANS: # pragma: no cover (safety) + raise KeyError(name) + return x._fact(name, lambda x: _compute_fact(x, name)) + + +def _compute_fact(x, name): + facts = x._facts + if name != "zero" and x.dtype.kind not in "fc": + return name == "finite" + if name in {"nan", "posinf", "neginf"} and facts.get("finite"): + return False + if name == "finite" and facts.get("nan"): + return False + return tensor_fact(x._T(), name) + + +def tensor_fact(T, name): + """``fact`` of the stored values of the Tensor ``T``, by one pass, not cached.""" + op, combine, empty = _FACT_SCANS[name] + if T.nvals == 0: + return empty + dt = T.dtype + if not callable(op): + B = T.apply(binary.eq[dt], right=gb_scalar(op, np_dtype(dt))).new() + else: + B = T.apply(op[dt]).new() + return bool(B.reduce_scalar(combine[lookup_dtype(bool)]).new().value) + + +def keep_facts(src, dst, names=("finite", "nan", "zero")): + """Give ``dst`` the facts ``src`` already knows among ``names`` (for results whose stored + values keep them: ``abs`` and ``negative`` keep all three, a rounding keeps finiteness). + """ + if facts := src._facts: + kept = {k: v for k, v in facts.items() if k in names} + if kept: + if dst._facts is None: + dst._facts = {} + dst._facts.update(kept) + return dst + + +def map_values(func, tensors, dtype): + """A Tensor with the entries of ``tensors`` (all of one shape and pattern) whose values + are ``func(*value_arrays)`` in ``dtype``, computed by numpy over the stored values + (``Tensor._map_values``: the values in place, the pattern copied in GraphBLAS). + """ + return tensors[0]._map_values(func, tensors[1:], gb_dtype(dtype)) diff --git a/graphblas/array_api/linalg.py b/graphblas/array_api/linalg.py new file mode 100644 index 000000000..3f7634209 --- /dev/null +++ b/graphblas/array_api/linalg.py @@ -0,0 +1,342 @@ +"""The Array API's ``linalg`` extension (``xp.linalg``). + +matmul, matrix_transpose, tensordot, vecdot, diagonal, trace, vector_norm, outer, cross, the +Frobenius, 1 and inf matrix norms, and matrix_power with n >= 0 work on the stored entries. +The decompositions and solvers (cholesky, det, eig, eigh, eigvals, eigvalsh, inv, pinv, qr, +slogdet, solve, svd, svdvals, matrix_rank, and matrix_norm of order 2, -2 or "nuc") are dense +algebra: they go through numpy densely, under the densify limit. +""" + +import math +from typing import NamedTuple + +import numpy as np + +from ..core.tensor import Tensor +from . import _utils +from ._array import Array +from ._linalg import matmul, matrix_transpose, tensordot, vecdot + +__all__ = [ + "cholesky", + "cross", + "det", + "diagonal", + "eig", + "eigh", + "eigvals", + "eigvalsh", + "inv", + "matmul", + "matrix_norm", + "matrix_power", + "matrix_rank", + "matrix_transpose", + "outer", + "pinv", + "qr", + "slogdet", + "solve", + "svd", + "svdvals", + "tensordot", + "trace", + "vecdot", + "vector_norm", +] + + +class EighResult(NamedTuple): + eigenvalues: Array + eigenvectors: Array + + +class EigResult(NamedTuple): + eigenvalues: Array + eigenvectors: Array + + +class QRResult(NamedTuple): + Q: Array + R: Array + + +class SlogdetResult(NamedTuple): + sign: Array + logabsdet: Array + + +class SVDResult(NamedTuple): + U: Array + S: Array + Vh: Array + + +def _check(x, name, ndim=2): + if not isinstance(x, Array): + raise TypeError(f"{name} needs an Array, not {type(x).__name__}") + if x.ndim < ndim: + raise ValueError(f"{name} needs an array of at least {ndim} dimensions") + + +def _float(x): + """``x`` in a floating-point dtype (numpy's linalg computes integers in float64).""" + from ._elementwise import cast + + return x if x.dtype.kind in "fc" else cast(x, np.dtype(np.float64)) + + +def _dense_call(name, func, *arrays): + from ._creation import _from_numpy + + for x in arrays: + _utils.check_dense(x.size, f"linalg.{name} (dense algebra)") + out = func(*[x.todense() for x in arrays]) + if isinstance(out, tuple): + return tuple(_from_numpy(np.asarray(o)) for o in out) + return _from_numpy(np.asarray(out)) + + +# Sparse +def diagonal(x, /, *, offset=0): + """The diagonals of the last two axes: O(entries).""" + _check(x, "diagonal") + n, m = x.shape[-2:] + offset = int(offset) + length = max(0, min(n, m - offset) if offset >= 0 else min(n + offset, m)) + shape = (*x.shape[:-2], length) + T = x._T() + if T.nvals == 0 or length == 0: + return Array._new(Tensor(T.dtype, shape), x._fill) + coords, values = T.to_coo(sort=False) + i, j = coords[-2], coords[-1] + keep = j - i == offset + k = i if offset >= 0 else j + out = [c[keep] for c in coords[:-2]] + [k[keep]] + if not keep.any(): + return Array._new(Tensor(T.dtype, shape), x._fill) + return Array._new(Tensor.from_coo(out, values[keep], T.dtype, shape=shape), x._fill) + + +def trace(x, /, *, offset=0, dtype=None): + from ._statistical import sum as sum_ + + return sum_(diagonal(x, offset=offset), axis=-1, dtype=dtype) + + +def outer(x1, x2, /): + from ._elementwise import multiply + from ._manipulation import expand_dims + + for x in (x1, x2): + _check(x, "outer", ndim=1) + if x.ndim != 1: + raise ValueError("outer needs 1-d arrays") + return multiply(expand_dims(x1, axis=1), expand_dims(x2, axis=0)) + + +def cross(x1, x2, /, *, axis=-1): + """The cross product of 3-vectors along ``axis`` (numpy's products and differences).""" + from ._elementwise import multiply, subtract + from ._manipulation import stack + + for x in (x1, x2): + _check(x, "cross", ndim=1) + if not -min(x1.ndim, x2.ndim) <= axis < 0: + raise ValueError(f"cross's axis must be negative and within both arrays, not {axis}") + if x1.shape[axis] != 3 or x2.shape[axis] != 3: + raise ValueError("cross needs vectors of size 3 along the axis") + np.broadcast_shapes(x1.shape, x2.shape) + + def part(x, i): + return x[(Ellipsis, i) + (slice(None),) * (-1 - axis)] + + a = [part(x1, i) for i in range(3)] + b = [part(x2, i) for i in range(3)] + c = [ + subtract(multiply(a[1], b[2]), multiply(a[2], b[1])), + subtract(multiply(a[2], b[0]), multiply(a[0], b[2])), + subtract(multiply(a[0], b[1]), multiply(a[1], b[0])), + ] + return stack(c, axis=axis) + + +def vector_norm(x, /, *, axis=None, keepdims=False, ord=2): # noqa: A002 (the Array API's name) + """O(entries): the norms are reductions of ``abs(x)``.""" + from ._elementwise import abs as abs_ + from ._elementwise import pow as pow_ + from ._elementwise import sqrt, square + from ._searching import count_nonzero + from ._statistical import max as max_ + from ._statistical import min as min_ + from ._statistical import sum as sum_ + + if not isinstance(x, Array): + raise TypeError(f"vector_norm needs an Array, not {type(x).__name__}") + a = abs_(_float(x)) + if axis is None and x.ndim: + axis = tuple(range(x.ndim)) + if ord == 2: + return sqrt(sum_(square(a), axis=axis, keepdims=keepdims)) + if ord == 1: + return sum_(a, axis=axis, keepdims=keepdims) + if ord == math.inf: + return max_(a, axis=axis, keepdims=keepdims) + if ord == -math.inf: + return min_(a, axis=axis, keepdims=keepdims) + if ord == 0: + from ._elementwise import cast + + return cast(count_nonzero(a, axis=axis, keepdims=keepdims), a.dtype) + p = float(ord) + with np.errstate(divide="ignore"): + inv = 1.0 / p + return pow_(sum_(pow_(a, p), axis=axis, keepdims=keepdims), inv) + + +def matrix_norm(x, /, *, keepdims=False, ord="fro"): # noqa: A002 (the Array API's name) + from ._elementwise import abs as abs_ + from ._elementwise import sqrt, square + from ._manipulation import expand_dims + from ._statistical import max as max_ + from ._statistical import min as min_ + from ._statistical import sum as sum_ + + _check(x, "matrix_norm") + a = abs_(_float(x)) + if ord == "fro": + out = sqrt(sum_(square(a), axis=(-2, -1))) + elif ord in (1, -1): + cols = sum_(a, axis=-2) + out = max_(cols, axis=-1) if ord == 1 else min_(cols, axis=-1) + elif ord in (math.inf, -math.inf): + rows = sum_(a, axis=-1) + out = max_(rows, axis=-1) if ord == math.inf else min_(rows, axis=-1) + elif ord in (2, -2, "nuc"): + s = svdvals(x) + if ord == 2: + out = max_(s, axis=-1) + elif ord == -2: + out = min_(s, axis=-1) + else: + out = sum_(s, axis=-1) + else: + raise ValueError(f"invalid matrix_norm order {ord!r}") + if keepdims: + out = expand_dims(expand_dims(out, axis=-1), axis=-1) + return out + + +def matrix_power(x, n, /): + """``x`` to the integer power ``n``: repeated sparse products for n >= 0; inv for n < 0.""" + _check(x, "matrix_power") + if x.shape[-1] != x.shape[-2]: + raise ValueError("matrix_power needs square matrices") + n = int(n) + if n < 0: + return matrix_power(inv(x), -n) + if n == 0: + from ._creation import eye + from ._manipulation import broadcast_to + + return broadcast_to(eye(x.shape[-1], dtype=x.dtype), x.shape) + result = None + base = x + while n: + if n & 1: + result = base if result is None else matmul(result, base) + n >>= 1 + if n: + base = matmul(base, base) + return result + + +# Dense algebra (numpy under the densify limit) +def cholesky(x, /, *, upper=False): + _check(x, "cholesky") + + def func(a): + L = np.linalg.cholesky(a) + return np.conj(np.swapaxes(L, -1, -2)) if upper else L + + return _dense_call("cholesky", func, _float(x)) + + +def det(x, /): + _check(x, "det") + return _dense_call("det", np.linalg.det, _float(x)) + + +def eigh(x, /): + _check(x, "eigh") + return EighResult(*_dense_call("eigh", lambda a: tuple(np.linalg.eigh(a)), _float(x))) + + +def eigvalsh(x, /): + _check(x, "eigvalsh") + return _dense_call("eigvalsh", np.linalg.eigvalsh, _float(x)) + + +def eig(x, /): + _check(x, "eig") + return EigResult(*_dense_call("eig", lambda a: tuple(np.linalg.eig(a)), _float(x))) + + +def eigvals(x, /): + _check(x, "eigvals") + return _dense_call("eigvals", np.linalg.eigvals, _float(x)) + + +def inv(x, /): + _check(x, "inv") + return _dense_call("inv", np.linalg.inv, _float(x)) + + +def pinv(x, /, *, rtol=None): + _check(x, "pinv") + if isinstance(rtol, Array): + rtol = rtol.todense() + return _dense_call("pinv", lambda a: np.linalg.pinv(a, rtol=rtol), _float(x)) + + +def matrix_rank(x, /, *, rtol=None): + _check(x, "matrix_rank") + if isinstance(rtol, Array): + rtol = rtol.todense() + return _dense_call("matrix_rank", lambda a: np.linalg.matrix_rank(a, rtol=rtol), _float(x)) + + +def qr(x, /, *, mode="reduced"): + _check(x, "qr") + return QRResult(*_dense_call("qr", lambda a: tuple(np.linalg.qr(a, mode=mode)), _float(x))) + + +def slogdet(x, /): + _check(x, "slogdet") + return SlogdetResult(*_dense_call("slogdet", lambda a: tuple(np.linalg.slogdet(a)), _float(x))) + + +def solve(x1, x2, /): + _check(x1, "solve") + _check(x2, "solve", ndim=1) + a, b = _float(x1), _float(x2) + + def func(a, b): + if b.ndim == 1: # (the standard's rule: a vector only when x2 is 1-d) + return np.linalg.solve(a, b[..., None])[..., 0] + return np.linalg.solve(a, b) + + return _dense_call("solve", func, a, b) + + +def svd(x, /, *, full_matrices=True): + _check(x, "svd") + out = _dense_call( + "svd", lambda a: tuple(np.linalg.svd(a, full_matrices=full_matrices)), _float(x) + ) + return SVDResult(*out) + + +def svdvals(x, /): + _check(x, "svdvals") + return _dense_call("svdvals", lambda a: np.linalg.svd(a, compute_uv=False), _float(x)) diff --git a/graphblas/core/tensor.py b/graphblas/core/tensor.py index 399ebc82c..ef43e6862 100644 --- a/graphblas/core/tensor.py +++ b/graphblas/core/tensor.py @@ -28,6 +28,7 @@ (77M entries), Uber (3.3M), NIPS (3.1M) and Enron (54M) are FROSTT tensors (frostt.io). """ +import contextlib import itertools import math import operator @@ -803,7 +804,10 @@ def dup(self, dtype=None, *, layout=None, clear=False, name=None): if layout is not None: M, lay = tl.relayout(M, self._shape, lay, tl.as_layout(layout, self._shape)) if dtype != self.dtype: - M = M.dup(dtype) + # (Matrix.dup would hold an n x 1 Matrix by column, as SuiteSparse holds a Vector) + D = tl.new_matrix(dtype, M._nrows, M._ncols) + D << M + M = D if M is self._M: T = Tensor._from_M(M, self._shape, lay, name=name, store=self._store) else: @@ -1116,6 +1120,81 @@ def apply(self, op, right=None, *, left=None): expr_repr="{0.name}.apply({op})", ) + def _map_values(self, func, others=(), dtype=None): + """A Tensor with this Tensor's entries whose values are ``func(*value_arrays)``: the + stored values of this Tensor and of ``others`` (Tensors of the same shape and pattern), + as numpy arrays in one order, in the result's ``dtype`` (default: this Tensor's). + + Private, for ``graphblas.array_api`` (numpy over the values of entries GraphBLAS has + no operator for). The pattern is copied in GraphBLAS and the values written into the + copy's own array through a GxB_Container, so the entries are never rebuilt from + coordinates; the inputs' arrays are views, not copies. + """ + tensors = [self, *others] + gbdt = self.dtype if dtype is None else lookup_dtype(dtype) + shape, layout = self._shape, self._layout + Ms = [T._M for T in tensors] + if self._M._nvals == 0: + return Tensor._from_M(tl.new_matrix(gbdt, Ms[0]._nrows, Ms[0]._ncols), shape, layout) + if any(T._layout != layout for T in others): + Ms = [Ms[0]] + [tl.relayout(T._M, shape, T._layout, layout)[0] for T in others] + # An iso Matrix has one value; the template for the result is one that holds them all, + # if any does, so that its value array has the length of func's output. + isos = [M.ss.is_iso for M in Ms] + template = Ms[isos.index(False)] if False in isos else Ms[0] + W = tl.new_matrix(gbdt, template._nrows, template._ncols) + W << template + with contextlib.ExitStack() as stack: + # (tensors sharing one Matrix, copy-on-write, unload it once) + unique = {id(M): M for M in Ms} + conts = {key: stack.enter_context(tl._Unloaded(M)) for key, M in unique.items()} + wcont = stack.enter_context(tl._Unloaded(W)) + formats = {c.format for c in conts.values()} | {wcont.format} + nvals = {c.nvals for c in conts.values()} | {wcont.nvals} + if len(formats) == 1 and len(nvals) == 1: + values = { + key: _value_array(conts[key], M.dtype, stack) for key, M in unique.items() + } + arrays = [values[id(M)] for M in Ms] + out = _value_array(wcont, gbdt, stack) + with np.errstate(all="ignore"): + result = np.asarray(func(*arrays)) + if out is not None and result.shape == out.shape: + out[...] = result + return Tensor._from_M(W, shape, layout) + # (formats differ: a fallback through coordinates, which aligns the values by position) + return _map_values_coo(func, tensors, gbdt) # pragma: no cover (safety) + + @classmethod + def _concat(cls, tensors, axis): + """``tensors`` (of one dtype, and of one shape but along ``axis``) joined along + ``axis``, by one GxB_Matrix_concat of their stored Matrices: row blocks when ``axis`` + leads the row group of the first one's layout, column blocks when it leads the column + group. Otherwise every part moves to a layout with ``axis`` leading the rows first. + The others take the first one's layout. Private, for ``graphblas.array_api``. + """ + T0 = tensors[0] + shape = list(T0._shape) + shape[axis] = sum(T._shape[axis] for T in tensors) + shape = tuple(shape) + rows, cols = T0._layout + if rows[:1] == (axis,) or cols[:1] == (axis,): + layout = T0._layout + else: + layout = Layout( + (axis, *(a for a in rows if a != axis)), tuple(a for a in cols if a != axis) + ) + by_rows = layout.rows[:1] == (axis,) + Ms = [ + T._M if T._layout == layout else tl.relayout(T._M, T._shape, T._layout, layout)[0] + for T in tensors + if T._shape[axis] # (a part of size 0 along the axis adds nothing) + ] + R = tl.new_matrix(T0.dtype, tl.prod(shape, layout.rows), tl.prod(shape, layout.cols)) + if Ms: + R.ss.concat([[M] for M in Ms] if by_rows else [Ms]) + return cls._from_M(R, shape, layout) + # Reductions def reduce(self, op=monoid.plus, axis=None, *, keepdims=False): """Reduce over ``axis`` (None: all axes, giving a 0-d Tensor) with the monoid ``op``, @@ -1251,6 +1330,26 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * ) +def _value_array(cont, dtype, stack): + """The value array of an unloaded Container as a numpy view (None when it is empty).""" + return stack.enter_context(tl._VectorArray(cont.x, np.dtype(dtype.np_type))).array + + +def _map_values_coo(func, tensors, dtype): + coords = None + arrays = [] + for T in tensors: + idx, vals = T.to_coo() + if coords is None: + coords = idx + arrays.append(vals) + with np.errstate(all="ignore"): + result = np.asarray(func(*arrays), dtype=dtype.np_type) + if result.ndim == 0: + result = np.broadcast_to(result, coords[0].shape if coords else (1,)) + return Tensor.from_coo(coords, result, dtype, shape=tensors[0]._shape) + + def _ewise_infix(left, right, op, kind): ltensor = _as_tensor(left, "infix operator") rtensor = _as_tensor(right, "infix operator") diff --git a/graphblas/core/tensor_contract.py b/graphblas/core/tensor_contract.py index b10ef17ba..38100e5af 100644 --- a/graphblas/core/tensor_contract.py +++ b/graphblas/core/tensor_contract.py @@ -137,8 +137,19 @@ def _contract(updater, a, b, op, a_axes, b_axes, a_batch=(), b_batch=()): if best is None: raise ValueError("the contracted or free axes do not fit in a Matrix dimension (2**60)") _, sa, sb, beta = best - A = _in_layout(a, Layout.from_state(sa), opt_nthreads) - B = _in_layout(b, Layout.from_state(sb), opt_nthreads) + A = B = None + if not nb and updater.kwargs.get("mask") is not None: + # Under a mask, an operand stored as the transpose of what the product needs goes in + # through the transpose descriptor, not a transposed copy: SuiteSparse can then take + # its masked dot product (C = A*B' with both held by row: an SDDMM). Given the + # copy, it ran a masked saxpy over the whole product (n = 2e4, k = 64, a mask of 2e5 + # entries: 2.0 s against 0.9 ms) + A = _transposed(a, sa) + B = _transposed(b, sb) + if A is None: + A = _in_layout(a, Layout.from_state(sa), opt_nthreads) + if B is None: + B = _in_layout(b, Layout.from_state(sb), opt_nthreads) out_batch = {x: a_batch.index(x) for x in a_batch} out_a = {x: nb + i for i, x in enumerate(free_a)} out_b = {y: nb + len(free_a) + i for i, y in enumerate(free_b)} @@ -148,7 +159,8 @@ def _contract(updater, a, b, op, a_axes, b_axes, a_batch=(), b_batch=()): layout = Layout(rows, tuple(out_b[y] for y in sb[0][nb + k :])) if census.on: census.note("contract", "batched, block-diagonal" if nb else "one mxm", - moved=[A is not a._M, B is not b._M], cost=best[0]) # fmt: skip + moved=[getattr(A, "_matrix", A) is not a._M, + getattr(B, "_matrix", B) is not b._M], cost=best[0]) # fmt: skip if not nb: return A.mxm(B, op), layout, shape nblocks = tl.prod(a._shape, a_batch) @@ -172,6 +184,17 @@ def _contract(updater, a, b, op, a_axes, b_axes, a_batch=(), b_batch=()): return A.mxm(B, op), layout, shape +def _transposed(t, state): + """``t``'s Matrix transposed by descriptor (no copy) when that puts it in the layout + ``state``, else None. + """ + rows, cols = t._layout + seq, p = state + if rows and cols and seq == cols + rows and p == len(cols): + return t._M.T + return None + + def _matmul_expr(self, other, op): if (tensor := _as_tensor(other, "matmul")) is None: self._expect_type(other, Tensor, within="matmul", argname="other", op=op) diff --git a/graphblas/core/tensor_ewise.py b/graphblas/core/tensor_ewise.py index 6102f1dd6..c40dd7e09 100644 --- a/graphblas/core/tensor_ewise.py +++ b/graphblas/core/tensor_ewise.py @@ -190,7 +190,10 @@ def _bcast_recipe(updater, method, a, b, op, untyped_op, extra, shape): # result, so "covers the result" must also mean "has its axes". a = _lift_rank(a, len(shape)) b = _lift_rank(b, len(shape)) - if method == "ewise_mult": + if method == "ewise_mult" and math.prod(shape) == 0: + # (no positions: the diagonal scaling below would multiply mismatched shapes) + _ewise_recipe(updater, method, _expand(a, shape), _expand(b, shape), op, extra) + elif method == "ewise_mult": _bcast_mult(updater, a, b, op, untyped_op, shape) else: if census.on: diff --git a/graphblas/core/tensor_reduce.py b/graphblas/core/tensor_reduce.py index 1129868f1..83351f2de 100644 --- a/graphblas/core/tensor_reduce.py +++ b/graphblas/core/tensor_reduce.py @@ -276,13 +276,18 @@ def _reduce_matrix(M, shape, layout, A, op, opts): M = M.ss.reshape(tl.prod(shape, seq[:q]), tl.prod(shape, seq[q:])) layout = Layout(seq[:q], seq[q:]) continue + # Into many outputs, a sum on 2 threads is as fast as on 1 or faster; on 4 and 8 it + # can be 2-4x slower (FROSTT Uber, 4392 x 1.96e6: 4.0 / 4.0 / 17 / 15 ms at 1 / 2 / + # 4 / 8 threads; 1e4 x 1e6 with 1e7 entries: 18 / 18 / 11 / 39 ms). MAX and MIN + # scale instead (5e6 x 5e6 with 1e7 entries: 122 / 120 / 51 / 50 ms; 1e7 x 1e7 + # with 1e8: 0.66 / 0.66 / 0.32 / 0.28 s; no loss on the others), so they keep + # every thread. (SuiteSparse 10.5.1, arm64.) + capped = wide and getattr(op, "name", None) not in {"max", "min"} if census.on: - census.note("reduce", "row group: reduce_columnwise" + (", 2 threads" if wide + census.note("reduce", "row group: reduce_columnwise" + (", 2 threads" if capped and nthreads > 2 else ""), nvals=nvals) # fmt: skip v = Vector(dtype, M._ncols) - # Into many outputs, 2 threads scatter as fast as 1 or faster; 4 and 8 can be - # 3-4x slower (Uber 4.0 / 4.9 / 17 / 15 ms at 2 / 1 / 4 / 8 threads) - v(**(dict(opts, nthreads=min(2, nthreads)) if wide else opts)) << M.reduce_columnwise( + v(**(dict(opts, nthreads=min(2, nthreads)) if capped else opts)) << M.reduce_columnwise( op ) M = _vector_to_row(v) diff --git a/graphblas/graphblas.yaml b/graphblas/graphblas.yaml index f3885d73a..4711f66b2 100644 --- a/graphblas/graphblas.yaml +++ b/graphblas/graphblas.yaml @@ -8,3 +8,5 @@ mapnumpy: True jit_c_control: null # Tensor: largest number of entries an implicit broadcast may replicate (ewise_add / union) tensor_max_expansion: 1000000000 +# graphblas.array_api: most elements an array may hold densely (np.asarray, cumulative_sum, ...) +array_api_densify_limit: 100000000 diff --git a/graphblas/tests/test_array_api_creation.py b/graphblas/tests/test_array_api_creation.py new file mode 100644 index 000000000..48778f4eb --- /dev/null +++ b/graphblas/tests/test_array_api_creation.py @@ -0,0 +1,613 @@ +"""graphblas.array_api: creation, data types, the Array object, basic indexing and the shape +functions, against numpy. +""" + +import copy +import math +import pickle + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp +from graphblas import Matrix, Scalar, Tensor, Vector + +from .test_array_api_elementwise import ( + DTYPES, + assert_array_equal, + check_invariants, + inputs, + make, + np_call, +) + + +def test_namespace(): + assert xp.__array_api_version__ == "2025.12" + for name in xp.__all__: + assert hasattr(xp, name), name + assert xp.float64 is np.dtype("float64") + assert xp.bool == np.dtype(bool) + x = xp.asarray([1, 2]) + assert x.__array_namespace__() is xp + assert x.__array_namespace__(api_version="2024.12") is xp + with pytest.raises(ValueError, match="version"): + x.__array_namespace__(api_version="2019.12") + + +# asarray +def test_asarray_stores_what_differs_from_the_fill(): + x = xp.asarray([1.0, 0.0, -0.0, np.nan, 2.0]) + assert x.fill_value == 0.0 + assert x.nnz == 4 # -0.0 and NaN are stored + assert_array_equal(x, np.array([1.0, 0.0, -0.0, np.nan, 2.0])) + y = xp.asarray([1.0, np.nan, np.nan], fill_value=np.nan) + assert y.nnz == 1 + assert_array_equal(y, np.array([1.0, np.nan, np.nan])) + z = xp.asarray([[1.0, 2.0], [-0.0, 0.0]], fill_value=-0.0) + assert z.nnz == 3 + assert_array_equal(z, np.array([[1.0, 2.0], [-0.0, 0.0]])) + c = xp.asarray([complex(0, -0.0), 0j, complex(np.nan, 0)]) + assert c.nnz == 2 + assert_array_equal(c, np.array([complex(0, -0.0), 0j, complex(np.nan, 0)])) + + +@pytest.mark.parametrize( + ("obj", "dtype"), + [ + (True, "bool"), + (3, "int64"), + (2.5, "float64"), + (1j, "complex128"), + ([1, 2], "int64"), + ([1.0, 2], "float64"), + ([True, False], "bool"), + ([[1, 0], [0, 3]], "int64"), + ([], "float64"), + (np.float32(1.5), "float32"), + (np.arange(4, dtype=np.uint16), "uint16"), + ], +) +def test_asarray_dtypes(obj, dtype): + x = xp.asarray(obj) + assert x.dtype == np.dtype(dtype) + assert_array_equal(x, np.asarray(obj, dtype=dtype)) + + +def test_asarray_dtype_and_copy(): + x = xp.asarray([1, 2, 3], dtype=xp.float32) + assert x.dtype == xp.float32 + assert xp.asarray(x) is x + assert xp.asarray(x, copy=False) is x + y = xp.asarray(x, copy=True) + assert y is not x + y += 1 + assert_array_equal(x, np.array([1, 2, 3], np.float32)) + assert_array_equal(xp.asarray(x, dtype=xp.int8), np.array([1, 2, 3], np.int8)) + with pytest.raises(ValueError, match="copy"): + xp.asarray(x, dtype=xp.int8, copy=False) + with pytest.raises(ValueError, match="copy"): + xp.asarray([1, 2], copy=False) + with pytest.raises(TypeError): + xp.asarray([1, 2], dtype=np.float16) + with pytest.raises(TypeError): + xp.asarray(np.array(["a"])) + with pytest.raises(OverflowError): + xp.asarray([300], dtype=xp.int8) + with pytest.raises(ValueError, match="device"): + xp.asarray([1], device="gpu") + + +def test_asarray_nested_arrays(): + x = xp.asarray([xp.asarray(1.0), xp.asarray(2.0)]) + assert_array_equal(x, np.array([1.0, 2.0])) + x = xp.asarray([[1, 2], xp.asarray([3, 4])]) + assert_array_equal(x, np.array([[1, 2], [3, 4]])) + + +def test_asarray_graphblas_objects(): + v = Vector.from_coo([0, 3], [1.5, 2.5], size=5) + x = xp.asarray(v) + assert x.nnz == 2 + assert_array_equal(x, np.array([1.5, 0, 0, 2.5, 0])) + x = xp.asarray(v, fill_value=np.inf) + assert_array_equal(x, np.array([1.5, np.inf, np.inf, 2.5, np.inf])) + A = Matrix.from_coo([0, 1], [1, 0], [5, 6], nrows=2, ncols=3) + assert_array_equal(xp.asarray(A), np.array([[0, 5, 0], [6, 0, 0]])) + assert_array_equal(xp.asarray(A.T), np.array([[0, 6], [5, 0], [0, 0]])) + assert_array_equal( + xp.asarray(A, dtype=xp.float32), np.array([[0, 5, 0], [6, 0, 0]], np.float32) + ) + T = Tensor.from_coo([[0, 1], [1, 0], [2, 3]], [1.0, 2.0], shape=(2, 2, 4)) + x = xp.asarray(T, fill_value=-1.0) + expected = np.full((2, 2, 4), -1.0) + expected[0, 1, 2], expected[1, 0, 3] = 1.0, 2.0 + assert_array_equal(x, expected) + T[0, 0, 0] = 9.0 # the Array has its own data + assert_array_equal(x, expected) + assert_array_equal(xp.asarray(Scalar.from_value(3)), np.array(3)) + assert_array_equal(xp.asarray(Scalar(float)), np.array(0.0)) + assert isinstance(x.tensor, Tensor) + assert x.tensor.nvals == 2 + + +def test_asarray_scipy_and_pydata_sparse(): + sp = pytest.importorskip("scipy.sparse") + A = sp.csr_array(np.array([[0.0, 1.5], [2.0, 0.0]])) + assert_array_equal(xp.asarray(A), np.array([[0.0, 1.5], [2.0, 0.0]])) + sparse = pytest.importorskip("sparse") + s = sparse.COO.from_numpy(np.array([[1.0, 2.0], [2.0, 3.0]]), fill_value=2.0) + x = xp.asarray(s) + assert x.fill_value == 2.0 + assert x.nnz == 2 + assert_array_equal(x, np.array([[1.0, 2.0], [2.0, 3.0]])) + + +def test_from_coo(): + x = xp.from_coo(([0, 2], [1, 0]), [1.0, 2.0], (3, 2), fill_value=5.0) + assert_array_equal(x, np.array([[5.0, 1.0], [5.0, 5.0], [2.0, 5.0]])) + # a repeated coordinate sums its values (or for bool), as pydata/sparse and scipy do + x = xp.from_coo(([0, 0, 1], [1, 1, 0]), [1.0, 2.0, 5.0], (2, 2)) + assert_array_equal(x, np.array([[0.0, 3.0], [5.0, 0.0]])) + x = xp.from_coo(([0, 0],), [True, False], (3,)) + assert_array_equal(x, np.array([True, False, False])) + x = xp.from_coo(([1, 1],), np.array([100, 100], np.int8), (2,)) + assert_array_equal(x, np.array([0, -56], np.int8)) # (wraps, as numpy's int8 sum) + # one value for every coordinate: stored once (iso), summed where coordinates repeat + x = xp.from_coo(([0, 2, 0],), 7, (3,)) + assert_array_equal(x, np.array([14, 0, 7])) + assert_array_equal(xp.from_coo(([0, 0],), True, (2,)), np.array([True, False])) + assert_array_equal(xp.from_coo(([1, 0],), 2.5, (2,)), np.array([2.5, 2.5])) + x = xp.from_coo(np.array([[0, 2], [1, 0]]), 7, (3, 2)) + assert x.dtype == np.dtype(int) + assert_array_equal(x, np.array([[0, 7], [0, 0], [7, 0]])) + x = xp.from_coo(([],), [], (4,), dtype=xp.int8) + assert_array_equal(x, np.zeros(4, np.int8)) + coords, values = xp.from_coo(([3, 1],), [1.0, 2.0], (5,)).to_coo() + np.testing.assert_array_equal(coords[0], [1, 3]) + np.testing.assert_array_equal(values, [2.0, 1.0]) + + +# Creation functions +@pytest.mark.parametrize("shape", [(), 3, (0,), (2, 3), (2, 0, 4)]) +@pytest.mark.parametrize("dtype", [None, "bool", "int16", "uint8", "float32", "complex64"]) +def test_full_family(shape, dtype): + dt = None if dtype is None else np.dtype(dtype) + for name, value in [("zeros", 0), ("ones", 1), ("empty", 0)]: + x = getattr(xp, name)(shape, dtype=dt) + assert x.nnz == 0 + assert_array_equal(x, getattr(np, name if name != "empty" else "zeros")(shape, dtype=dt)) + like = getattr(xp, f"{name}_like")(x) + assert like.nnz == 0 + assert_array_equal(like, np.full(x.shape, value, dtype=x.dtype)) + if dt is None: + for value in [True, 3, 2.5, 1j, -0.0, np.nan]: + assert_array_equal(xp.full(shape, value), np.full(shape, value)) + else: + assert_array_equal(xp.full(shape, 1, dtype=dt), np.full(shape, 1, dtype=dt)) + assert_array_equal(xp.full_like(xp.zeros(shape), 2, dtype=dt), np.full(shape, 2, dtype=dt)) + with pytest.raises(OverflowError): + xp.full(shape, 1000, dtype=xp.int8) + + +@pytest.mark.parametrize("dtype", [None, "int32", "bool", "complex128"]) +def test_eye(dtype): + for args in [(3,), (3, 5), (5, 3), (0,), (4, 4)]: + for k in [-6, -2, -1, 0, 1, 3, 5]: + assert_array_equal(xp.eye(*args, k=k, dtype=dtype), np.eye(*args, k=k, dtype=dtype)) + + +@pytest.mark.parametrize( + "args", + [(5,), (2, 9), (2, 9, 3), (9, 2, -2), (0,), (1.5, 4.0, 0.5), (10, 0, -3), (0, 1, 0.1)], +) +@pytest.mark.parametrize("dtype", [None, "int16", "float32"]) +def test_arange(args, dtype): + if dtype == "int16" and any(isinstance(a, float) for a in args): + return + expected = np.arange(*args, dtype=dtype) + if dtype is None and all(isinstance(a, int) for a in args): + expected = expected.astype(np.int64) + assert_array_equal(xp.arange(*args, dtype=dtype), expected) + with pytest.raises(ValueError, match="step"): + xp.arange(0, 3, 0) + + +@pytest.mark.parametrize( + ("start", "stop", "num"), [(0, 1, 5), (2.5, -1, 4), (0, 1, 1), (0, 1, 0), (1j, 2, 3)] +) +@pytest.mark.parametrize("endpoint", [True, False]) +def test_linspace(start, stop, num, endpoint): + expected = np.linspace(start, stop, num, endpoint=endpoint) + assert_array_equal(xp.linspace(start, stop, num, endpoint=endpoint), expected) + if not isinstance(start, complex): + assert_array_equal( + xp.linspace(start, stop, num, endpoint=endpoint, dtype=xp.float32), + np.linspace(start, stop, num, endpoint=endpoint, dtype=np.float32), + ) + + +def test_tril_triu(): + rng = np.random.default_rng(0) + for shape in [(4, 4), (3, 5), (2, 3, 4), (0, 3)]: + a = np.where(rng.random(shape) < 0.5, 0.0, rng.normal(size=shape)) + for fill in [0.0, 2.0]: + dense = np.where(a == 0, fill, a) + x = make(dense, fill) + for k in [-3, -1, 0, 1, 2, 5]: + assert_array_equal(xp.tril(x, k=k), np.tril(dense, k=k)) + assert_array_equal(xp.triu(x, k=k), np.triu(dense, k=k)) + with pytest.raises(ValueError, match="2 dimensions"): + xp.tril(xp.zeros(3)) + + +@pytest.mark.parametrize("indexing", ["xy", "ij"]) +def test_meshgrid(indexing): + a, b, c = np.array([1.0, 0.0, 2.0]), np.array([0.0, 5.0]), np.array([7.0, 0.0, 0.0, 1.0]) + for arrays in [(a,), (a, b), (a, b, c)]: + out = xp.meshgrid(*[make(v, 0.0) for v in arrays], indexing=indexing) + expected = np.meshgrid(*arrays, indexing=indexing) + assert len(out) == len(expected) + for o, e in zip(out, expected, strict=True): + assert_array_equal(o, e) + assert xp.meshgrid() == () + + +def test_from_dlpack(): + a = np.array([[1.0, 0.0], [0.0, 2.0]]) + assert_array_equal(xp.from_dlpack(a), a) + x = make(a, 0.0) + assert_array_equal(xp.from_dlpack(x), a) + np.testing.assert_array_equal(np.from_dlpack(x), a) + assert xp.from_dlpack(x, copy=False) is x + assert xp.from_dlpack(x, copy=True) is not x + with pytest.raises(ValueError, match="copy"): + xp.from_dlpack(a, copy=False) + + +# Data type functions +def test_astype_all_pairs(): + for src in DTYPES: + for dense, x in inputs(src, (3, 4), seed=1): + for dst in DTYPES: + # (numpy's astype from complex to real keeps the real part, with a warning) + real = dense.real if src.kind == "c" and dst.kind not in "cb" else dense + expected = np_call(real.astype, dst) + if src.kind == "f" and dst.kind in "iu" and not np.isfinite(dense).all(): + # numpy's cast of NaN and inf to an integer is undefined; any value will do + out = xp.astype(x, dst) + assert out.dtype == dst + finite = np.isfinite(dense) + np.testing.assert_array_equal(out.todense()[finite], expected[finite]) + continue + if src.kind == "f" and dst.kind in "iu": + big = np.abs(dense) >= 2.0 ** (8 * dst.itemsize - 1) + out = xp.astype(x, dst) + np.testing.assert_array_equal(out.todense()[~big], expected[~big]) + continue + assert_array_equal(xp.astype(x, dst), expected) + x = xp.asarray([1, 2]) + assert xp.astype(x, xp.int64, copy=False) is x + assert xp.astype(x, xp.int64) is not x + + +def test_can_cast(): + assert xp.can_cast(xp.int8, xp.int16) + assert not xp.can_cast(xp.int16, xp.int8) + assert xp.can_cast(xp.asarray([1], dtype=xp.float32), xp.float64) + assert not xp.can_cast(xp.float64, xp.float32) + assert xp.can_cast(xp.bool, xp.uint8) + + +def test_finfo_iinfo(): + for dt in [xp.float32, xp.float64, xp.complex64, xp.complex128]: + info = xp.finfo(dt) + ref = np.finfo(dt) + assert info.bits == ref.bits + assert info.eps == float(ref.eps) + assert info.max == float(ref.max) + assert info.min == float(ref.min) + assert info.smallest_normal == float(ref.smallest_normal) + assert isinstance(info.eps, float) + assert xp.finfo(xp.asarray(1, dtype=dt)) == info + assert xp.finfo(xp.complex64).dtype == xp.float32 + for dt in [xp.int8, xp.uint64, xp.int64]: + info = xp.iinfo(dt) + assert (info.bits, info.min, info.max) == ( + np.iinfo(dt).bits, + int(np.iinfo(dt).min), + int(np.iinfo(dt).max), + ) + assert info.dtype == dt + with pytest.raises(TypeError): + xp.finfo(xp.int8) + with pytest.raises(TypeError): + xp.iinfo(xp.float32) + + +def test_isdtype(): + assert xp.isdtype(xp.int8, "signed integer") + assert xp.isdtype(xp.uint8, "integral") + assert not xp.isdtype(xp.bool, "numeric") + assert xp.isdtype(xp.complex64, ("real floating", "complex floating")) + assert xp.isdtype(xp.float32, xp.float32) + assert not xp.isdtype(xp.float32, xp.float64) + with pytest.raises(ValueError, match="kind"): + xp.isdtype(xp.int8, "integer") + with pytest.raises(TypeError): + xp.isdtype("int8", "integral") + + +def test_result_type(): + assert xp.result_type(xp.int8, xp.uint8) == xp.int16 + assert xp.result_type(xp.asarray([1], dtype=xp.float32), 1.5) == xp.float32 + assert xp.result_type(xp.int8, 1) == xp.int8 + assert xp.result_type(xp.float32, xp.int8, xp.uint8) == np.result_type( + np.result_type(np.float32, np.int8), np.uint8 + ) + with pytest.raises(ValueError, match="at least one"): + xp.result_type(1, 2.0) + + +def test_info(): + info = xp.__array_namespace_info__() + assert info.default_device() == "cpu" + assert info.devices() == ("cpu",) + assert info.default_dtypes()["real floating"] == xp.float64 + assert info.default_dtypes(device="cpu")["indexing"] == xp.int64 + assert set(info.dtypes()) == {str(dt) for dt in DTYPES} + assert set(info.dtypes(kind="integral")) == { + "int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64" + } # fmt: skip + assert set(info.dtypes(kind=("bool", "complex floating"))) == { + "bool", + "complex64", + "complex128", + } + assert info.capabilities()["data-dependent shapes"] + + +# The Array object +def test_array_attributes(): + a = np.arange(24.0).reshape(2, 3, 4) + x = make(a, 0.0) + assert x.shape == (2, 3, 4) + assert x.ndim == 3 + assert x.size == 24 + assert x.dtype == xp.float64 + assert x.device == "cpu" + assert x.to_device("cpu") is x + with pytest.raises(ValueError, match="device"): + x.to_device("gpu") + assert_array_equal(x.T, a.T) + assert_array_equal(x.mT, np.swapaxes(a, -1, -2)) + with pytest.raises(ValueError, match="2 dimensions"): + xp.zeros(3).mT + assert x.nnz == 23 + assert x.fill_value == 0.0 + assert repr(x) == "Array(shape=(2, 3, 4), dtype=float64, nnz=23, fill_value=0.0)" + assert repr(xp.asarray(1.5)) == "Array(1.5, dtype=float64)" + with pytest.raises(TypeError, match="asarray"): + xp.Array() + assert len(x) == 2 + assert [y.shape for y in x] == [(3, 4), (3, 4)] + with pytest.raises(TypeError, match="len"): + _ = len(xp.asarray(1)) + with pytest.raises(TypeError): + _ = iter(xp.asarray(1)) + with pytest.raises(TypeError): + hash(x) + + +def test_metadata_is_stored(): + # shape, dtype, ndim and size never read the Tensor (one may not be computed yet) + x = make(np.arange(24, dtype=np.int32).reshape(2, 3, 4), 0) + x._tensor = None # (any read of it would fail) + assert (x.shape, x.dtype, x.ndim, x.size) == ((2, 3, 4), xp.int32, 3, 24) + assert x.fill_value == 0 + y = xp.zeros((5, 0), dtype=xp.bool) + y += xp.ones((5, 0), dtype=xp.bool) + assert (y.shape, y.dtype) == ((5, 0), xp.bool) + + +def test_python_scalars(): + assert bool(xp.asarray(True)) is True + assert int(xp.asarray(-3)) == -3 + assert float(xp.asarray(2.5)) == 2.5 + assert complex(xp.asarray(1 + 2j)) == 1 + 2j + assert [10, 20, 30][xp.asarray(1)] == 20 + assert math.isnan(float(xp.asarray(np.nan))) + with pytest.raises(TypeError): + [1, 2][xp.asarray(1.0)] + with pytest.raises(TypeError): + bool(xp.asarray([1, 2])) + with pytest.raises(TypeError): + float(xp.asarray(1j)) + with pytest.raises(ValueError, match="NaN"): + int(xp.asarray(np.nan)) + + +def test_array_protocol_is_guarded(): + x = make(np.array([1.0, 0.0, 2.0]), 0.0) + np.testing.assert_array_equal(np.asarray(x), [1.0, 0.0, 2.0]) + assert np.asarray(x, dtype=np.float32).dtype == np.float32 + with pytest.raises(ValueError, match="copies"): + np.asarray(x, copy=False) + big = xp.zeros((10**5, 10**5)) + with pytest.raises(xp.DensifyError, match="array_api_densify_limit"): + np.asarray(big) + with gb.config.set(array_api_densify_limit=2), pytest.raises(xp.DensifyError): + x.todense() + assert big.nnz == 0 + assert xp.sum(big) == 0.0 + + +def test_copy_module(): + a = np.array([[1.0, 0.0], [2.0, 3.0]]) + x = make(a, 0.0) + for func in [copy.copy, copy.deepcopy]: + y = func(x) + y[0, 0] = 5.0 + assert_array_equal(x, a) + assert y[0, 0] == 5.0 + z = func(x @ x) + assert_array_equal(z, a @ a) + + +def test_pickle_roundtrips(): + # (review finding L7) through coordinates, values, shape, dtype and fill; explicit fills + # stay stored, and a pending product is computed first + rng = np.random.default_rng(19) + for dense, fill in [ + (np.array([1.0, 0.0, 2.0]), 0.0), + (np.array([[1, 7], [7, 3]], np.int16), 7), + (np.array([True, False, True]), True), + (np.array([1j, 0, np.nan]), 0j), + (np.array(2.5), 0.0), + (np.zeros((0, 3)), 1.0), + ]: + x = make(dense, fill, rng) + y = pickle.loads(pickle.dumps(x)) + assert_array_equal(y, dense) + assert y.nnz == x.nnz + assert y.fill_value == x.fill_value or np.isnan(fill) + a = make(np.array([[1.0, 2.0], [0.0, 3.0]]), 0.0) + P = a @ a + assert_array_equal(pickle.loads(pickle.dumps(P)), a.todense() @ a.todense()) + + +# Indexing and shape functions +@pytest.mark.parametrize( + "key", + [ + 0, + -1, + (1, 2), + (1, 2, 3), + slice(None), + slice(1, None), + slice(None, None, -1), + slice(None, None, 2), + (slice(None), 1), + (Ellipsis, 1), + (1, Ellipsis), + (Ellipsis,), + (None, 0), + (0, None, slice(None, None, -2)), + (Ellipsis, None), + (None, Ellipsis, None), + (slice(5, 1, -1), slice(0, 0)), + (), + (1, slice(0, 2), -1), + ], +) +def test_getitem_basic(key): + rng = np.random.default_rng(2) + a = np.where(rng.random((3, 4, 5)) < 0.5, 0.0, np.round(rng.normal(size=(3, 4, 5)), 1)) + for fill in [0.0, -1.5]: + dense = np.where(a == 0, fill, a) + x = make(dense, fill, rng) + assert_array_equal(x[key], dense[key]) + + +def test_getitem_errors_and_0d(): + x = make(np.arange(6.0).reshape(2, 3), 0.0) + with pytest.raises(IndexError): + x[2] + with pytest.raises(IndexError): + x[0, -4] + with pytest.raises(IndexError): + x[0, 0, 0] + with pytest.raises(IndexError): + x[..., 0, ...] + z = xp.asarray(2.5) + assert_array_equal(z[()], np.array(2.5)) + assert_array_equal(z[...], np.array(2.5)) + assert_array_equal(z[None], np.array([2.5])) + assert_array_equal(x[xp.asarray(1)], np.arange(3.0, 6.0)) + + +@pytest.mark.parametrize( + ("shape", "new"), + [ + ((2, 3, 4), (6, 4)), + ((2, 3, 4), (4, -1)), + ((2, 3, 4), (24,)), + ((2, 3, 4), (2, 3, 4)), + ((6,), (1, 6, 1)), + ((1,), ()), + ((), (1, 1)), + ((), -1), + ((0, 3), (3, 0)), + ], +) +def test_reshape(shape, new): + for dtype in [np.dtype("int16"), np.dtype("float64")]: + for dense, x in inputs(dtype, shape, seed=3): + assert_array_equal(xp.reshape(x, new), np.reshape(dense, new)) + assert_array_equal(xp.reshape(x, new, copy=True), np.reshape(dense, new)) + with pytest.raises(ValueError, match="reshape"): + xp.reshape(xp.zeros((2, 3)), (4, 2)) + with pytest.raises(ValueError, match="-1"): + xp.reshape(xp.zeros((2, 3)), (-1, -1)) + + +def test_reshape_copy_false(): + x = make(np.arange(24.0).reshape(2, 3, 4), 0.0) + out = xp.reshape(x, (2, 12), copy=False) + assert_array_equal(out, np.arange(24.0).reshape(2, 12)) + + +def test_permute_expand_squeeze(): + rng = np.random.default_rng(4) + a = np.where(rng.random((2, 1, 3)) < 0.5, 0, rng.integers(1, 9, (2, 1, 3))) + x = make(a, 0) + for axes in [(0, 1, 2), (2, 0, 1), (-1, -2, -3)]: + assert_array_equal(xp.permute_dims(x, axes), np.transpose(a, axes)) + with pytest.raises(ValueError, match="permutation"): + xp.permute_dims(x, (0, 0, 1)) + for axis in [0, 1, 3, -1, -4, (0, 2), (0, -1)]: + assert_array_equal(xp.expand_dims(x, axis=axis), np.expand_dims(a, axis)) + with pytest.raises(IndexError): + xp.expand_dims(x, axis=5) + assert_array_equal(xp.squeeze(x, axis=1), np.squeeze(a, axis=1)) + assert_array_equal(xp.squeeze(x, axis=(1,)), np.squeeze(a, axis=1)) + with pytest.raises(ValueError, match="squeeze"): + xp.squeeze(x, axis=0) + one = make(np.array([[7]]), 0) + assert_array_equal(xp.squeeze(one, axis=(0, 1)), np.array(7)) + assert_array_equal(xp.expand_dims(xp.asarray(3), axis=0), np.array([3])) + assert_array_equal(xp.permute_dims(xp.asarray(3), ()), np.array(3)) + check_invariants(xp.squeeze(one, axis=(0, 1))) + + +def test_asarray_of_arrays_stays_sparse(): + # a list of Arrays is stacked, not densified (review finding L7) + import graphblas as gb + + x = xp.zeros((10**5, 10**5)) + x[1, 2] = 3.0 + s = xp.asarray([x, x]) + assert s.shape == (2, 10**5, 10**5) + assert s.nnz == 2 + assert_array_equal( + xp.asarray([[make(np.array([1, 0]), 0), make(np.array([0, 2]), 0)]]), + np.array([[[1, 0], [0, 2]]]), + ) + assert_array_equal(xp.asarray([xp.asarray([1, 2]), [3, 4]]), np.array([[1, 2], [3, 4]])) + ragged = [make(np.array([1, 2]), 0), make(np.array([3]), 0)] + with pytest.raises(ValueError, match="inhomogeneous"): # (as numpy refuses it) + xp.asarray(ragged) + assert_array_equal( + xp.asarray([make(np.array([1.0]), 0.0)], dtype=xp.float32), np.array([[1.0]], np.float32) + ) + del gb + + +def test_asarray_byte_order(): + # (review finding L9) + assert_array_equal(xp.asarray(np.array([1, 2], dtype=">i8")), np.array([1, 2])) + assert_array_equal( + xp.asarray(np.array([1.5, 0.0], dtype=">f4")), np.array([1.5, 0.0], np.float32) + ) diff --git a/graphblas/tests/test_array_api_elementwise.py b/graphblas/tests/test_array_api_elementwise.py new file mode 100644 index 000000000..12cf698ce --- /dev/null +++ b/graphblas/tests/test_array_api_elementwise.py @@ -0,0 +1,1103 @@ +"""graphblas.array_api element-wise functions against numpy, position by position. + +The helpers here (``make``, ``assert_array_equal``, ``inputs``) are shared by the other +``test_array_api_*`` modules. +""" + +import math +import operator +import warnings +import zlib + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp +from graphblas import backend +from graphblas.core.tensor import _has_container as tensor_has_container + +if backend != "suitesparse": + pytest.skip("graphblas.array_api needs the suitesparse backend", allow_module_level=True) +if not tensor_has_container(): + pytest.skip("graphblas.array_api needs SuiteSparse:GraphBLAS 10", allow_module_level=True) + + +DTYPES = [ + np.dtype(name) + for name in ( + "bool", + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", + "float32", + "float64", + "complex64", + "complex128", + ) +] +FLOAT_SPECIALS = [np.nan, np.inf, -np.inf, -0.0, 0.0, 1.0, -1.0, 0.5, -2.5, 3.0, 1e30] + + +def random_values(rng, dtype, shape): + """Dense values of ``dtype`` with the special values of the kind.""" + if dtype.kind == "b": + return rng.random(shape) < 0.5 + if dtype.kind in "iu": + info = np.iinfo(dtype) + vals = rng.integers(-4 if dtype.kind == "i" else 0, 5, shape).astype(dtype) + edge = rng.random(shape) < 0.1 + extremes = np.where(rng.random(shape) < 0.5, np.array(info.min, dtype), info.max) + return np.where(edge, extremes, vals) + real = rng.choice(np.array(FLOAT_SPECIALS), size=shape) + real = np.where(rng.random(shape) < 0.5, real, np.round(rng.normal(size=shape) * 3, 2)) + if dtype.kind == "c": + imag = rng.choice(np.array(FLOAT_SPECIALS), size=shape) + imag = np.where(rng.random(shape) < 0.6, 0.0, imag) + out = np.empty(shape, dtype) + with np.errstate(all="ignore"): + out.real = real + out.imag = imag + return out + with np.errstate(all="ignore"): + return real.astype(dtype) + + +def fills_for(dtype): + if dtype.kind == "b": + return [False, True] + if dtype.kind in "iu": + return [0, 3] + if dtype.kind == "c": + return [0, complex(1.5, -2), complex(np.nan, 0)] + return [0.0, -0.0, 2.5, np.nan, -np.inf] + + +def same_value(a, b): + a, b = np.asarray(a), np.asarray(b) + if a.dtype.kind == "c": + return same_value(a.real, b.real) and same_value(a.imag, b.imag) + if a.dtype.kind == "f": + if np.isnan(a) or np.isnan(b): + return bool(np.isnan(a) and np.isnan(b)) + return bool(a == b and np.signbit(a) == np.signbit(b)) + return bool(np.equal(a, b)) + + +def make(dense, fill, rng=None, *, explicit=0.3): + """An Array whose value is ``dense`` at every position: the positions that differ from + ``fill`` are stored, and (with ``rng``) some that equal it are stored too. + """ + dense = np.asarray(dense) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", np.exceptions.ComplexWarning) + fill = np.asarray(fill).astype(dense.dtype)[()] + if dense.ndim == 0: + return xp.asarray(dense) + flat = dense.reshape(-1) + keep = np.array([not same_value(v, fill) for v in flat], dtype=bool) + if rng is not None: + keep |= rng.random(flat.size) < explicit + idx = np.nonzero(keep.reshape(dense.shape)) + out = xp.from_coo(idx, dense[idx], dense.shape, fill_value=fill, dtype=dense.dtype) + assert_array_equal(out, dense) + return out + + +def check_invariants(x): + assert isinstance(x, xp.Array) + assert np.asarray(x.fill_value).dtype == x.dtype + T = x._T() + assert T.shape == x.shape + assert np.dtype(T.dtype.np_type) == x.dtype + if x.ndim == 0: + assert T.nvals == 0 + T._check_invariants() + + +def assert_array_equal(x, expected, *, zero_signs=True): + """``x`` holds ``expected`` exactly (dtype, shape, NaN, signed zeros), but for the sign of + a zero at a position ``x`` does not store, which the contract leaves unspecified (DESIGN + D3b). A 0-d array's value is exact. ``zero_signs=False`` ignores every zero's sign (as + for max and min, where the Array API leaves the order of -0.0 and 0.0 unspecified). + """ + expected = np.asarray(expected) + check_invariants(x) + assert x.dtype == expected.dtype, (x.dtype, expected.dtype) + assert x.shape == expected.shape, (x.shape, expected.shape) + dense = x.todense() + assert dense.dtype == expected.dtype + unstored = np.full(x.shape, not zero_signs) + if x.ndim and zero_signs: + unstored[...] = True + unstored[x.to_coo()[0]] = False + _assert_identical(dense, expected, unstored) + + +def _assert_identical(a, b, unstored): + if a.dtype.kind == "c": + _assert_identical(a.real, b.real, unstored) + _assert_identical(a.imag, b.imag, unstored) + return + if a.dtype.kind == "f": + nan_a, nan_b = np.isnan(a), np.isnan(b) + sign = (np.signbit(a) != np.signbit(b)) & ~(unstored & (a == 0) & (b == 0)) + bad = (nan_a != nan_b) | (~nan_a & ((a != b) | sign)) + else: + bad = a != b + if bad.any(): + idx = np.argwhere(bad)[:5] + detail = [(tuple(int(i) for i in ix), a[tuple(ix)], b[tuple(ix)]) for ix in idx] + raise AssertionError(f"{bad.sum()} positions differ (position, got, expected): {detail}") + + +# Functions computed by libm (SuiteSparse's unary ops) for floating dtypes: special values +# (NaN, inf, zeros and their signs) exactly as numpy has them, other values within MAXULP units +# in the last place. libm and numpy's own loops are each within 1 to 2 ulp of the exact value +# for these (glibc's libm-test-ulps lists at most 2 for the real ones in double; numpy's float32 +# SIMD loops document up to 1.5), so two of them differ by up to 4; on macOS arm64 the largest +# difference seen was 1 (tanh, float32 sin and cos). Complex values stay on numpy. +LIBM = {"acos", "acosh", "asin", "asinh", "atan", "atanh", "cos", "cosh", "exp", "expm1", "log", + "log1p", "log2", "log10", "sin", "sinh", "tan", "tanh"} # fmt: skip +MAXULP = 4 + + +def assert_libm_close(x, expected): + """``x`` holds ``expected`` but for libm's rounding: special values exact, the rest within + MAXULP (complex: each part's special values exact, the rest within MAXULP * eps of the + magnitude). + """ + expected = np.asarray(expected) + check_invariants(x) + assert x.dtype == expected.dtype, (x.dtype, expected.dtype) + assert x.shape == expected.shape, (x.shape, expected.shape) + got = x.todense() + if got.dtype.kind not in "fc": + _assert_identical(got, expected, np.zeros(got.shape, bool)) + return + if got.dtype.kind == "f": + parts = [(got, expected)] + else: + parts = [(got.real, expected.real), (got.imag, expected.imag)] + with np.errstate(all="ignore"): + scale = np.abs(expected) + for g, e in parts: + special = ~np.isfinite(g) | ~np.isfinite(e) | (g == 0) | (e == 0) + _assert_identical(g[special], e[special], np.zeros(special.sum(), bool)) + g, e = g[~special], e[~special] + if got.dtype.kind == "f": + np.testing.assert_array_max_ulp(g, e, maxulp=MAXULP) + else: + # (a part much smaller than the whole is only as exact as the whole) + s = scale[~special] + s = np.where(np.isfinite(s), s, np.abs(e)) + assert (np.abs(g - e) <= MAXULP * np.finfo(e.dtype).eps * s).all(), (g, e) + + +def np_call(func, *args, **kwargs): + with np.errstate(all="ignore"), warnings.catch_warnings(): + warnings.simplefilter("ignore") + return func(*args, **kwargs) + + +def inputs(dtype, shape, seed=0): + """(dense, Array) pairs over the fills of the dtype, some with explicit fills.""" + rng = np.random.default_rng(seed) + for i, fill in enumerate(fills_for(dtype)): + dense = random_values(rng, dtype, shape) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + fillv = np.asarray(fill).astype(dtype)[()] + if dense.ndim: + sparse = rng.random(shape) < 0.5 + dense = np.where(sparse, fillv, dense).astype(dtype) + yield dense, make(dense, fillv, rng if i % 2 else None) + + +# (name, numpy function) +UNARY = [ + ("abs", np.abs), + ("acos", np.arccos), + ("acosh", np.arccosh), + ("asin", np.arcsin), + ("asinh", np.arcsinh), + ("atan", np.arctan), + ("atanh", np.arctanh), + ("bitwise_invert", np.invert), + ("ceil", np.ceil), + ("conj", np.conj), + ("cos", np.cos), + ("cosh", np.cosh), + ("exp", np.exp), + ("expm1", np.expm1), + ("floor", np.floor), + ("imag", np.imag), + ("isfinite", np.isfinite), + ("isinf", np.isinf), + ("isnan", np.isnan), + ("log", np.log), + ("log1p", np.log1p), + ("log2", np.log2), + ("log10", np.log10), + ("logical_not", np.logical_not), + ("negative", np.negative), + ("positive", np.positive), + ("real", np.real), + ("reciprocal", np.reciprocal), + ("round", np.round), + ("sign", np.sign), + ("signbit", np.signbit), + ("sin", np.sin), + ("sinh", np.sinh), + ("sqrt", np.sqrt), + ("square", np.square), + ("tan", np.tan), + ("tanh", np.tanh), + ("trunc", np.trunc), +] + +BINARY = [ + ("add", np.add), + ("atan2", np.arctan2), + ("bitwise_and", np.bitwise_and), + ("bitwise_left_shift", np.left_shift), + ("bitwise_or", np.bitwise_or), + ("bitwise_right_shift", np.right_shift), + ("bitwise_xor", np.bitwise_xor), + ("copysign", np.copysign), + ("divide", np.divide), + ("equal", np.equal), + ("floor_divide", np.floor_divide), + ("greater", np.greater), + ("greater_equal", np.greater_equal), + ("hypot", np.hypot), + ("less", np.less), + ("less_equal", np.less_equal), + ("logaddexp", np.logaddexp), + ("logical_and", np.logical_and), + ("logical_or", np.logical_or), + ("logical_xor", np.logical_xor), + ("maximum", np.maximum), + ("minimum", np.minimum), + ("multiply", np.multiply), + ("nextafter", np.nextafter), + ("not_equal", np.not_equal), + ("pow", np.power), + ("remainder", np.remainder), + ("subtract", np.subtract), +] + +SUPPORTED = set(DTYPES) + + +def expected_or_error(func, *args): + """Numpy's result, or the exception class numpy raises; None when numpy's result dtype is + one graphblas.array_api lacks (float16). + """ + try: + out = np_call(func, *args) + except (TypeError, ValueError, OverflowError) as exc: + return type(exc) + out = np.asarray(out) + if out.dtype not in SUPPORTED: + return None + return out + + +def test_all_elementwise_functions_are_tested(): + names = {name for name, _ in UNARY + BINARY} | {"clip"} + assert names == set(xp._elementwise.__all__) + + +@pytest.mark.parametrize(("name", "func"), UNARY, ids=[n for n, _ in UNARY]) +@pytest.mark.parametrize("shape", [(), (7,), (3, 4), (2, 3, 4)]) +def test_unary(name, func, shape): + for dtype in DTYPES: + for dense, x in inputs(dtype, shape, seed=zlib.crc32(name.encode()) % 1000): + expected = expected_or_error(func, dense) + if expected is None: + with pytest.raises(TypeError): + getattr(xp, name)(x) + continue + if isinstance(expected, type): + with pytest.raises(expected): + getattr(xp, name)(x) + continue + out = getattr(xp, name)(x) + try: + (assert_libm_close if name in LIBM else assert_array_equal)(out, expected) + except AssertionError as exc: + raise AssertionError(f"{name} of {dtype} fill {x.fill_value!r}: {exc}") from None + + +BINARY_SHAPES = [ + ((5,), (5,)), + ((3, 4), (3, 4)), + ((3, 4), (4,)), + ((3, 1), (1, 4)), + ((2, 3, 4), (3, 1)), + ((), (3, 4)), + ((3, 4), ()), + ((), ()), +] +DTYPE_PAIRS = [(dt, dt) for dt in DTYPES] + [ + (np.dtype("int8"), np.dtype("int64")), + (np.dtype("uint8"), np.dtype("int8")), + (np.dtype("int32"), np.dtype("float32")), + (np.dtype("int64"), np.dtype("float64")), + (np.dtype("float32"), np.dtype("float64")), + (np.dtype("float32"), np.dtype("complex128")), + (np.dtype("bool"), np.dtype("int16")), + (np.dtype("uint64"), np.dtype("int64")), +] + + +def _binary_cases(dt1, dt2, shapes, seed): + for i, (dense1, x1) in enumerate(inputs(dt1, shapes[0], seed=seed)): + for j, (dense2, x2) in enumerate(inputs(dt2, shapes[1], seed=seed + 1)): + if (i + j) % 2 == 0 or i == j: + yield dense1, x1, dense2, x2 + + +@pytest.mark.parametrize(("name", "func"), BINARY, ids=[n for n, _ in BINARY]) +@pytest.mark.parametrize("shapes", BINARY_SHAPES, ids=str) +def test_binary(name, func, shapes): + for dt1, dt2 in DTYPE_PAIRS: + for dense1, x1, dense2, x2 in _binary_cases( + dt1, dt2, shapes, zlib.crc32(name.encode()) % 1000 + ): + expected = expected_or_error(func, dense1, dense2) + if expected is None: + with pytest.raises(TypeError): + getattr(xp, name)(x1, x2) + continue + if isinstance(expected, type): + with pytest.raises(expected): + np_call(getattr(xp, name), x1, x2) + continue + out = getattr(xp, name)(x1, x2) + try: + assert_array_equal(out, expected, zero_signs=name not in MAX_MIN) + except AssertionError as exc: + raise AssertionError( + f"{name} of {dt1} fill {x1.fill_value!r} and {dt2} fill {x2.fill_value!r}: " + f"{exc}" + ) from None + + +MAX_MIN = {"maximum", "minimum"} +PY_SCALARS = [0, 1, 2, -3, 200, 2**40, True, False, 0.0, -0.0, 0.5, -1.5, np.nan, np.inf, 1j] + + +@pytest.mark.parametrize(("name", "func"), BINARY, ids=[n for n, _ in BINARY]) +@pytest.mark.parametrize("shape", [(), (3, 4)]) +def test_binary_python_scalar(name, func, shape): + for dtype in DTYPES: + dense, x = next(iter(inputs(dtype, shape, seed=7))) + for s in PY_SCALARS: + for args, dargs in (((x, s), (dense, s)), ((s, x), (s, dense))): + expected = expected_or_error(func, *dargs) + if expected is None: + with pytest.raises(TypeError): + getattr(xp, name)(*args) + continue + if isinstance(expected, type): + with pytest.raises(expected): + np_call(getattr(xp, name), *args) + continue + out = np_call(getattr(xp, name), *args) + try: + assert_array_equal(out, expected, zero_signs=name not in MAX_MIN) + except AssertionError as exc: + raise AssertionError(f"{name}{tuple(map(repr, dargs))}: {exc}") from None + + +OPERATORS = [ + (operator.add, "add"), + (operator.sub, "subtract"), + (operator.mul, "multiply"), + (operator.truediv, "divide"), + (operator.floordiv, "floor_divide"), + (operator.mod, "remainder"), + (operator.pow, "pow"), + (operator.and_, "bitwise_and"), + (operator.or_, "bitwise_or"), + (operator.xor, "bitwise_xor"), + (operator.lshift, "bitwise_left_shift"), + (operator.rshift, "bitwise_right_shift"), + (operator.eq, "equal"), + (operator.ne, "not_equal"), + (operator.lt, "less"), + (operator.le, "less_equal"), + (operator.gt, "greater"), + (operator.ge, "greater_equal"), +] + + +@pytest.mark.parametrize(("op", "name"), OPERATORS, ids=[n for _, n in OPERATORS]) +def test_operators(op, name): + rng = np.random.default_rng(1) + a = rng.integers(1, 5, (3, 4)) + b = rng.integers(1, 5, (4,)) + x, y = make(a, 0), make(b, 2) + assert_array_equal(op(x, y), np_call(op, a, b)) + assert_array_equal(op(x, 3), np_call(op, a, 3)) + assert_array_equal(op(3, y), np_call(op, 3, b)) + # numpy arrays defer to the Array + out = op(a, y) + assert isinstance(out, xp.Array) + assert_array_equal(out, np_call(op, a, b)) + + +def test_unary_operators(): + a = np.array([1, -2, 0, 3]) + x = make(a, 0) + assert_array_equal(-x, -a) + assert_array_equal(+x, +a) + assert_array_equal(abs(x), abs(a)) + assert_array_equal(~x, ~a) + + +@pytest.mark.parametrize( + ("op", "name"), + [(operator.iadd, "add"), (operator.imul, "multiply"), (operator.isub, "subtract")], +) +def test_inplace_operators(op, name): + a = np.array([[1.0, 0.0, 2.0], [0.0, -0.0, np.nan]]) + x = make(a, 0.0) + orig = x + x = op(x, make(np.array([1.0, 2.0, 0.0]), 0.0)) + assert x is orig + assert_array_equal(x, np_call(op, a.copy(), np.array([1.0, 2.0, 0.0]))) + x = op(x, 2) + assert x is orig + with pytest.raises(TypeError): + op(make(np.array([1, 2]), 0), 1.5) + with pytest.raises(ValueError, match="shape"): + op(make(np.array([1.0, 2.0]), 0.0), make(np.ones((3, 2)), 0.0)) + + +@pytest.mark.parametrize("dtype", ["int8", "uint8", "float32", "float64", "int64"]) +def test_clip(dtype): + rng = np.random.default_rng(3) + dt = np.dtype(dtype) + for dense, x in inputs(dt, (3, 4), seed=4): + for lo, hi in [(1, None), (None, 2), (0, 3), (-1, 1)]: + if dt.kind == "u" and lo == -1: + lo = 0 + expected = np.clip(dense, lo, hi) if (lo, hi) != (None, None) else dense + out = xp.clip(x, min=lo, max=hi) + assert_array_equal(out, np.asarray(expected).astype(dt)) + lo_dense = random_values(rng, dt, (4,)) + lo_arr = make(lo_dense, 0) + with np.errstate(all="ignore"): + expected = np.maximum(dense, lo_dense).astype(dt) + assert_array_equal(xp.clip(x, min=lo_arr), expected) + x = make(np.array([1, 2, 3], dtype=np.int8), 0) + assert_array_equal(xp.clip(x, min=-1000, max=1000), np.array([1, 2, 3], dtype=np.int8)) + assert_array_equal(xp.clip(x), np.array([1, 2, 3], dtype=np.int8)) + + +def test_signed_zero_union(): + # 0.0 fill: an unstored position is +0.0, a stored -0.0 differs + x = make(np.array([-0.0, 0.0, 1.0, -1.0]), 0.0) + y = make(np.array([0.0, -0.0, 0.0, 0.0]), 0.0) + assert x.nnz == 3 + assert y.nnz == 1 + a, b = x.todense(), y.todense() + for name, func in [("add", np.add), ("multiply", np.multiply), ("subtract", np.subtract)]: + assert_array_equal(getattr(xp, name)(x, y), func(a, b)) + assert_array_equal(getattr(xp, name)(y, x), func(b, a)) + + +def test_integer_multiply_intersects(): + # Both fills 0: integer products are 0 where either side is unstored, so the result + # stores only the intersection of entries. + x = make(np.array([1, 0, 2, 0, 3]), 0) + y = make(np.array([0, 4, 5, 0, 6]), 0) + out = x * y + assert out.nnz == 2 + assert_array_equal(out, np.array([0, 0, 10, 0, 18])) + + +def test_float_multiply_of_nonnegative_values_intersects(): + x = make(np.array([1.0, 0.0, 2.0, 0.0, 3.0]), 0.0) + y = make(np.array([0.0, 4.0, 5.0, 0.0, 6.0]), 0.0) + out = x * y + assert out.nnz == 2 + assert_array_equal(out, x.todense() * y.todense()) + # a negative value makes -0.0 where the other operand is unstored + x *= -1 + out = x * y + assert_array_equal(out, x.todense() * y.todense()) + assert out.nnz > 2 + z = make(np.array([0.0, np.inf, 1.0, 0.0, 0.0]), 0.0) + assert_array_equal(z * y, np_call(np.multiply, z.todense(), y.todense())) + + +@pytest.mark.parametrize("name", ["maximum", "minimum"]) +def test_float_maximum_minimum(name): + # Without NaN GraphBLAS's max and min are numpy's; with it, numpy computes. The order of + # -0.0 and 0.0 is unspecified. + rng = np.random.default_rng(8) + func = getattr(np, name) + + def check(x, expected): + assert_array_equal(x, expected, zero_signs=False) + + for fill1, fill2 in [(0.0, 0.0), (2.5, -1.0), (np.inf, 0.0), (0.0, -0.0), (np.nan, 1.0)]: + a = np.where(rng.random((4, 5)) < 0.5, fill1, np.round(rng.normal(size=(4, 5)), 1)) + b = np.where(rng.random(5) < 0.5, fill2, np.round(rng.normal(size=5), 1)) + x, y = make(a, fill1), make(b, fill2) + check(getattr(xp, name)(x, y), func(a, b)) + for s in [0.0, -0.0, 1.5, np.nan, -np.inf, 3]: + check(getattr(xp, name)(x, s), func(a, s)) + check(getattr(xp, name)(s, y), func(s, b)) + a[0, 0] = -0.0 + a[1, 1] = np.nan + x = make(a, fill1) + check(getattr(xp, name)(x, y), func(a, b)) + check(getattr(xp, name)(y, x), func(b, a)) + x = make(np.array([1.0, 0.0, 2.0]), 0.0) + y = make(np.array([0.5, 0.5, 0.5]), 0.0) + check(getattr(xp, name)(x, y), func(x.todense(), y.todense())) + x *= -0.0 # (the cached facts about x's values must not survive) + check(getattr(xp, name)(x, y), func(x.todense(), y.todense())) + + +def test_float_multiply_keeps_union(): + x = make(np.array([-1.0, 0.0, 2.0, np.inf]), 0.0) + y = make(np.array([0.0, 4.0, 5.0, 0.0]), 0.0) + assert_array_equal(x * y, np_call(np.multiply, x.todense(), y.todense())) + + +def test_broadcast_expansion_is_guarded(): + x = make(np.arange(1, 101, dtype=np.float64).reshape(100, 1), 0.0) + y = make(np.arange(1, 101, dtype=np.float64).reshape(1, 100), 0.0) + with gb.config.set(array_api_densify_limit=1000), pytest.raises(xp.DensifyError): + x + y + assert_array_equal(x + y, x.todense() + y.todense()) + + +def test_values_mapped_on_every_storage_format(): + # numpy maps the stored values through the Container in place: check bitmap, full, + # hypersparse and iso storage + rng = np.random.default_rng(5) + for density in [0.0, 0.05, 0.5, 1.0]: + for shape in [(40,), (30, 20), (6, 5, 4)]: + a = np.round(rng.normal(size=shape), 2) + a[rng.random(shape) >= density] = 0.0 + x = make(a, 0.0) + assert_array_equal(xp.sin(x), np.sin(a)) + assert_array_equal(xp.pow(x, x), np_call(np.power, a, a)) + for fmt in ["bitmap", "sparse", "hypersparse", "full"]: + M = x._T()._M + if fmt == "full" and M.nvals != M.nrows * M.ncols: + continue + M.ss.config["sparsity_control"] = fmt + assert_array_equal(xp.exp(x), np.exp(a)) + assert_array_equal(xp.floor_divide(x, 0.3), np_call(np.floor_divide, a, 0.3)) + iso = xp.from_coo(([0, 3, 5],), 2.5, (8,), fill_value=0.0) + assert iso._T().ss.is_iso + assert_array_equal(xp.log(iso), np_call(np.log, iso.todense())) + other = make(np.arange(8.0), 0.0) + assert_array_equal(xp.atan2(iso, other), np.arctan2(iso.todense(), other.todense())) + assert_array_equal(xp.atan2(other, iso), np.arctan2(other.todense(), iso.todense())) + + +def test_unseen_fill_does_not_raise(): + # x stores every position, so its fill (-1) never meets y's fill as an exponent + x = xp.from_coo(([0, 1, 2],), [2, 3, 4], (3,), fill_value=-1) + y = xp.full((3,), -2) + assert_array_equal(xp.pow(y, x), np.power(np.full(3, -2), [2, 3, 4])) + with pytest.raises(ValueError, match="negative"): + xp.pow(make(np.array([2, 0, 3]), 0), -1) + + +def test_zero_d_results_store_nothing(): + x = xp.asarray(2.5) + y = xp.add(x, 1) + check_invariants(y) + assert float(y) == 3.5 + z = make(np.array([1.0, 2.0]), 0.0)[1] + check_invariants(z) + assert float(z) == 2.0 + assert z._T().nvals == 0 + + +def test_results_do_not_share_writes(): + x = make(np.array([1.0, 0.0, 2.0]), 0.0) + y = +x + y += 1 + assert_array_equal(x, np.array([1.0, 0.0, 2.0])) + assert_array_equal(y, np.array([2.0, 1.0, 3.0])) + + +@pytest.mark.parametrize("dtype", ["float32", "float64", "complex64", "complex128"]) +def test_libm_functions_on_wide_values(dtype): + # Values over the whole range, the edges of each domain and the special values + rng = np.random.default_rng(11) + dt = np.dtype(dtype) + info = np.finfo(dt) + edges = [0.0, -0.0, 1.0, -1.0, 0.5, -0.5, 2.0, info.tiny, -info.tiny, info.tiny / 8, + info.max, -info.max, 88.0, 89.0, 709.0, 711.0, -745.0, np.pi, np.nan, np.inf, + -np.inf] # fmt: skip + n = 3000 + v = np.concatenate([ + np.array(edges), + rng.normal(size=n) * 3, + np.exp(rng.uniform(-60, 60, n)) * rng.choice([-1.0, 1.0], n), + rng.uniform(-1.2, 1.2, n), + ]) # fmt: skip + with np.errstate(all="ignore"): + if dt.kind == "c": + w = np.empty(v.size, dt) + w.real, w.imag = v, rng.permutation(v) + v = w + v = v.astype(dt) + x = make(v, 0) + for name in sorted(LIBM): + func = dict(UNARY)[name] + try: + assert_libm_close(getattr(xp, name)(x), np_call(func, v)) + except AssertionError as exc: + raise AssertionError(f"{name} of {dtype}: {exc}") from None + + +def test_math_consistency_of_fill(): + # An entry equal to the fill maps to the new fill, bit for bit (DESIGN D10: no value + # depends on which positions are stored), also where libm computes the entries + for name, _ in UNARY: + if name in {"logical_not", "bitwise_invert"}: + continue + x = xp.from_coo(([0, 1],), [0.7, 0.7], (4,), fill_value=0.7) + out = getattr(xp, name)(x) + dense = out.todense() + assert all(same_value(v, out.fill_value) for v in dense), name + + +def test_isnan_of_integers_stores_nothing(): + x = make(np.array([1, 0, 2]), 0) + for name, value in [("isnan", False), ("isinf", False), ("isfinite", True)]: + out = getattr(xp, name)(x) + assert out.nnz == 0 + assert out.fill_value == value + + +def test_math_constants(): + assert xp.e == math.e + assert xp.pi == math.pi + assert math.isinf(xp.inf) + assert math.isnan(xp.nan) + assert xp.newaxis is None + + +def test_unsupported_operands(): + x = make(np.array([1.0, 2.0]), 0.0) + with pytest.raises(TypeError): + xp.add(x, "a") + assert (x == "a") is False + with pytest.raises(TypeError): + xp.sin(make(np.array([True, False]), False)) # numpy would give float16 + + +def test_intersection_of_broadcast_operands_is_guarded(): + # an intersection of two broadcast operands is an outer product (review finding M4) + x = make(np.arange(1, 101).reshape(100, 1), 0) + y = make(np.arange(1, 101).reshape(1, 100), 0) + b = make(np.ones((100, 1), bool), False) + for func, args, nnz in [ + (xp.multiply, (x, y), 10000), + (xp.bitwise_and, (x, y), 10000), + (xp.logical_and, (b, b.mT), 10000), + ]: + with gb.config.set(array_api_densify_limit=1000), pytest.raises(xp.DensifyError): + func(*args) + assert func(*args).nnz == nnz + + +def test_absorbing_fills_intersect(): + # A fill that absorbs the op (inf + y is inf) on both operands: the result is the + # intersection of their entries, and a broadcast operand is not replicated + rng = np.random.default_rng(31) + n = 40 + g = np.where(rng.random((n, n)) < 0.1, rng.random((n, n)) + 1, np.inf) + d = np.where(rng.random(n) < 0.5, rng.random(n), np.inf) + G, D = make(g, np.inf, rng), make(d, np.inf, rng) + for func, nf in [(xp.add, np.add), (xp.maximum, np.maximum)]: + R = func(xp.expand_dims(D, 1), G) + assert_array_equal(R, nf(d[:, None], g)) + assert R.nnz <= G.nnz + assert R.fill_value == np.inf + R = xp.minimum(make(-g, -np.inf), make(-d[:, None], -np.inf)) + assert_array_equal(R, np.minimum(-g, -d[:, None])) + assert R.nnz <= G.nnz + R = xp.add(make(-g, -np.inf), make(-d, -np.inf)) # (-inf absorbs add too) + assert_array_equal(R, -g + -d) + assert R.nnz <= G.nnz + # integers' extremes, and bools + big = np.iinfo(np.int64).max + gi = np.where(np.isinf(g), big, (np.where(np.isinf(g), 0, g) * 10).astype(np.int64)) + di = np.where(np.isinf(d), big, (np.where(np.isinf(d), 0, d) * 10).astype(np.int64)) + R = xp.maximum(make(gi, big), make(di, big)) + assert_array_equal(R, np.maximum(gi, di)) + assert R.nnz <= G.nnz + small = np.iinfo(np.int8).min + R = xp.minimum(make(np.int8(-3) * (gi % 2), 0), make(np.full((n, n), small, np.int8), small)) + assert_array_equal(R, np.minimum(np.int8(-3) * (gi % 2), np.int8(small))) + gb_, db_ = np.isinf(g), np.isinf(d) + for func, nf in [(xp.logical_or, np.logical_or), (xp.add, np.add), (xp.maximum, np.maximum)]: + R = func(make(gb_, True), make(db_, True)) + assert_array_equal(R, nf(gb_, db_)) + assert R.nnz <= (~gb_).sum() + R = xp.minimum(make(~gb_, False), make(~db_, False)) + assert_array_equal(R, np.minimum(~gb_, ~db_)) + # An operand that stores every position counts as holding the absorbing fill + full = rng.random(n) + 1 + R = xp.expand_dims(make(full, 0.0), 1) + G + assert_array_equal(R, full[:, None] + g) + assert R.nnz <= G.nnz + assert R.fill_value == np.inf + # NaN, or for add the opposite infinity, anywhere: the union, exact + for bad in [np.nan, -np.inf]: + d2 = d.copy() + d2[0] = bad + R = xp.add(xp.expand_dims(make(d2, np.inf), 1), G) + with np.errstate(invalid="ignore"): + assert_array_equal(R, d2[:, None] + g) + g2 = g.copy() + g2[0, 0] = -np.inf + assert_array_equal(xp.add(make(d, np.inf), make(g2, np.inf)), d + g2) + # Fills that do not both absorb, and two fully stored operands: the union as before + assert_array_equal(xp.add(make(d, np.inf), make(d, 0.0)), d + d) + R = xp.add(make(full, 0.0), make(full, 1.0)) + assert_array_equal(R, full + full) + with np.errstate(invalid="ignore"): + assert_array_equal(xp.subtract(make(d, np.inf), make(d, np.inf)), d - d) + + +def test_fully_stored_bool_results_take_fill_false(): + # A bool result that stores every position never shows its fill: it takes fill False, + # and its False entries go when at most a quarter of the entries are True (equal(deg, 0) + # of a fully stored deg stores the zero-degree positions only) + deg = xp.asarray(np.array([3, 1, 0, 2, 5, 0, 4, 1, 2, 7])) + assert deg.nnz == 8 + assert deg.fill_value == 0 + full = xp.asarray(np.array([3, 1, 9, 2, 5, 6, 4, 1, 2, 7]), fill_value=-1) + assert full.nnz == 10 + for out, expected, nnz, fill in [ + (xp.equal(full, 1), full.todense() == 1, 2, False), + (xp.greater(full, 1), full.todense() > 1, 10, False), # (8 of 10 True: kept) + (xp.less(full, 5), full.todense() < 5, 10, False), + (xp.isfinite(xp.astype(full, xp.float64)), np.ones(10, bool), 10, False), + (xp.astype(full == 9, xp.bool), full.todense() == 9, 1, False), + ]: + assert_array_equal(out, expected) + assert (out.nnz, out.fill_value) == (nnz, fill) + # a 0-d result is left as it is + assert_array_equal(xp.equal(xp.asarray(3), 3), np.array(True)) + + +def test_bool_normalization_decides_by_a_sample(monkeypatch): + # (review finding L10) a sample of the leading entries says whether dropping the False + # entries is likely; values never change, only what is stored + from graphblas.array_api import _elementwise + + monkeypatch.setattr(_elementwise, "_SAMPLE", 4) + v = np.arange(40.0).reshape(10, 4) + full = xp.asarray(v + 1) + for out, expected, nnz in [ + (xp.greater(full, 4), v + 1 > 4, 40), # (sample mostly True: kept, fill False) + (xp.greater(full, 36), v + 1 > 36, 4), # (few True: dropped) + (xp.less(full, 37), v + 1 < 37, 40), # (sample all True, most True: kept) + (xp.less(full, 5), v + 1 < 5, 40), # (the sample is all True: kept, though few are) + (xp.greater(full, 3), v + 1 > 3, 40), # (a quarter of the sample: the select decides) + ]: + assert_array_equal(out, expected) + assert (out.nnz, out.fill_value) == (nnz, False) + + +def test_one_sided_stores_a_small_broadcast_operand(): + # (review finding R2) a broadcast operand no larger than the other's entries is stored in + # full and taken by the broadcast product, with no coordinates exported + rng = np.random.default_rng(21) + a = np.where(rng.random((50, 40)) < 0.6, rng.random((50, 40)) + 0.5, 0.0) + y = np.where(rng.random(50) < 0.1, rng.random(50), 0.0) + A, Y = make(a, 0.0), xp.exp(make(y, 0.0)) + assert Y.nnz < 50 + assert Y.fill_value == 1.0 + out = A * Y[:, None] + assert out.nnz == A.nnz + assert_array_equal(out, a * np.exp(y)[:, None]) + + +def test_partly_stored_bool_results_with_mostly_false_take_fill_false(): + # equal(deg, 0) of a deg that leaves its zero degrees unstored has fill True and stores + # a False at every other position; when those are more than half of all positions, it + # takes fill False and stores the True positions only + deg = xp.asarray(np.array([3, 1, 0, 2, 5, 0, 4, 1, 2, 7])) + out = xp.equal(deg, 0) + assert_array_equal(out, deg.todense() == 0) + assert (out.nnz, out.fill_value) == (2, False) + # with stored True entries too (explicit fills of a fill-True operand) + x = make(np.array([True, False, False, False, True, False]), True, np.random.default_rng(3)) + out = xp.logical_and(x, x) + assert_array_equal(out, x.todense()) + assert out.fill_value == np.False_ or out.nnz * 2 <= out.size + for dense in [np.eye(6, dtype=bool), np.array([[True, False], [False, False]])]: + out = xp.logical_not(xp.logical_not(make(dense, False))) + assert_array_equal(out, dense) + assert out.fill_value == np.False_ + assert out.nnz == dense.sum() + # mostly True, or not more than half stored: left with fill True + sparse_deg = xp.asarray(np.array([0, 0, 0, 2, 0, 0])) + assert (xp.equal(sparse_deg, 0).nnz, xp.equal(sparse_deg, 0).fill_value) == (1, True) + y = make( + np.array([True, True, True, False, True]), True, np.random.default_rng(4), explicit=0.9 + ) + assert xp.logical_and(y, y).fill_value == np.True_ or y.nnz * 2 <= y.size + + +def test_mostly_true_masks_stay_sparse_in_contractions(): + # (review finding H3) a mostly-True fully stored mask keeps fill False, so a contraction + # with a sparse bool matrix stays sparse (it raised DensifyError with fill True) + rng = np.random.default_rng(15) + n = 300 + a = np.zeros((n, n), bool) + a[rng.integers(0, n, 900), rng.integers(0, n, 900)] = True + A = make(a, False) + d = rng.integers(1, 10, n).astype(float) + m = xp.greater(xp.asarray(d), 1.5) + assert m.fill_value == False # noqa: E712 + mv = d > 1.5 + with gb.config.set(array_api_densify_limit=n * 10): + assert_array_equal(m @ A, mv @ a) + assert_array_equal(xp.einsum("i,ij->j", m, A), np.einsum("i,ij->j", mv, a)) + assert_array_equal( + xp.matmul(m, A, semiring="lxor_land"), + np.logical_xor.reduce(mv[:, None] & a, axis=0), + ) + out = xp.matmul(xp.astype(m, xp.uint8), xp.astype(A, xp.uint8), semiring="max_times") + assert_array_equal(out, np.max(mv[:, None].astype(np.uint8) * a, axis=0)) + + +def test_clip_with_array_bounds_ties(): + # numpy's clip gives the bound on a tie, so the sign of a zero is the bound's (review M6) + x = make(np.array([0.0, -0.0, 0.0, 1.0, np.nan]), 0.0) + lo = make(np.array([-0.0, 0.0, 0.0, -0.0, 0.0]), 0.0) + hi = make(np.array([-0.0, 0.0, 2.0, 0.5, 1.0]), 0.0) + a, b, c = x.todense(), lo.todense(), hi.todense() + assert_array_equal(xp.clip(x, min=lo, max=hi), np_call(np.clip, a, b, c)) + assert_array_equal(xp.clip(x, min=lo), np_call(np.clip, a, b, None)) + assert_array_equal(xp.clip(x, max=hi), np_call(np.clip, a, None, c)) + assert_array_equal(xp.clip(x, min=lo, max=0.5), np_call(np.clip, a, b, 0.5)) + # the bounds may broadcast x, as in numpy + lo2 = make(np.array([[0.5] * 5, [-1.0] * 5]), 0.0) + assert_array_equal(xp.clip(x, min=lo2), np_call(np.clip, a, lo2.todense(), None)) + s = xp.asarray(np.uint8(0)) + assert_array_equal(xp.clip(s, max=xp.zeros((0,), dtype=xp.uint8)), np.zeros(0, np.uint8)) + assert_array_equal(xp.clip(xp.asarray(2.0), min=make(np.array([1.0, 3.0]), 0.0)), [2.0, 3.0]) + + +def test_signed_zeros_give_way(): + # (review finding H1, DESIGN D3a, D3b) a product with an unstored zero is the fill + # whatever its sign: a row scaling stores only A's entries + rng = np.random.default_rng(0) + A = make(np.where(rng.random((50, 50)) < 0.05, rng.normal(size=(50, 50)), 0.0), 0.0) + d = make(rng.normal(size=(50, 1)), 0.0) + out = d * A + assert out.nnz == A.nnz + assert_array_equal(out, d.todense() * A.todense()) + # inf and NaN stay exact + inf = make(np.array([[np.inf, 0.0]]), 0.0) + assert_array_equal(inf * make(np.array([[0.0, 1.0]]), 0.0), np.array([[np.nan, 0.0]])) + # divide: the numerator's entries, when no value of the denominator is 0 or NaN + a = np.where(rng.random((4, 5)) < 0.5, 0.0, rng.normal(size=(4, 5))) + b = rng.normal(size=5) + 3.0 + for fb in [1.0, -2.0]: + q = make(a, 0.0) / make(b, fb) + assert q.nnz == np.count_nonzero(a) + assert_array_equal(q, a / b) + q = make(a, 0.0) / make(np.r_[b[:4], 0.0], 1.0) + assert_array_equal(q, np_call(np.divide, a, np.r_[b[:4], 0.0])) + + +@pytest.mark.parametrize("name", ["multiply", "divide"]) +def test_products_match_numpy_but_for_unstored_zero_signs(name): + func = getattr(np, name) + for shapes in [((5,), (5,)), ((3, 4), (4,)), ((3, 1), (1, 4))]: + for dt1, dt2 in [(np.dtype("float64"),) * 2, (np.dtype("float32"), np.dtype("int8"))]: + for dense1, x1, dense2, x2 in _binary_cases(dt1, dt2, shapes, 5): + out = getattr(xp, name)(x1, x2) + assert_array_equal(out, np_call(func, dense1, dense2)) + + +def test_operands_storing_every_position_keep_products_sparse(): + # An operand that stores every position never shows its fill (review finding N1): A / d + # and A * (1 / deg)[:, None] store A's entries, whatever d's fill + rng = np.random.default_rng(3) + A = make(np.where(rng.random((40, 40)) < 0.05, rng.normal(size=(40, 40)), 0.0), 0.0) + d = xp.asarray(rng.normal(size=40) + 5.0) + deg = xp.asarray(rng.integers(1, 9, 40).astype(float)) + q = A / d + assert q.nnz == A.nnz + assert_array_equal(q, A.todense() / d.todense()) + inv = 1 / deg + assert inv.fill_value == np.inf + p = A * xp.reshape(inv, (40, 1)) + assert p.nnz == A.nnz + assert p.fill_value == 0.0 + assert_array_equal(p, A.todense() * inv.todense()[:, None]) + # integers too: a full operand's fill takes no part + i = make(np.where(rng.random((5, 6)) < 0.5, 0, rng.integers(1, 5, (5, 6))), 0) + full = xp.asarray(rng.integers(1, 5, 6), fill_value=7) + assert (i * full).nnz == i.nnz + assert_array_equal(i * full, i.todense() * full.todense()) + assert (A / xp.exp(d)).nnz == A.nnz + assert_array_equal(A / xp.exp(d), A.todense() / np.exp(d.todense())) + signed = xp.asarray(rng.normal(size=40)) + assert (A / signed).nnz == A.nnz + assert_array_equal(A / signed, A.todense() / signed.todense()) + + +def test_one_sided_products(): + # one operand's fill is 0, the other's is not: the result lies within the first's entries + # (a broadcast second operand is stored in full first, not replicated) + rng = np.random.default_rng(4) + A = make(np.where(rng.random((30, 30)) < 0.05, rng.normal(size=(30, 30)), 0.0), 0.0) + v = make(np.where(rng.random(30) < 0.5, 0.0, rng.normal(size=30)), 0.0) + e = xp.exp(v) # (fill 1, some positions unstored) + w = make(np.where(rng.random(30) < 0.3, 1.0, rng.normal(size=30) + 4), 1.0) + for out, expected in [ + (A * xp.reshape(e, (1, 30)), A.todense() * e.todense()[None, :]), + (xp.reshape(e, (30, 1)) * A, e.todense()[:, None] * A.todense()), + (A / xp.reshape(w, (30, 1)), A.todense() / w.todense()[:, None]), + (A / e, A.todense() / e.todense()), + ]: + assert out.nnz <= A.nnz + assert_array_equal(out, expected) + # a second operand as large as the result: the masked union, no full copy + B = make(np.where(rng.random((60, 60)) < 0.05, rng.normal(size=(60, 60)), 0.0), 0.0) + big = make(np.where(rng.random((60, 60)) < 0.5, 2.0, rng.normal(size=(60, 60)) + 5), 2.0) + out = B / big + assert out.nnz == B.nnz + assert_array_equal(out, B.todense() / big.todense()) + + +def test_unstored_zero_signs_in_the_oracle(): + # The tests' oracle: a stored -0.0 must match, an unstored one may read 0.0 + x = make(np.array([-0.0, 1.0]), 0.0) + assert x.nnz == 2 + with pytest.raises(AssertionError, match="differ"): + assert_array_equal(x, np.array([0.0, 1.0])) + y = make(np.array([0.0, 1.0]), 0.0) + assert_array_equal(y, np.array([-0.0, 1.0])) + with pytest.raises(AssertionError, match="differ"): + assert_array_equal(xp.asarray(-0.0), np.array(0.0)) + + +def test_one_sided_broadcast_operand_is_gathered(): + # (review findings rev1 M9, rev2 M5) a broadcast operand that shows its fill is read at the + # other operand's entries only, O(those entries): X of shape (1000, 1e6) with 4 entries + # times exp(v)[None, :] (fill 1) stores 4 entries, under a densify limit of 1000 + rng = np.random.default_rng(12) + n = 10**6 + rows, cols = [0, 10, 500, 999], [5, 77, n - 1, 123_456] + X = xp.from_coo((rows, cols), [1.0, 2.0, 3.0, 4.0], (1000, n)) + v = xp.from_coo((np.sort(rng.choice(n, n // 2, replace=False)),), rng.normal(size=n // 2), (n,)) + w = make(np.where(rng.random(n) < 0.5, 2.0, rng.random(n) + 1), 2.0) + vals = np.array([1.0, 2.0, 3.0, 4.0]) + e = xp.exp(v) + ed, wd = e.todense()[cols], w.todense()[cols] + with gb.config.set(array_api_densify_limit=1000): + for out, expected in [ + (X * e[None, :], vals * ed), + (e[None, :] * X, ed * vals), + (X / w, vals / wd), + (X * xp.asarray(-2.0), vals * -2.0), + ]: + assert out.nnz == 4 + assert out.fill_value == 0.0 + coords, values = out.to_coo() + assert coords[0].tolist() == rows + np.testing.assert_array_equal(values, expected) + # numpy's answer at every position, on small ones: broadcast along either axis, 3-d, a + # 0-d operand, explicit fills, an X with no entries + A = make(np.where(rng.random((6, 50)) < 0.05, rng.normal(size=(6, 50)), 0.0), 0.0, rng) + u = make(np.where(rng.random(50) < 0.5, 1.0, rng.normal(size=50)), 1.0, rng) + c = make(np.where(rng.random((6, 1)) < 0.5, 3.0, rng.normal(size=(6, 1))), 3.0) + B = make(np.where(rng.random((4, 6, 5)) < 0.2, rng.normal(size=(4, 6, 5)), 0.0), 0.0) + g = make(np.where(rng.random((6, 1)) < 0.5, 2.0, rng.normal(size=(6, 1)) + 3), 2.0) + for out, expected in [ + (A * u[None, :], A.todense() * u.todense()[None, :]), + (A * c, A.todense() * c.todense()), + (B * c, B.todense() * c.todense()), + (B / g, B.todense() / g.todense()), + (xp.zeros((6, 50)) * u, np.zeros((6, 50)) * u.todense()), + ]: + assert out.nnz <= max(A.nnz, B.nnz) + assert_array_equal(out, expected) + + +def test_clip_fill_uses_numpys_array_loop(): + # numpy's 0-d clip breaks a tie differently from its array loop (review finding N3) + x = xp.full((3,), 0.0) + lo = xp.full((3,), -0.0) + a, b = np.zeros(3), np.full(3, -0.0) + assert_array_equal(xp.clip(x, min=lo, max=0.5), np.clip(a, b, 0.5)) + assert_array_equal(xp.clip(x, min=-0.0, max=0.5), np.clip(a, -0.0, 0.5)) + + +def _spec_pow(x1, x2): + """The Array API's pow of two Python floats, case by case (its special cases).""" + odd = math.isfinite(x2) and x2 % 2 == 1 + if math.isnan(x2): + return 1.0 if x1 == 1 else math.nan + if x2 == 0: + return 1.0 + if math.isnan(x1): + return math.nan + if math.isinf(x2): + if abs(x1) == 1: + return 1.0 + return math.inf if (abs(x1) > 1) == (x2 > 0) else 0.0 + if x1 == math.inf: + return math.inf if x2 > 0 else 0.0 + if x1 == -math.inf: + if x2 > 0: + return -math.inf if odd else math.inf + return -0.0 if odd else 0.0 + if x1 == 0: + negative = math.copysign(1.0, x1) < 0 and odd + if x2 > 0: + return -0.0 if negative else 0.0 + return -math.inf if negative else math.inf + if x1 < 0 and x2 != int(x2): + return math.nan + return x1**x2 + + +def test_pow_special_cases(): + # The Array API's special cases of pow, with arrays, 0-d arrays and Python scalars as + # exponents (numpy takes sqrt for an exponent of 0.5 that repeats: pow(-inf, 0.5) was NaN + # there, review finding L11) + bases = [np.nan, np.inf, -np.inf, 0.0, -0.0, 2.0, 0.5, 1.0, -1.0, -2.0, -0.5] + exponents = [np.nan, 0.0, -0.0, np.inf, -np.inf, 0.5, 2.0, 3.0, -1.0, -2.0, -3.0, 1.5, -0.5] + for dt in [np.float64, np.float32]: + b = np.array(bases, dt) + for e in exponents: + expected = np.array([_spec_pow(float(v), e) for v in b], dt) + x = make(b, dt(-np.inf)) + for exponent in [e, xp.asarray(dt(e)), xp.full((len(b),), dt(e))]: + assert_array_equal(np_call(xp.pow, x, exponent), expected) + assert_array_equal(np_call(xp.pow, make(b, dt(0.0)), e), expected) + # a base that repeats: a scalar, or a 0-d array + for v in bases: + e = make(np.array(exponents, dt), dt(0.5)) + expected = np.array([_spec_pow(v, float(k)) for k in e.todense()], dt) + assert_array_equal(np_call(xp.pow, xp.asarray(dt(v)), e), expected) diff --git a/graphblas/tests/test_array_api_errors.py b/graphblas/tests/test_array_api_errors.py new file mode 100644 index 000000000..6041aa04a --- /dev/null +++ b/graphblas/tests/test_array_api_errors.py @@ -0,0 +1,348 @@ +"""graphblas.array_api: argument checking, rare storage cases and extension paths.""" + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp +from graphblas.array_api import _utils +from graphblas.core.tensor import Tensor + +from .test_array_api_elementwise import assert_array_equal, make + + +def test_array_arguments_are_checked(): + a = np.array([1.0, 2.0]) + for func in [xp.reshape, xp.squeeze, xp.permute_dims, xp.expand_dims]: + with pytest.raises(TypeError, match="needs an Array"): + func(a, 0) + for func in [xp.sum, xp.diff, xp.cumulative_sum, xp.tril, xp.zeros_like, xp.clip]: + with pytest.raises(TypeError, match="needs an Array"): + func(a) + with pytest.raises(TypeError, match="needs an Array"): + xp.astype(a, xp.float32) + with pytest.raises(TypeError, match="arrays or Python scalars"): + xp.add(make(a, 0.0), [1, 2]) + + +def test_scalars_alone(): + assert_array_equal(xp.sin(1.0), np.sin(np.float64(1.0))) + assert_array_equal(xp.add(1, 2), np.array(3)) + assert_array_equal(xp.multiply(np.float32(2), 3), np.array(6, np.float32)) + + +def test_shape_errors(): + x = make(np.arange(6.0).reshape(2, 3), 0.0) + with pytest.raises(ValueError, match="negative"): + xp.zeros((2, -1)) + with pytest.raises(ValueError, match="invalid shape"): + xp.reshape(x, (-2, 3)) + with pytest.raises(ValueError, match="cannot reshape"): + xp.reshape(x, (4, -1)) + with pytest.raises(ValueError, match="repeated"): + xp.expand_dims(x, axis=(0, 0)) + with pytest.raises(ValueError, match="repeated"): + xp.squeeze(xp.zeros((1, 1)), axis=(0, -2)) + assert_array_equal(xp.squeeze(x, axis=()), np.arange(6.0).reshape(2, 3)) + with pytest.raises(ValueError, match="do not broadcast"): + x + make(np.ones(4), 0.0) + with pytest.raises(ValueError, match="1-d"): + xp.meshgrid(x) + with pytest.raises(ValueError, match="one data type"): + xp.meshgrid(xp.zeros(2), xp.zeros(2, dtype=xp.int8)) + with pytest.raises(ValueError, match="indexing"): + xp.meshgrid(xp.zeros(2), indexing="yx") + + +def test_creation_errors(): + with pytest.raises(TypeError, match="Python scalar"): + xp.full(3, "a") + with pytest.raises(ValueError, match="keeps its fill"): + xp.asarray(xp.zeros(2), fill_value=1.0) + assert_array_equal(xp.asarray(np.array([1, 2], np.float16)), np.array([1, 2], np.float32)) + assert_array_equal(xp.full(2, np.float32(1.5)), np.full(2, 1.5, np.float32)) + assert_array_equal(xp.full(2, xp.asarray(3, dtype=xp.int8)), np.full(2, 3, np.int8)) + assert_array_equal(xp.full_like(xp.zeros(2), xp.asarray(4.0)), np.full(2, 4.0)) + assert_array_equal(xp.from_coo(([1, 3],), [1, 2]), np.array([0, 1, 0, 2])) + assert_array_equal(xp.from_coo((), 5, ()), np.array(5)) + assert_array_equal(xp.from_coo((), [], (), dtype=xp.int8), np.array(0, np.int8)) + + +def test_sparse_input_errors(): + sp = pytest.importorskip("scipy.sparse") + A = sp.coo_array((np.array([1.0, 2.0]), (np.array([0, 0]), np.array([1, 1]))), shape=(2, 2)) + assert_array_equal(xp.asarray(A), np.array([[0.0, 3.0], [0.0, 0.0]])) # duplicates summed + with pytest.raises(ValueError, match="fill value 0"): + xp.asarray(A, fill_value=1.0) + sparse = pytest.importorskip("sparse") + s = sparse.GCXS.from_numpy(np.array([[1.0, 0.0], [0.0, 3.0]])) + assert_array_equal(xp.asarray(s), np.array([[1.0, 0.0], [0.0, 3.0]])) + with pytest.raises(ValueError, match="differs"): + xp.asarray(s, fill_value=1.0) + + +def test_dtype_errors(): + with pytest.raises(TypeError, match="not a data type"): + xp.asarray([1], dtype="not a dtype") + assert xp.isdtype(np.float32, "real floating") + with pytest.raises(ValueError, match="at least one"): + xp.result_type() + info = xp.__array_namespace_info__() + with pytest.raises(ValueError, match="kind"): + info.dtypes(kind="integer") + with pytest.raises(ValueError, match="device"): + info.default_dtypes(device="gpu") + + +def test_object_protocol_errors(): + x = make(np.array([1.0, 0.0]), 0.0) + with pytest.raises(BufferError, match="copies"): + x.__dlpack__(copy=False) + assert x.__dlpack_device__() == np.empty(0).__dlpack_device__() + np.testing.assert_array_equal(np.from_dlpack(x), [1.0, 0.0]) + with pytest.raises(TypeError, match="complex"): + int(xp.asarray(1j)) + with pytest.raises(TypeError): + x @ "a" + with pytest.raises(TypeError): + "a" @ x + with pytest.raises(TypeError): + x + "a" + with pytest.raises(TypeError): + "a" + x + np.testing.assert_array_equal(np.from_dlpack(xp.from_dlpack(x)), [1.0, 0.0]) + capsule = x.__dlpack__(dl_device=x.__dlpack_device__(), max_version=(1, 0)) + assert capsule is not None + with pytest.raises(TypeError, match="None"): + xp.astype(x, None) + + +def test_clip_mixed_bounds(): + a = np.array([1, 5, 0, 9], np.int16) + x = make(a, 0) + lo = make(np.array([2, 2, 2, 2], np.int32), 2) + out = xp.clip(x, min=lo, max=np.int16(6)) + assert out.dtype == xp.int16 + assert_array_equal(out, np.minimum(np.maximum(a, 2), 6).astype(np.int16)) + assert_array_equal(xp.clip(x, max=4.5), np.clip(a, None, 4.5).astype(np.int16)) + assert_array_equal(xp.clip(x, min=xp.asarray(np.int16(3))), np.clip(a, 3, None)) + hi = make(np.array([4, 4, 4, 4], np.int16), 4) + assert_array_equal(xp.clip(x, max=hi), np.minimum(a, 4)) + + +def test_complex_comparisons_mixed_precision(): + a = np.array([1 + 1j, 0, 2], np.complex64) + b = np.array([1 + 1j, 3, 0], np.complex128) + assert_array_equal(xp.equal(make(a, 0), make(b, 0)), a == b) + assert_array_equal(xp.not_equal(make(a, 0), make(b, 0)), a != b) + + +def test_bool_products(): + a = np.array([[True, False, True], [True, True, True]]) + x = make(a, False) + for axis in [None, 0, 1]: + assert_array_equal(xp.prod(x, axis=axis, dtype=xp.bool), np.prod(a, axis=axis, dtype=bool)) + y = make(a, True) + assert_array_equal(xp.prod(y, axis=1, dtype=xp.bool), np.prod(a, axis=1, dtype=bool)) + + +def test_statistical_errors(): + with pytest.raises(ValueError, match="at least one dimension"): + xp.cumulative_sum(xp.asarray(1.0)) + with pytest.raises(ValueError, match="at least one dimension"): + xp.diff(xp.asarray(1.0)) + with pytest.raises(ValueError, match="nonnegative"): + xp.diff(xp.zeros(3), n=-1) + + +def test_concat_errors(): + from graphblas.array_api._manipulation import _broadcast_to + + a = make(np.array([[1.0, 0.0]]), 0.0) + b = make(np.array([[2.0, 2.0]]), 2.0) + assert_array_equal(xp.concat([a, b]), np.array([[1.0, 0.0], [2.0, 2.0]])) + assert_array_equal(xp.concat([xp.zeros((1, 2)), xp.zeros((2, 2))]), np.zeros((3, 2))) + with pytest.raises(ValueError, match="at least one"): + xp.concat([]) + with pytest.raises(ValueError, match="at least 1"): + xp.concat([xp.asarray(1.0)]) + with pytest.raises(ValueError, match="do not concatenate"): + xp.concat([a, xp.zeros((1, 3))]) + with pytest.raises(ValueError, match="broadcast"): + _broadcast_to(a, (3, 3)) + with pytest.raises(ValueError, match="cannot broadcast"): + _broadcast_to(a, (2,)) + + +def test_map_values_storage_cases(): + T = Tensor.from_coo([[0, 1], [1, 2]], [1.0, 2.0], shape=(2, 3)) + U = T.relayout(((1,), (0,))) + out = _utils.map_values(np.add, [T, U], np.dtype(float)) + np.testing.assert_array_equal(out.to_dense(fill_value=0), [[0.0, 2.0, 0.0], [0.0, 0.0, 4.0]]) + out = _utils.map_values(np.add, [T, T.dup()], np.dtype(float)) # (one Matrix, shared) + np.testing.assert_array_equal(out.to_coo()[1], [2.0, 4.0]) + E = Tensor(T.dtype, (2, 3)) + assert _utils.map_values(np.sin, [E], np.dtype(float)).nvals == 0 + assert _utils.full_tensor(np.dtype(float), (0, 3), 1.0).nvals == 0 + + +def test_scalar_cache_is_bounded(): + for i in range(1100): + _utils.gb_scalar(i, np.dtype(np.int64)) + assert len(_utils._SCALARS) <= 1025 + assert _utils.gb_scalar(-0.0, np.dtype(float)).value == 0.0 + + +def test_facts_of_stored_values(): + x = make(np.array([1 + 0j, complex(-0.0, 1), complex(np.nan, 0)]), 0j) + assert [_utils.fact(x, k) for k in ["finite", "nan", "zero"]] == [False, True, False] + assert x._facts == {"finite": False, "nan": True, "zero": False} + y = make(np.array([-0.0, np.inf, 2.0]), 1.0) + assert [_utils.fact(y, k) for k in ["finite", "nan", "zero"]] == [False, False, True] + # NaN and finiteness exclude each other: one is known from the other + z = make(np.array([np.nan, 2.0]), 0.0) + assert _utils.fact(z, "nan") + assert not _utils.fact(z, "finite") + w = make(np.array([3.0, 2.0]), 0.0) + assert _utils.fact(w, "finite") + assert not _utils.fact(w, "nan") + # integers are finite without a look; empty arrays hold nothing + i = make(np.array([0, 5]), 1) + assert _utils.fact(i, "finite") + assert not _utils.fact(i, "nan") + assert _utils.fact(i, "zero") + e = xp.zeros(4) + assert [_utils.fact(e, k) for k in ["finite", "nan", "zero"]] == [True, False, False] + # infinities by sign; known absent from finiteness, and for integers + u = make(np.array([np.inf, 2.0]), 0.0) + assert [_utils.fact(u, k) for k in ["posinf", "neginf"]] == [True, False] + assert _utils.fact(make(np.array([-np.inf]), 0.0), "neginf") + assert not _utils.fact(w, "posinf") + assert not _utils.fact(i, "neginf") + # abs, negative, positive and conj keep all three; roundings keep two + v = make(np.array([-1.5, np.nan, 0.0]), 0.0) + for k in ["finite", "nan", "zero"]: + _utils.fact(v, k) + for func in [xp.abs, xp.negative, xp.positive, xp.conj]: + assert func(v)._facts == v._facts + for func in [xp.floor, xp.ceil, xp.trunc, xp.round]: + assert func(v)._facts == {"finite": False, "nan": True} + assert xp.exp(v)._facts is None + assert xp.abs(xp.zeros(3))._facts is None + # (review finding M4) |z| of a finite complex z may be inf: abs keeps only "zero" there + z = make(np.array([1.5e308 + 1.5e308j, 1, 0, 0]), 0j) + assert _utils.fact(z, "finite") + _utils.fact(z, "zero") + assert set(xp.abs(z)._facts) == {"zero"} + zz = make(np.array([0.0, 0.0, 2.0, 3.0]), 0.0) + with np.errstate(over="ignore", invalid="ignore"): + assert_array_equal(xp.abs(z) * zz, np.abs(z.todense()) * zz.todense()) + + +def test_thread3_argument_checks(): + x = make(np.arange(6.0).reshape(2, 3), 0.0) + for func in [xp.take, xp.take_along_axis]: + with pytest.raises(TypeError, match="Array"): + func(np.zeros(3), np.array([0])) + with pytest.raises(TypeError, match="integer"): + xp.take(x, np.array([0.5]), axis=0) + with pytest.raises(ValueError, match="dimension"): + xp.take(xp.asarray(1.0), np.array([0])) + with pytest.raises(TypeError, match="Arrays"): + xp.matmul(x, [[1.0], [2.0], [3.0]]) + with pytest.raises(ValueError, match="nonnegative"): + xp.tensordot(x, x, axes=-1) + with pytest.raises(ValueError, match="more axes"): + xp.tensordot(x, x, axes=3) + with pytest.raises(ValueError, match="as many"): + xp.tensordot(x, x, axes=([0], [0, 1])) + with pytest.raises(ValueError, match="dimension"): + xp.vecdot(xp.asarray(1.0), x) + with pytest.raises(ValueError, match="differs"): + xp.vecdot(x, make(np.ones((2, 2)), 0.0)) + assert_array_equal(xp.vecdot(x, x, axis=1), np.vecdot(x.todense(), x.todense(), axis=1)) + with pytest.raises(ValueError, match="repeated"): + xp.flip(x, axis=(0, 0)) + with pytest.raises(TypeError, match="Python scalars"): + xp.where(x > 0, x, "a") + with pytest.raises(TypeError, match="Array"): + xp.argmax(np.zeros(3)) + with pytest.raises(ValueError, match="negative"): + xp.linalg.cross(xp.zeros(3), xp.zeros(3), axis=0) + with pytest.raises(TypeError, match="Array"): + xp.linalg.vector_norm(np.zeros(3)) + a = np.array([[1.0, 2.0], [3.0, 4.0]]) + rtol = xp.asarray(1e-3) + assert_array_equal(xp.linalg.pinv(make(a, 0.0), rtol=rtol), np.linalg.pinv(a, rtol=1e-3)) + assert int(xp.linalg.matrix_rank(make(a, 0.0), rtol=rtol)) == 2 + assert_array_equal(x @ x.mT, x.todense() @ x.todense().T) + + +def test_thread3_edge_cases(): + x = make(np.arange(6.0).reshape(2, 3), 0.0) + assert_array_equal(x[np.array([], dtype=float)], np.zeros((0, 3))) + assert_array_equal(x[xp.asarray(True)], x.todense()[None]) + assert_array_equal(x[np.array(False)], x.todense()[np.array(False)]) + y = make(np.arange(6.0).reshape(2, 3), 0.0) + y[:, 0:0] = 5.0 + y[[0, 1], 0:0] = xp.zeros((2, 0)) + y[[], 1] = 3.0 + assert_array_equal(y, x.todense()) + e = xp.zeros((2, 3)) + e[0] = 0.0 # (deleting from an array that stores nothing) + assert e.nnz == 0 + assert_array_equal( + xp.take_along_axis(xp.zeros((2, 3)), np.zeros((2, 1), int), axis=1), np.zeros((2, 1)) + ) + assert_array_equal(xp.nonzero(xp.zeros(4))[0], np.zeros(0, np.int64)) + assert_array_equal(xp.sort(xp.zeros(3)), np.zeros(3)) + assert_array_equal(xp.sort(make(np.array([0.0, 2.0, 0.0]), 2.0)), np.array([0.0, 0.0, 2.0])) + z = make(np.array([[2.0, 2.0], [2.0, 2.0]]), 2.0, np.random.default_rng(0), explicit=1.0) + assert_array_equal(xp.sort(z, axis=0), np.full((2, 2), 2.0)) + full = make(np.array([[1.0, 2.0], [3.0, 1.0]]), 0.0) + assert_array_equal(xp.unique_inverse(full).inverse_indices, np.array([[0, 1], [2, 0]])) + assert_array_equal(xp.unique_values(xp.zeros((2, 0))), np.zeros(0)) + assert_array_equal( + xp.unique_inverse(xp.zeros((2, 0))).inverse_indices, np.zeros((2, 0), np.int64) + ) + assert_array_equal(xp.isin(full, full), np.ones((2, 2), bool)) + with np.errstate(all="ignore"): + assert_array_equal(xp.var(xp.zeros((0, 3)), axis=0, keepdims=True), np.full((1, 3), np.nan)) + # tensordot with a nonzero fill on a 0-d result, and int32 * int8 promotion + a = make(np.array([1.0, 0.0, 2.0]), 1.0) + assert_array_equal(xp.tensordot(a, a, axes=1), np.tensordot(a.todense(), a.todense(), axes=1)) + i32 = make(np.array([[1, 2], [3, 4]], np.int32), 0) + i8 = make(np.array([[1, 0], [0, 1]], np.int8), 0) + assert_array_equal( + xp.tensordot(i32, i8, axes=1), np.tensordot(i32.todense(), i8.todense(), axes=1) + ) + inf = make(np.array([[np.inf, 1.0], [2.0, np.inf]]), np.inf) + assert_array_equal( + xp.tensordot(inf, i32, axes=1), np.tensordot(inf.todense(), i32.todense(), axes=1) + ) + # both operands broadcast: an outer product of the entries, under the limit + u, v = make(np.ones((3, 1, 4)), 0.0), make(np.ones((5, 4)), 0.0) + assert_array_equal(xp.vecdot(u, v), np.vecdot(u.todense(), v.todense())) + with gb.config.set(array_api_densify_limit=4), pytest.raises(xp.DensifyError): + xp.vecdot(u, v) + # numpy operands + c = np.array([True, False, True]) + assert_array_equal(xp.where(c, np.array([1, 2, 3]), 0), np.where(c, [1, 2, 3], 0)) + assert_array_equal(xp.where(True, x, 0.0), x.todense()) + assert_array_equal(np.eye(3) @ x.mT, np.eye(3) @ x.todense().T) + + +def test_zero_d_tensors_are_copies_of_one_empty_tensor(): + # (a 0-d Array stores nothing; its Tensor is an O(1) copy-on-write copy of one per dtype, + # which no change to a copy reaches) + x = xp.asarray(1.5) + y = xp.multiply(x, 2.0) + assert float(y) == 3.0 + t = x.tensor + t[()] = 7.0 + assert t.nvals == 1 + template = _utils._EMPTY_0D[_utils.gb_dtype(np.float64)] + assert template.nvals == 0 + assert xp.asarray(2.5)._T().nvals == 0 + assert float(xp.asarray(2.5)) == 2.5 + assert _utils.empty_tensor(np.int8, ()).nvals == 0 diff --git a/graphblas/tests/test_array_api_indexing.py b/graphblas/tests/test_array_api_indexing.py new file mode 100644 index 000000000..0a3454c3d --- /dev/null +++ b/graphblas/tests/test_array_api_indexing.py @@ -0,0 +1,305 @@ +"""graphblas.array_api indexing (get and set), take and take_along_axis against numpy.""" + +import itertools +import math + +import numpy as np +import pytest + +import graphblas.array_api as xp + +from .test_array_api_elementwise import assert_array_equal, inputs, make + +rng = np.random.default_rng(0) +A = np.arange(60).reshape(3, 4, 5) +MASK = (A % 7) < 3 + + +def _keys(): + return [ + [0, 2], + np.array([2, -1, 0, 0]), + ([1, 0], [3, 3]), + ([[0], [2]], [1, 0, 1], [4, 4, 0]), + (2, [1, 3]), + ([1], 2, [0, 4]), + (0, Ellipsis, [4, 1]), + ([0, 1], slice(None), [1, 3]), + (slice(None), [2, 0], [1, 1]), + (slice(1, None), [3, 0, 3]), + (slice(None, None, -1), slice(1, 3), np.array([[0, 1], [2, 3]])), + (None, [1], slice(None), 2), + ([2, 1], None, 0), + (Ellipsis, [3, 0]), + MASK, + MASK[:, :, 0], + (slice(None), MASK[0]), + (MASK[:, 0, :].any(axis=1),), + (MASK[:, 0, 0], 1), + True, + False, + (np.array([], dtype=np.int64),), + (np.zeros((3, 4, 5), bool),), + (Ellipsis, None), + ] + + +@pytest.mark.parametrize("key", range(len(_keys()))) +def test_getitem_advanced(key): + key = _keys()[key] + for dtype in [np.dtype("float64"), np.dtype("int16"), np.dtype("bool"), np.dtype("complex64")]: + for dense, x in inputs(dtype, (3, 4, 5), seed=3): + assert_array_equal(x[key], dense[key]) + # index arrays given as Arrays + if ( + isinstance(key, np.ndarray) + or isinstance(key, tuple) + and any(isinstance(k, np.ndarray) for k in key) + ): + k2 = key if not isinstance(key, np.ndarray) else (key,) + k2 = tuple(xp.asarray(k) if isinstance(k, np.ndarray) else k for k in k2) + assert_array_equal(x[k2], dense[key]) + + +def test_boolean_mask_from_comparison(): + for fill in [0.0, 1.5]: + a = np.where(rng.random((4, 5)) < 0.5, fill, np.round(rng.normal(size=(4, 5)), 1)) + x = make(a, fill) + for mask in [x > 0, x < 0.5, x == fill, xp.isnan(x)]: + assert_array_equal(x[mask], a[mask.todense()]) + assert_array_equal(x[x > 100], a[a > 100]) + big = xp.zeros((10**5, 10**5)) + mask = big != 0 + assert big[mask].shape == (0,) + + +def test_getitem_errors(): + x = make(A.astype(float), 0.0) + with pytest.raises(IndexError): + x[[0, 3]] + with pytest.raises(IndexError): + x[[0], [-5]] + with pytest.raises(IndexError): + x[MASK[:2]] + with pytest.raises(IndexError): + xp.zeros((0, 0))[np.zeros((0, 1), bool)] + assert xp.zeros((2, 3))[np.zeros((0, 3), bool)].shape == (0,) + with pytest.raises(IndexError): + x[np.array([0.5])] + with pytest.raises(IndexError): + x[xp.asarray([0.5])] + with pytest.raises(IndexError): + x["a"] + with pytest.raises(IndexError): + x[..., 0, ...] + with pytest.raises(IndexError): + x[MASK, 0] + with pytest.raises(IndexError): + x[xp.asarray(1.5)] + + +SET_CASES = [ + ((0, 1), 9.0), + ((0, [1, 2]), 9.0), + (([0, 1], [0, 2]), 7.0), + (([0, 1, 0], [0, 2, 0], [1, 1, 1]), np.array([1.0, 2.0, 3.0])), + ((slice(None), [0, 0], 1), np.array([5.0, 6.0])), + (([1], slice(1, 3)), 0.0), + ((Ellipsis, 1), -1.0), + ((1,), np.arange(20.0).reshape(4, 5)), + ((slice(None), slice(1, 3)), np.array([1.0, 0.0, 2.0, 0.0, 3.0])), + ((slice(None, None, -1), 2), 0.0), + (([2, 0, 2],), np.arange(20.0).reshape(4, 5) * 0), + (MASK, 0.0), + (MASK, 4.0), + (MASK[:, :, 0], np.array([1.0, 0.0, 2.0, 0.0, 3.0])), + (([0, 2], [1, 3]), np.array([[1.0], [2.0]])), + ((), 3.5), + ((Ellipsis,), np.zeros(5)), + ((None, 0), 1.0), +] + + +@pytest.mark.parametrize("case", range(len(SET_CASES))) +def test_setitem(case): + key, value = SET_CASES[case] + for fill in [0.0, 2.5, -0.0]: + dense = np.where(A % 3 == 0, fill, A.astype(float)) + x = make(dense, fill) + expected = dense.copy() + expected[key] = value + for v in [value, make(np.asarray(value), 0.0), make(np.asarray(value), fill)]: + y = make(dense, fill) + y[key] = v + assert_array_equal(y, expected) + assert y.fill_value == x.fill_value + y = make(dense, fill) + y[MASK.tolist()] = 1.0 # (a nested list is an index array too) + expected = dense.copy() + expected[MASK] = 1.0 + assert_array_equal(y, expected) + + +def test_setitem_dtypes_and_errors(): + x = make(np.array([1, 0, 3], np.int8), 0) + x[0] = 1.9 + assert_array_equal(x, np.array([1, 0, 3], np.int8)) + with pytest.raises(OverflowError): + x[1] = 300 + with pytest.raises(TypeError, match="complex"): + x[1] = 1j + with pytest.raises(ValueError, match="cannot assign"): + x[:2] = xp.asarray([1, 2, 3]) + with pytest.raises(ValueError, match="cannot assign"): + x[[0, 1]] = xp.asarray([1, 2, 3]) + y = make(np.zeros((2, 2)), 0.0) + with pytest.raises(ValueError, match="cannot assign"): + y[[0, 1], [0, 1]] = np.ones(3) + z = xp.asarray(1.5) + z[False] = 3.0 + assert_array_equal(z, np.array(1.5)) + z[True] = 2.5 + assert_array_equal(z, np.array(2.5)) + q = make(np.array([1.0, 2.0]), 0.0) + q[False] = 3.0 + assert_array_equal(q, np.array([1.0, 2.0])) + z[()] = 2.0 + assert_array_equal(z, np.array(2.0)) + with pytest.raises(ValueError, match="0-d"): + z[...] = xp.asarray([1.0]) + # the cached summary of the values does not survive an assignment + w = make(np.array([1.0, 0.0, 2.0]), 0.0) + assert_array_equal(w * make(np.array([0.0, 3.0, 0.0]), 0.0), np.zeros(3)) + w[0] = -1.0 + assert_array_equal(w * make(np.array([0.0, 3.0, 0.0]), 0.0), np.array([-0.0, 0.0, 0.0])) + # an assignment does not change another array that shares its data + a = make(np.array([1.0, 0.0, 2.0]), 0.0) + b = xp.asarray(a, copy=True) + c = a[:] + a[[0, 2]] = 5.0 + assert_array_equal(b, np.array([1.0, 0.0, 2.0])) + assert_array_equal(c, np.array([1.0, 0.0, 2.0])) + + +@pytest.mark.parametrize("axis", [0, 1, -1, None]) +def test_take(axis): + for dtype in [np.dtype("float64"), np.dtype("uint8"), np.dtype("bool")]: + shape = (7,) if axis is None else (3, 4) + for dense, x in inputs(dtype, shape, seed=2): + n = dense.shape[0 if axis is None else axis] + for idx in [np.array([0, n - 1, 1, 1]), np.array([-1, 0]), np.array([], np.int64)]: + expected = np.take(dense, idx, axis=axis) + assert_array_equal(xp.take(x, xp.asarray(idx), axis=axis), expected) + assert_array_equal(xp.take(x, idx, axis=axis), expected) + with pytest.raises(IndexError): + xp.take(x, xp.asarray([n]), axis=axis) + with pytest.raises(ValueError, match="axis"): + xp.take(xp.zeros((2, 2)), xp.asarray([0])) + with pytest.raises(TypeError, match="integer"): + xp.take(xp.zeros(2), xp.asarray([0.0])) + + +@pytest.mark.parametrize("axis", [0, 1, -1]) +def test_take_along_axis(axis): + for dense, x in inputs(np.dtype("float32"), (3, 4), seed=5): + n = dense.shape[axis] + idx = np.argsort(dense, axis=axis) + assert_array_equal( + xp.take_along_axis(x, xp.asarray(idx), axis=axis), + np.take_along_axis(dense, idx, axis=axis), + ) + shape = [3, 4] + shape[axis] = 2 + idx = rng.integers(-n, n, shape) + assert_array_equal( + xp.take_along_axis(x, xp.asarray(idx), axis=axis), + np.take_along_axis(dense, idx, axis=axis), + ) + bshape = [1, 1] + bshape[axis] = 5 + idx = rng.integers(0, n, bshape) + assert_array_equal( + xp.take_along_axis(x, idx, axis=axis), np.take_along_axis(dense, idx, axis=axis) + ) + with pytest.raises(ValueError, match="dimensions"): + xp.take_along_axis(xp.zeros((2, 2)), xp.asarray([0])) + + +def _index_patterns(): + """Keys mixing ints, slices, integer arrays, ``...`` and ``None`` with at least one + integer array, as numpy accepts them for a 4-d array. + """ + parts = [slice(None), 1, np.array([0, 2]), np.array([[1], [0]]), Ellipsis, None, slice(0, 2)] + for n in range(1, 5): + for key in itertools.product(parts, repeat=n): + if not any(isinstance(k, np.ndarray) for k in key): + continue + if sum(k is Ellipsis for k in key) > 1: + continue + yield key + + +def test_mixed_keys_order_axes_as_numpy(): + # (shim finding B1) the advanced axes go first when the advanced indices are apart, a + # "..." that covers no axis included; in place when they are next to each other + count = 0 + for a, key in itertools.chain( + ((np.arange(36.0).reshape(3, 4, 3) + 1, k) for k in _index_patterns()), + ((np.arange(108.0).reshape(3, 4, 3, 3) + 1, k) for k in _index_patterns()), + ): + x = make(a, 0.0) + try: + expected = a[key] + except IndexError: + continue + count += 1 + assert_array_equal(x[key], expected) + if any(k is None for k in key): + continue # (a value for a key with new axes is not read with them yet) + b = a.copy() + value = np.arange(expected.size, dtype=np.float64).reshape(expected.shape) - 50 + b[key] = value + y = make(a, 0.0) + y[key] = make(value, 0.0) if value.ndim else xp.asarray(value) + assert_array_equal(y, b) + assert count > 2000 + x = make(np.arange(24.0).reshape(2, 3, 4), 0.0) + expected = np.array([[4.0, 16.0], [6.0, 18.0]]) + for index in [[0, 2], np.array([0, 2]), xp.asarray([0, 2])]: + assert_array_equal(x[:, 1, ..., index], expected) + y = make(np.arange(24.0).reshape(2, 3, 4), 0.0) + y[:, 1, ..., index] = make(-expected, 0.0) + b = np.arange(24.0).reshape(2, 3, 4) + b[:, 1, ..., [0, 2]] = -expected + assert_array_equal(y, b) + + +def test_setitem_with_new_axes(): + # (review finding L3) the value of a basic key with None is shaped like numpy's selection, + # new axes included + parts = [slice(None), 1, slice(0, 2), Ellipsis, None] + count = 0 + for a in [np.arange(12.0).reshape(3, 4) + 1, np.arange(24.0).reshape(2, 3, 4) + 1]: + for n in range(1, 5): + for key in itertools.product(parts, repeat=n): + if None not in key or sum(k is Ellipsis for k in key) > 1: + continue + try: + shape = a[key].shape + except IndexError: + continue + for vshape in [shape, shape[1:], (1,) * len(shape)]: + value = np.arange(math.prod(vshape), dtype=np.float64).reshape(vshape) - 50 + b = a.copy() + b[key] = value + y = make(a, 0.0) + y[key] = make(value, 0.0) if value.ndim else xp.asarray(value) + assert_array_equal(y, b) + count += 1 + assert count > 300 + z = xp.asarray(1.0) + z[None] = xp.asarray([5.0]) + assert_array_equal(z, np.array(5.0)) + x = make(np.ones((3, 4)), 0.0) + with pytest.raises(ValueError, match="new axes"): + x[None] = xp.ones((2, 3, 4)) diff --git a/graphblas/tests/test_array_api_lazy.py b/graphblas/tests/test_array_api_lazy.py new file mode 100644 index 000000000..c7564c117 --- /dev/null +++ b/graphblas/tests/test_array_api_lazy.py @@ -0,0 +1,672 @@ +"""graphblas.array_api: deferred contractions and their two push-down consumers (DESIGN D9).""" + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp +from graphblas.array_api import _lazy, _utils + +from .test_array_api_elementwise import assert_array_equal, check_invariants, make + +rng = np.random.default_rng(9) + + +def _rand(shape, fill=0.0, dtype=np.float64, density=0.5): + dtype = np.dtype(dtype) + if dtype.kind == "b": + a = rng.random(shape) < 0.5 + elif dtype.kind in "iu": + a = rng.integers(1, 9, shape) + else: + a = np.round(rng.normal(size=shape) * 4, 1) + return np.where(rng.random(shape) < density, a, fill).astype(dtype) + + +def assert_close(x, expected): + expected = np.asarray(expected) + check_invariants(x) + assert x.dtype == expected.dtype, (x.dtype, expected.dtype) + assert x.shape == expected.shape, (x.shape, expected.shape) + got = x.todense() + if got.dtype.kind in "fc": + tol = 1e-5 if got.dtype in (np.float32, np.complex64) else 1e-9 + np.testing.assert_allclose(got, expected, rtol=tol, atol=tol) + else: + np.testing.assert_array_equal(got, expected) + + +def pending(x): + return x._tensor is None + + +class _Seen(list): + hints = None + + +@pytest.fixture +def masks(monkeypatch): + """The masks the pending products were computed under (None: whole).""" + seen = _Seen() + compute = _lazy.Contraction.compute + + def recording(self, mask=None, **opts): + seen.append(mask) + hints.append(opts) + return compute(self, mask, **opts) + + hints = [] + seen.hints = hints # (the options each computation got) + + monkeypatch.setattr(_lazy.Contraction, "compute", recording) + return seen + + +def test_contractions_are_pending(masks): + a, b = _rand((4, 5)), _rand((5, 3)) + x1, x2 = make(a, 0.0), make(b, 0.0) + for P, expected in [ + (x1 @ x2, a @ b), + (xp.matmul(x1, x2), a @ b), + (xp.tensordot(x1, x2, axes=1), a @ b), + (xp.einsum("ij,jk->ik", x1, x2), a @ b), + (xp.vecdot(x1, x1), np.vecdot(a, a)), + (xp.matmul(x1, x2, semiring="lor_land"), (a != 0).astype(int) @ (b != 0) > 0), + ]: + assert pending(P) + # Metadata and repr do not compute it + assert P.shape == expected.shape + assert P.ndim == expected.ndim + assert P.size == expected.size + assert P.dtype == expected.dtype + assert P.fill_value == 0 + assert P.device == "cpu" + assert repr(P) == ( + f"Array(shape={expected.shape}, dtype={expected.dtype}, " + f"fill_value={P.fill_value.item()!r}, pending)" + ) + assert pending(P) + assert_close(P, expected) # (computes it) + assert not pending(P) + assert P._pending is None # (the snapshots are released) + T = P._T() + assert P._T() is T # (computed once) + assert masks == [None] * 6 + # Not deferred: 0-d results, empty operands, the kill switch, nonzero fills + assert not pending(xp.vecdot(make(a[0], 0.0), make(a[1], 0.0))) + assert not pending(xp.einsum("ij,ij->", x1, x1)) + assert not pending(x1 @ xp.zeros((5, 3))) + assert not pending(x1 @ make(np.where(b == 0, np.inf, b), np.inf)) # (dense: inf shows) + assert pending(x1 @ make(b, 1.0)) # (every position stored: its fill is taken as 0) + assert pending(x1 @ make(np.where(b == 0, 1.0, b), 1.0)) # (small: stored in full) + with gb.config.set(), pytest.MonkeyPatch.context() as mp: + mp.setattr(_lazy, "enabled", False) + P = x1 @ x2 + assert not pending(P) + assert_close(P, a @ b) + + +def test_snapshots(masks): + a, b = _rand((4, 5)), _rand((5, 3)) + x1, x2 = make(a, 0.0), make(b, 0.0) + P = x1 @ x2 + # The snapshots share the operands' data (copy-on-write: O(1)) + T1, T2 = P._pending.tensors + assert T1._M is x1._T()._M + assert T2._M is x2._T()._M + # Changing an operand afterwards changes neither the pending result nor the operand's + # earlier readers + x1[0, :] = 7.0 + x1[1, 1] = 0.0 + x2 += 1.0 + assert T1._M is not x1._T()._M + assert_close(P, a @ b) + a[0, :] = 7.0 + a[1, 1] = 0.0 + assert_close(x1 @ x2, a @ (b + 1.0)) + # The same Array on both sides + P = x1 @ x1.T + x1[...] = 0.0 + assert_close(P, a @ a.T) + # __setitem__ on a pending Array computes it first + P = make(a, 0.0) @ make(b, 0.0) + P[0, 0] = -1.0 + expected = a @ b + expected[0, 0] = -1.0 + assert_close(P, expected) + P = make(a, 0.0) @ make(b, 0.0) + P += 1.0 + assert_close(P, a @ b + 1.0) + Q = make(a[:, :4], 0.0) + Q @= make(2 * np.eye(4), 0.0) + assert not pending(Q) + assert_close(Q, 2 * a[:, :4]) + + +def test_pending_operands_and_consumers(): + a, b, c = _rand((4, 5)), _rand((5, 3)), _rand((3, 2)) + x1, x2, x3 = make(a, 0.0), make(b, 0.0), make(c, 0.0) + # A pending operand is computed (one level of deferral) + P = x1 @ x2 + R = P @ x3 + assert not pending(P) + assert pending(R) + assert_close(R, a @ b @ c) + # Everything else computes it + for func, np_func in [ + (lambda p: p + 1, lambda p: p + 1), + (xp.sum, np.sum), + (lambda p: p.T, lambda p: p.T), + (lambda p: p[[0, 2]], lambda p: p[[0, 2]]), + (lambda p: xp.asarray(np.asarray(p)), lambda p: p), + (lambda p: xp.stack(list(p)), lambda p: p), + ]: + P = x1 @ x2 + assert_close(func(P), np_func(a @ b)) + assert not pending(P) + # (the stored entries: where some term meets) + structure = (a != 0).astype(int) @ (b != 0).astype(int) > 0 + P = x1 @ x2 + assert P.nnz == structure.sum() + assert not pending(P) + P = x1 @ x2 + assert_close(xp.asarray(P.tensor), a @ b) + assert not pending(P) + + +def _sddmm_inputs(n=12, k=3, dtype=np.float64): + S = _rand((n, n), 0.0, dtype, density=0.2) + A = _rand((n, k), 0.0, dtype, density=1.0) + B = _rand((n, k), 0.0, dtype, density=1.0) + return S, A, B + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32, np.int64, np.complex128]) +def test_sddmm(dtype, masks): + s, a, b = _sddmm_inputs(dtype=dtype) + S, A, B = make(s, 0), make(a, 0), make(b, 0) + expected = s * (a @ b.T) + for func in [ + lambda P: xp.multiply(S, P), + lambda P: xp.multiply(P, S), + lambda P: S * P, + lambda P: P * S, + ]: + P = xp.matmul(A, xp.matrix_transpose(B)) + masks.clear() + R = func(P) + assert len(masks) == 1 + assert masks[0] is not None # (computed under S's structure) + assert pending(P) # (the masked product is not kept) + assert_close(R, expected) + assert R.nnz <= S.nnz + # Other contractions, and S with explicit zeros + S = make(s, 0, rng, explicit=0.3) + for P in [ + xp.einsum("ik,jk->ij", A, B), + xp.tensordot(A, B, axes=([1], [1])), + xp.vecdot(A[:, None, :], B[None, :, :]), + ]: + masks.clear() + assert_close(xp.multiply(S, P), expected) + assert masks[0] is not None + # A pending S is computed, then serves as the mask + P = A @ xp.matrix_transpose(B) + Q = make(np.eye(a.shape[0], dtype=dtype), 0) @ S + assert_close(xp.multiply(Q, P), expected) + + +def test_sddmm_bool_and_semirings(masks): + s, a, b = _sddmm_inputs(dtype=np.bool_) + S, A, B = make(s, False), make(a, False), make(b, False) + masks.clear() + assert_close(xp.multiply(S, A @ xp.matrix_transpose(B)), s & (a @ b.T)) + assert masks[0] is not None + # A fill other than 0 is computed whole: min_plus's inf + s, a, b = _sddmm_inputs() + S = make(s, 0.0) + A, B = make(np.where(a == 0, np.inf, a), np.inf), make(np.where(b == 0, np.inf, b), np.inf) + masks.clear() + P = xp.matmul(A, xp.matrix_transpose(B), semiring="min_plus") + expected = np.min(A.todense()[:, None, :] + B.todense()[None, :, :], axis=2) + with np.errstate(invalid="ignore"): + assert_close(xp.multiply(S, P), s * expected) + assert masks == [None] + + +def test_sddmm_falls_back(masks): + s, a, b = _sddmm_inputs() + A, B = make(a, 0.0), make(b, 0.0) + expected = a @ b.T + # S's fill is not 0, S broadcasts, or S is a Python scalar: the product is computed whole + for S, s_np in [(make(s, 1.0), s), (make(s[:1], 0.0), s[:1]), (2.0, 2.0)]: + masks.clear() + P = A @ xp.matrix_transpose(B) + assert_close(xp.multiply(S, P), s_np * expected) + assert [m is None for m in masks] == [True] + # (an ndarray is an Array with fill 0) + masks.clear() + assert_close(xp.multiply(s, A @ xp.matrix_transpose(B)), s * expected) + assert [m is None for m in masks] == [False] + # An operand holding inf or NaN: exact, as numpy (0 * inf is NaN outside S) + for bad in [np.inf, np.nan]: + a2 = a.copy() + a2[0, 0] = bad + A2 = make(a2, 0.0) + masks.clear() + R = xp.multiply(make(s, 0.0), A2 @ xp.matrix_transpose(B)) + with np.errstate(invalid="ignore"): + assert_close(R, s * (a2 @ b.T)) + assert [m is None for m in masks] == [True] + + +def test_sddmm_overflow_outside_s(masks): + # A product that may overflow to inf outside S's entries is computed whole: 0 * inf there + # is NaN, as numpy has it + s = np.array([[1.0, 0.0], [0.0, 0.0]]) + a = np.full((2, 1), 1e200) + with np.errstate(invalid="ignore", over="ignore"): + expected = s * (a @ a.T) + assert np.isnan(expected[0, 1]) + for func in [ + lambda A, B: A @ xp.matrix_transpose(B), + lambda A, B: xp.einsum("ik,jk->ij", A, B), + lambda A, B: xp.vecdot(A[:, None, :], B[None, :, :]), + ]: + masks.clear() + R = xp.multiply(make(s, 0.0), func(make(a, 0.0), make(a, 0.0))) + assert_array_equal(R, expected) + assert [m is None for m in masks] == [True] + # The bound counts the terms of each sum (and keeps half the range for rounding): one + # term of 6e153 * 6e153 is 3.6e307, three may reach 1.1e308 + a = np.full((3, 3), 6e153) + for k, pushed in [(1, True), (3, False)]: + masks.clear() + A = make(a[:, :k], 0.0) + with np.errstate(over="ignore", invalid="ignore"): + expected = s[:1, :1] * (a[:1, :k] @ a[:1, :k].T) + R = xp.multiply(make(s[:1, :1], 0.0), A[:1] @ xp.matrix_transpose(A[:1])) + assert_array_equal(R, expected) + assert [m is not None for m in masks] == [pushed] + # Large values well within the bound, float32 too, and complex + for dtype, value in [(np.float64, 1e100), (np.float32, 1e15), (np.complex128, 1e100j)]: + a = np.full((3, 2), value, dtype=dtype) + s3 = np.eye(3, dtype=dtype) + masks.clear() + R = xp.multiply(make(s3, 0), make(a, 0) @ xp.matrix_transpose(make(a, 0))) + assert_array_equal(R, s3 * (a @ a.T)) + assert [m is not None for m in masks] == [True] + masks.clear() + a = np.full((3, 2), 1e20, dtype=np.float32) + with np.errstate(over="ignore", invalid="ignore"): + expected = np.eye(3, dtype=np.float32) * (a @ a.T) + R = xp.multiply(make(np.eye(3, dtype=np.float32), 0), make(a, 0) @ make(a.T, 0)) + assert_array_equal(R, expected) + assert [m is None for m in masks] == [True] + + +def test_sddmm_with_inf_stays_sparse(masks): + # (review finding H2) an operand holding inf, or values whose products may overflow, + # computes the product whole: a sparse product, as without deferral, with no densify + # check on its logical size + rng = np.random.default_rng(16) + n = 400 + r, c = rng.integers(0, n, 400), rng.integers(0, n, 400) + a = np.zeros((n, n)) + a[r, c] = rng.random(400) + 0.5 + a[r[0], c[0]] = np.inf + s = np.zeros((n, n)) + s[rng.integers(0, n, 400), rng.integers(0, n, 400)] = 1.0 + A = xp.from_coo(np.nonzero(a), a[np.nonzero(a)], (n, n)) + S = xp.from_coo(np.nonzero(s), s[np.nonzero(s)], (n, n)) + with gb.config.set(array_api_densify_limit=n * n // 10): + R = xp.multiply(S, A @ A) + # (the product as computed without deferral: D6 skips an unstored zero times inf) + P = A @ A + with np.errstate(invalid="ignore"): + assert_close(R, s * P.todense()) + assert [m is None for m in masks] == [True, True] + assert np.isinf(P.todense()).any() + + +def test_max_abs(): + assert _lazy._max_abs(xp.zeros(3)._T()) == 0.0 + assert _lazy._max_abs(make(np.array([-5, 3], np.int64), 0)._T()) == 5.0 + assert _lazy._max_abs(make(np.array([True, False]), False)._T()) == 1.0 + assert _lazy._max_abs(make(np.array([3 + 4j, 1j]), 0)._T()) == 5.0 + assert _lazy._max_abs(make(np.array([-2.5, np.nan, 1.0]), 0.0)._T()) == 2.5 + + +def test_finite_facts(): + x = make(np.array([1.0, np.inf]), 0.0) + assert not _lazy._all_finite(x._T()) + assert not _lazy._all_finite(make(np.array([np.nan, 1.0]), 0.0)._T()) + assert not _lazy._all_finite(make(np.array([np.inf, -np.inf]), 0.0)._T()) + assert not _lazy._all_finite(make(np.array([1 + 0j, complex(0, np.inf)]), 0.0)._T()) + assert _lazy._all_finite(make(np.array([1.0, -2.0]), 0.0)._T()) + assert _lazy._all_finite(make(np.array([5, 7]), 0)._T()) + assert _lazy._all_finite(xp.zeros(3)._T()) + # Facts an Array already knows are taken as they are + y = make(np.array([1.0, 2.0]), 0.0) + assert _lazy._known_finite(y) is None + _utils.fact(y, "finite") + assert _lazy._known_finite(y) is True + assert _lazy._known_finite(make(np.array([1, 2]), 0)) is True + P = xp.matmul(y, make(np.array([[1.0], [2.0]]), 0.0)) + assert P._pending.finite == [True, True] # (every contraction asks: DESIGN D12a) + + +SLICES = [ + (1,), + (slice(1, 3),), + (slice(None), 2), + (slice(None, None, -2), slice(1, 3)), + (Ellipsis, 1), + (None, slice(0, 2)), + (2, None), +] + + +@pytest.mark.parametrize("key", SLICES, ids=str) +def test_getitem_slices_the_operands(key, masks): + a, b = _rand((4, 5)), _rand((5, 3)) + x1, x2 = make(a, 0.0), make(b, 0.0) + expected = (a @ b)[key] + for P in [x1 @ x2, xp.tensordot(x1, x2, axes=1), xp.einsum("ij,jk->ik", x1, x2)]: + masks.clear() + R = P[key] + assert pending(P) + if R.ndim and not any(k is None for k in key): + assert pending(R) + assert_close(R, expected) + assert masks == [None] # (only the slice was computed) + assert_close(P, a @ b) + + +def test_getitem_cases(masks): + a, b = _rand((2, 4, 5)), _rand((2, 5, 3)) + x1, x2 = make(a, 0.0), make(b, 0.0) + # Batched (equal batch axes), 1-d operands, three operands + for P, expected in [ + (x1 @ x2, a @ b), + (x1[0, 0] @ x2[0], a[0, 0] @ b[0]), + (x1[0] @ x2[0, :, 0], a[0] @ b[0, :, 0]), + ( + xp.einsum("bij,bjk,kl->bil", x1, x2, make(b[0].T, 0.0)), + np.einsum("bij,bjk,kl->bil", a, b, b[0].T), + ), + (xp.tensordot(x1, x2, axes=([0, 2], [0, 1])), np.tensordot(a, b, ([0, 2], [0, 1]))), + ]: + key = ( + (1,) + (slice(1, None),) * (expected.ndim - 1) + if expected.ndim > 1 + else (slice(1, None),) + ) + masks.clear() + R = P[key] + assert pending(P) + assert_close(R, expected[key]) + # Broadcast batch axes: no slicing of the operands, computed whole + y1 = make(a[:1], 0.0) + P = y1 @ x2 + assert_close(P[1, 1:], (a[:1] @ b)[1, 1:]) + assert not pending(P) + # vecdot has no slicing; all-int keys, full slices and index arrays compute whole + for P, expected, key in [ + (xp.vecdot(x1, x1), np.vecdot(a, a), (0, slice(1, 3))), + (x1 @ x2, a @ b, (0, 1, 2)), + (x1 @ x2, a @ b, (slice(None),)), + (x1 @ x2, a @ b, (Ellipsis,)), + (x1 @ x2, a @ b, (np.array([0, 1]),)), + (x1[0, 0] @ x2[0], a[0, 0] @ b[0], (1,)), + ]: + assert pending(P) + assert_close(P[key], expected[key]) + assert not pending(P) + # Slices of empty extent, and slicing then the SDDMM + P = make(a[0], 0.0) @ make(b[0], 0.0) + assert_close(P[2:2], (a[0] @ b[0])[2:2]) + s = _rand((2, 3), 0.0, density=0.5) + masks.clear() + R = xp.multiply(make(s, 0.0), P[1:3]) + assert_close(R, s * (a[0] @ b[0])[1:3]) + assert [m is not None for m in masks] == [True] + + +def test_iteration_computes_once(masks): + a, b = _rand((4, 5)), _rand((5, 3)) + P = make(a, 0.0) @ make(b, 0.0) + rows = list(P) + assert masks == [None] + for row, expected in zip(rows, a @ b, strict=True): + assert_close(row, expected) + + +def test_sddmm_dot_products(masks): + # A @ B with B held by row (k x n) under a sparse mask: dot products at the mask's + # entries, not SuiteSparse's masked saxpy over the whole product (found by the SAPS run) + rng = np.random.default_rng(13) + n, k = 60, 4 + s = np.where(rng.random((n, n)) < 0.02, rng.random((n, n)), 0.0) + a, b = rng.random((n, k)) + 0.5, rng.random((k, n)) + 0.5 + S, A, B = make(s, 0.0), xp.asarray(a), xp.asarray(b) + for P, expected in [ + (A @ B, s * (a @ b)), + (xp.tensordot(A, B, axes=1), s * (a @ b)), + (xp.einsum("ik,kj->ij", A, B), s * (a @ b)), + (xp.einsum("ik,jk->ij", A, xp.asarray(b.T)), s * (a @ b)), + ((A @ B)[: n // 2], s[: n // 2] * (a @ b)[: n // 2]), + ]: + masks.clear() + masks.hints.clear() + S2 = S if expected.shape == s.shape else S[: n // 2] + assert_close(xp.multiply(S2, P), expected) + assert masks.hints == [{"axb_method": "dot"}] + # Not for a dense mask, a vecdot, or an einsum whose product is permuted after + for S2, P, expected in [ + (make(np.ones((n, n)), 0.0), A @ B, a @ b), + (S, xp.vecdot(A[:, None, :], xp.asarray(b.T)[None, :, :]), s * (a @ b)), + (S, xp.einsum("ik,kj->ji", A, B), s * (a @ b).T), + ]: + masks.clear() + masks.hints.clear() + assert_close(xp.multiply(S2, P), expected) + assert masks.hints == [{}] + + +@pytest.mark.parametrize( + "op", + ["iadd", "isub", "imul", "itruediv", "ifloordiv", "imod", "ipow", "imatmul", "iand", "ior", + "ixor", "ilshift", "irshift"], +) # fmt: skip +def test_inplace_operators_on_a_pending_array(op): + # (review finding H1) an in-place operator replaces a pending product: later consumers + # (the SDDMM, slicing) must read the new value, not the old product + import operator + + func = getattr(operator, op) + rng = np.random.default_rng(14) + dtype = np.int64 if op in {"iand", "ior", "ixor", "ilshift", "irshift"} else np.float64 + a = rng.integers(1, 5, (4, 4)).astype(dtype) + s = np.zeros((4, 4), dtype) + s[0, 0], s[1, 2], s[3, 3] = 2, 3, 1 + s2 = np.eye(4, dtype=dtype) + P = make(a, 0) @ make(a, 0) + assert pending(P) + with np.errstate(all="ignore"): + expected = func((a @ a).copy(), s if op != "imatmul" else s2) + P = func(P, make(s, 0) if op != "imatmul" else make(s2, 0)) + assert not pending(P) + assert P._pending is None + assert_close(P, expected) + assert_close(xp.multiply(make(s2, 0), P), s2 * expected) + assert_close(P[1:], expected[1:]) + + +def test_sddmm_dot_hint_counts_batches(masks): + # (review finding L9) a batched product pairs entries within a batch only: the saxpy + # estimate divides by the number of batches, so sparse batches keep SuiteSparse's choice + rng = np.random.default_rng(17) + for dens_ab, dens_s, hint in [(1.0, 0.02, {"axb_method": "dot"}), (0.02, 0.5, {})]: + a = np.where(rng.random((8, 16, 4)) < dens_ab, rng.random((8, 16, 4)) + 0.5, 0.0) + b = np.where(rng.random((8, 4, 16)) < dens_ab, rng.random((8, 4, 16)) + 0.5, 0.0) + s = np.where(rng.random((8, 16, 16)) < dens_s, 1.0, 0.0) + masks.clear() + masks.hints.clear() + assert_close(xp.multiply(make(s, 0.0), make(a, 0.0) @ make(b, 0.0)), s * (a @ b)) + assert masks.hints == [hint] + # a matmul whose batch axes broadcast (no einsum terms): the batches of the result + a = rng.random((1, 16, 4)) + 0.5 + b = rng.random((8, 4, 16)) + 0.5 + s = np.where(rng.random((8, 16, 16)) < 0.005, 1.0, 0.0) + masks.hints.clear() + assert_close(xp.multiply(make(s, 0.0), xp.asarray(a) @ xp.asarray(b)), s * (a @ b)) + assert masks.hints == [{"axb_method": "dot"}] + + +def test_product_times_itself_computes_once(masks): + # (review finding L6) P * P: no push-down, one computation + a = _rand((5, 5)) + P = make(a, 0.0) @ make(a, 0.0) + assert_close(P * P, (a @ a) * (a @ a)) + assert masks == [None] + # (without the result's shape, one batch is assumed) + P = make(a, 0.0) @ make(a, 0.0) + assert _lazy._dot_hint(P._pending, 1) == {"axb_method": "dot"} + + +THREE = [ + ("ij,jk,ik->", [(5, 6), (6, 7), (5, 7)]), # (triangle counting) + ("ij,jk,ik->i", [(5, 6), (6, 7), (5, 7)]), + ("ij,jk,ik->ki", [(5, 6), (6, 7), (5, 7)]), + ("ik,ij,jk->ik", [(5, 7), (5, 6), (6, 7)]), # (an SDDMM) + ("qn,qf,nf->qn", [(5, 7), (5, 3), (7, 3)]), + ("ij,j,i->", [(5, 6), (6,), (5,)]), + ("bij,bjk,bik->b", [(2, 3, 4), (2, 4, 5), (2, 3, 5)]), +] + + +@pytest.mark.parametrize(("subscripts", "shapes"), THREE, ids=[s for s, _ in THREE]) +def test_einsum_three_operands_masked(subscripts, shapes, masks): + # One operand's indices are what the other two leave: their product is computed at that + # operand's stored positions only + for dtype in [np.float64, np.int64, np.bool_, np.complex128]: + arrays = [_rand(s, 0, dtype, density=0.4) for s in shapes] + fill = np.zeros((), dtype)[()] + Xs = [make(a, fill, rng) for a in arrays] + masks.clear() + R = xp.einsum(subscripts, *Xs) + assert_close(R, np.einsum(subscripts, *arrays)) + if all(xp.prune(x).nnz for x in Xs): # (an empty operand: an empty pair, no product) + assert masks[0] is not None + # Semirings whose multiply distributes over the monoid + from .test_array_api_semiring import SEMIRINGS, semiring_einsum + + for name in ["min_plus", "max_plus", "min_max", "max_min"]: + reduce, mult, e = SEMIRINGS[name] + arrays = [_rand(s, e, density=0.4) for s in shapes] + masks.clear() + Xs = [make(a, e, rng) for a in arrays] + R = xp.einsum(subscripts, *Xs, semiring=name) + assert_close(R, semiring_einsum(subscripts, *arrays, reduce=reduce, mult=mult)) + assert R.fill_value == e or not R.ndim + if all(xp.prune(x).nnz for x in Xs): + assert masks[0] is not None + + +def _structural_three(a, b, s, mult=np.multiply): + """``ij,jk,ik->ik`` with a term only where all three operands are nonzero (DESIGN D6).""" + stored = (a[:, :, None] != 0) & (b[None] != 0) & (s[:, None, :] != 0) + with np.errstate(invalid="ignore", over="ignore"): + terms = mult(mult(a[:, :, None], b[None]), s[:, None, :]) + return np.where(stored, terms, 0).sum(axis=1) + + +def test_einsum_three_operands_nonfinite(masks, monkeypatch): + # An operand holding inf or NaN: explicit fills go first (D12a), and the terms are the + # full einsum's (D6): no NaN where the masking operand stores nothing. The values do not + # depend on which path ran. + from graphblas.array_api import _linalg + + a, b, s = _rand((4, 5)), _rand((5, 6)), _rand((4, 6)) + for bad in [np.inf, -np.inf, np.nan]: + for which in range(3): + arrays = [a.copy(), b.copy(), s.copy()] + arrays[which][0, 1] = bad + Xs = [make(x, 0.0, rng) for x in arrays] + masks.clear() + R = xp.einsum("ij,jk,ik->ik", *Xs) + # (only the third operand covers the output "ik": when it is the non-finite one, + # the full einsum runs) + assert any(m is not None for m in masks) == (which != 2) + if which != 2: + assert_close(R, _structural_three(*arrays)) + with monkeypatch.context() as m: + m.setattr(_linalg, "_masked_three", lambda *args: None) + assert_close(R, xp.einsum("ij,jk,ik->ik", *Xs).todense()) + # With plus_times, a non-finite operand is never the one the pair is computed under + # (inf * (x + y) differs from inf * x + inf * y), and complex goes whole + c = [x.astype(np.complex128) for x in (a, b, s)] + c[0][0, 1] = np.inf + masks.clear() + R = xp.einsum("ij,jk,ik->ik", *[make(x, 0j) for x in c]) + assert all(m is None for m in masks) + assert_close(R, _structural_three(*c)) + # A NaN under a float min or max monoid: the full einsum, which propagates it (D12b) + inf = np.inf + arrays = [np.where(x == 0, inf, x) for x in (a, b, s)] + arrays[0][0, 0] = np.nan + masks.clear() + R = xp.einsum("ij,jk,ik->ik", *[make(x, inf) for x in arrays], semiring="min_plus") + assert all(m is None for m in masks) + x, y, z = arrays[0][:, :, None], arrays[1][None], arrays[2][:, None, :] + stored = (x != inf) & (y != inf) & (z != inf) # (a NaN is stored: the term is NaN) + assert_close(R, np.min(np.where(stored, x + y + z, inf), axis=1)) + + +def test_einsum_three_operands_not_masked(masks): + a, b, s = _rand((4, 5)), _rand((5, 6)), _rand((4, 6)) + A, B, S = make(a, 0.0), make(b, 0.0), make(s, 0.0) + # No operand is what the other two leave, or no index is shared + for subscripts, arrays in [ + ("ij,jk,kl->il", [a, b, _rand((6, 3))]), + ("ij,kl,ik->", [a, _rand((3, 2)), _rand((4, 3))]), + ]: + masks.clear() + R = xp.einsum(subscripts, *[make(x, 0.0) for x in arrays]) + assert_close(R, np.einsum(subscripts, *arrays)) + assert all(m is None for m in masks) + # A semiring whose multiply does not distribute over its monoid: plus_min (unsigned) + p, q, r = (np.abs(x).astype(np.uint8) for x in (a, b, s)) + masks.clear() + R = xp.einsum("ij,jk,ik->ik", make(p, 0), make(q, 0), make(r, 0), semiring="plus_min") + assert all(m is None for m in masks) + assert_close(R, _structural_three(p, q, r, np.minimum).astype(np.uint8)) + # A 0-d operand (its value is the identity here): the full einsum + inf = np.inf + m = [np.where(x == 0, inf, x) for x in (a, a)] + masks.clear() + R = xp.einsum( + "ij,ij,->", make(m[0], inf), make(m[1], inf), xp.asarray(inf), semiring="min_plus" + ) + assert all(x is None for x in masks) + assert_close(R, np.array(inf)) # (every term holds the identity inf) + # Deferral off: the pair is computed whole, then intersected + _lazy.enabled = False + try: + assert_close(xp.einsum("ij,jk,ik->", A, B, S), np.einsum("ij,jk,ik->", a, b, s)) + finally: + _lazy.enabled = True + + +def test_compute(): + a, b = _rand((4, 5)), _rand((5, 3)) + A, B = make(a, 0.0), make(b, 0.0) + P = A @ B + assert pending(P) + assert xp.compute(P) is P + assert not pending(P) + assert_close(P, a @ b) + assert xp.compute(A) is A + with pytest.raises(TypeError, match="needs Arrays"): + xp.compute(a) diff --git a/graphblas/tests/test_array_api_linalg.py b/graphblas/tests/test_array_api_linalg.py new file mode 100644 index 000000000..fc2f96a9d --- /dev/null +++ b/graphblas/tests/test_array_api_linalg.py @@ -0,0 +1,306 @@ +"""graphblas.array_api linear algebra (main namespace and the linalg extension) against numpy.""" + +import math + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp + +from .test_array_api_elementwise import assert_array_equal, check_invariants, make, np_call + +rng = np.random.default_rng(21) + + +def _rand(shape, fill, dtype=np.float64, density=0.5): + if np.dtype(dtype).kind == "b": + a = rng.random(shape) < 0.5 + return np.where(rng.random(shape) < density, a, bool(fill)) + if np.dtype(dtype).kind in "iu": + a = rng.integers(0, 5, shape) + else: + a = np.round(rng.normal(size=shape), 2) + return np.where(rng.random(shape) < density, a, fill).astype(dtype) + + +def assert_close(x, expected): + expected = np.asarray(expected) + check_invariants(x) + assert x.dtype == expected.dtype, (x.dtype, expected.dtype) + assert x.shape == expected.shape + got = x.todense() + if got.dtype.kind in "fc": + tol = 1e-5 if got.dtype.itemsize * (2 if got.dtype.kind == "f" else 1) == 8 else 1e-10 + np.testing.assert_allclose(got, expected, rtol=tol, atol=tol) + else: + np.testing.assert_array_equal(got, expected) + + +MATMUL_SHAPES = [ + ((3, 4), (4, 5)), + ((4,), (4, 5)), + ((3, 4), (4,)), + ((4,), (4,)), + ((2, 3, 4), (4, 5)), + ((2, 1, 3, 4), (3, 4, 2)), + ((0, 4), (4, 2)), +] + + +@pytest.mark.parametrize("shapes", MATMUL_SHAPES, ids=str) +def test_matmul(shapes): + for dtype in [np.float64, np.int64, np.float32, np.complex128, np.uint8, np.bool_]: + for fa, fb in [(0, 0), (1, 0), (0, 2), (1, 3)]: + if np.dtype(dtype).kind == "b" and (fa or fb) and (fa, fb) != (1, 3): + continue + a, b = _rand(shapes[0], fa, dtype), _rand(shapes[1], fb, dtype) + x1, x2 = make(a, fa), make(b, fb) + expected = np.matmul(a, b) + assert_close(xp.matmul(x1, x2), expected) + assert_close(x1 @ x2, expected) + with pytest.raises(ValueError, match="align"): + xp.matmul(xp.zeros((2, 3)), xp.zeros((2, 3))) + with pytest.raises(ValueError, match="at least one"): + xp.matmul(xp.asarray(1.0), xp.zeros(2)) + with pytest.raises(ValueError, match="broadcast"): + xp.matmul(xp.zeros((2, 2, 3)), xp.zeros((3, 3, 2))) + + +def test_matmul_zero_fill_is_structural(): + # (DESIGN D6: as scipy.sparse; numpy would give NaN for 0 * inf) + a = make(np.array([[np.inf, 0.0], [0.0, 1.0]]), 0.0) + b = make(np.array([[0.0, 1.0], [1.0, 0.0]]), 0.0) + assert_close(a @ b, np.array([[0.0, np.inf], [1.0, 0.0]])) + # a fill that is not finite: numpy, densely + c = make(np.array([[np.inf, 1.0], [np.inf, np.inf]]), np.inf) + assert_array_equal(np_call(xp.matmul, c, b), np_call(np.matmul, c.todense(), b.todense())) + with gb.config.set(array_api_densify_limit=2), pytest.raises(xp.DensifyError): + c @ b + + +def test_matmul_inplace(): + a = make(np.eye(3), 0.0) + b = make(np.arange(9.0).reshape(3, 3), 0.0) + a @= b + assert_close(a, np.arange(9.0).reshape(3, 3)) + + +@pytest.mark.parametrize( + "axes", [2, 1, 0, ([1, 2], [1, 0]), ([2], [0]), ([0, 1, 2], [2, 1, 0]), ((), ())] +) +def test_tensordot(axes): + for dtype in [np.float64, np.int32, np.complex64]: + for fa, fb in [(0, 0), (2, 0), (0, 1), (1, 1)]: + a = _rand((3, 4, 3), fa, dtype) + b = _rand((3, 4, 3), fb, dtype) if axes != 1 else _rand((3, 2), fb, dtype) + if axes == 2: + b = _rand((4, 3, 2), fb, dtype) + if axes == 0: + b = _rand((2,), fb, dtype) + expected = np.tensordot(a, b, axes=axes) + out = xp.tensordot(make(a, fa), make(b, fb), axes=axes) + if np.dtype(dtype) == np.complex64: + np.testing.assert_allclose(out.todense(), expected, rtol=1e-5, atol=1e-5) + assert out.dtype == expected.dtype + else: + assert_close(out, expected) + with pytest.raises(ValueError, match="differ"): + xp.tensordot(xp.zeros((2, 3)), xp.zeros((2, 3)), axes=1) + + +@pytest.mark.parametrize("axis", [-1, -2]) +def test_vecdot(axis): + for dtype in [np.float64, np.int16, np.complex128, np.bool_]: + for fa, fb in [(0, 0), (1, 0), (0, 2), (1, 2)]: + if np.dtype(dtype).kind == "b" and (fa or fb) and (fa, fb) != (1, 2): + continue + for sa, sb in [((3, 4), (3, 4)), ((2, 3, 4), (3, 4)), ((3, 1), (1, 3, 1))]: + a, b = _rand(sa, fa, dtype), _rand(sb, fb, dtype) + if -axis > min(a.ndim, b.ndim) or a.shape[axis] != b.shape[axis]: + continue + out = xp.vecdot(make(a, fa), make(b, fb), axis=axis) + assert_close(out, np.vecdot(a, b, axis=axis)) + with pytest.raises(ValueError, match="axis"): + xp.vecdot(xp.zeros(3), xp.zeros(3), axis=-2) + a = np.array([1 + 2j, np.inf, 1j]) + b = np.array([2j, 1, 1 - 1j]) + assert_array_equal(np_call(xp.vecdot, make(a, np.inf), make(b, 0)), np_call(np.vecdot, a, b)) + + +def test_matrix_transpose(): + a = _rand((2, 3, 4), 0.0) + assert_array_equal(xp.matrix_transpose(make(a, 0.0)), np.swapaxes(a, -1, -2)) + assert_array_equal(xp.linalg.matrix_transpose(make(a, 1.0)), np.swapaxes(a, -1, -2)) + with pytest.raises(ValueError, match="2 dimensions"): + xp.matrix_transpose(xp.zeros(3)) + + +@pytest.mark.parametrize("offset", [-3, -1, 0, 1, 2, 5]) +def test_diagonal_trace(offset): + for fill in [0.0, 1.5]: + a = _rand((2, 3, 4), fill) + x = make(a, fill) + assert_array_equal( + xp.linalg.diagonal(x, offset=offset), np.linalg.diagonal(a, offset=offset) + ) + assert_close(xp.linalg.trace(x, offset=offset), np.linalg.trace(a, offset=offset)) + b = _rand((4, 4), fill, np.int8) + assert_close( + xp.linalg.trace(make(b, int(fill)), offset=offset, dtype=xp.int8), + np.linalg.trace(b, offset=offset, dtype=np.int8), + ) + + +def test_outer_cross(): + a, b = _rand((4,), 0.0), _rand((3,), 1.0) + assert_array_equal(xp.linalg.outer(make(a, 0.0), make(b, 1.0)), np.outer(a, b)) + u, v = _rand((2, 3), 0.0, density=0.7), _rand((3,), 2.0, density=0.7) + assert_array_equal(xp.linalg.cross(make(u, 0.0), make(v, 2.0)), np.cross(u, v)) + u, v = _rand((3, 2), 0.0, density=0.7), _rand((3, 2), 0.0, density=0.7) + assert_array_equal( + xp.linalg.cross(make(u, 0.0), make(v, 0.0), axis=-2), np.cross(u, v, axis=-2) + ) + with pytest.raises(ValueError, match="size 3"): + xp.linalg.cross(xp.zeros(2), xp.zeros(2)) + with pytest.raises(ValueError, match="1-d"): + xp.linalg.outer(xp.zeros((2, 2)), xp.zeros(2)) + + +@pytest.mark.parametrize("order", [2, 1, 0, math.inf, -math.inf, 3, -1, 0.5]) +def test_vector_norm(order): + for dtype, fill in [(np.float64, 0.0), (np.float64, -2.0), (np.complex128, 0), (np.int32, 0)]: + a = _rand((3, 4), fill, dtype) + x = make(a, fill) + for axis in [None, 0, -1, (0, 1)]: + for keepdims in [False, True]: + expected = np_call( + np.linalg.vector_norm, a, axis=axis, keepdims=keepdims, ord=order + ) + assert_close( + np_call(xp.linalg.vector_norm, x, axis=axis, keepdims=keepdims, ord=order), + expected, + ) + + +@pytest.mark.parametrize("order", ["fro", 1, -1, math.inf, -math.inf, 2, -2, "nuc"]) +def test_matrix_norm(order): + for fill in [0.0, 1.0]: + a = _rand((2, 3, 4), fill) + for keepdims in [False, True]: + expected = np.linalg.matrix_norm(a, ord=order, keepdims=keepdims) + assert_close( + xp.linalg.matrix_norm(make(a, fill), ord=order, keepdims=keepdims), expected + ) + with pytest.raises(ValueError, match="order"): + xp.linalg.matrix_norm(xp.zeros((2, 2)), ord=3) + + +@pytest.mark.parametrize("n", [0, 1, 2, 5, -1, -2]) +def test_matrix_power(n): + a = np.array([[2.0, 1.0, 0.0], [0.0, 1.0, 0.0], [1.0, 0.0, 3.0]]) + for fill in [0.0, 1.0]: + dense = np.where(a == 0, fill, a) + expected = np.linalg.matrix_power(dense, n) + assert_close(xp.linalg.matrix_power(make(dense, fill), n), expected) + b = np.stack([a, a.T]) + assert_close(xp.linalg.matrix_power(make(b, 0.0), 3), np.linalg.matrix_power(b, 3)) + with pytest.raises(ValueError, match="square"): + xp.linalg.matrix_power(xp.zeros((2, 3)), 2) + + +def test_dense_algebra(): + a = np.array([[4.0, 1.0, 0.0], [1.0, 3.0, 0.0], [0.0, 0.0, 2.0]]) + x = make(a, 0.0) + la = np.linalg + assert_close(xp.linalg.cholesky(x), la.cholesky(a)) + assert_close(xp.linalg.cholesky(x, upper=True), la.cholesky(a).T) + assert_close(xp.linalg.det(x), la.det(a)) + w, _ = xp.linalg.eigh(x) + assert_close(w, la.eigh(a).eigenvalues) + assert_close(xp.linalg.eigvalsh(x), la.eigvalsh(a)) + assert_close(xp.linalg.eigvals(x), la.eigvals(a)) + assert xp.linalg.eig(x).eigenvectors.shape == (3, 3) + assert_close(xp.linalg.inv(x), la.inv(a)) + assert_close(xp.linalg.pinv(x), la.pinv(a)) + assert_close(xp.linalg.matrix_rank(x), la.matrix_rank(a)) + q, r = xp.linalg.qr(x) + assert_close(q @ r, a) + sign, logdet = xp.linalg.slogdet(x) + assert_close(sign, la.slogdet(a).sign) + assert_close(logdet, la.slogdet(a).logabsdet) + b = np.array([1.0, 0.0, 2.0]) + assert_close(xp.linalg.solve(x, make(b, 0.0)), la.solve(a, b)) + B = np.stack([a, a]) + assert_close(xp.linalg.solve(make(B, 0.0), make(B, 0.0)), la.solve(B, B)) + s = xp.linalg.svd(x, full_matrices=False).S + assert_close(s, la.svd(a).S) + assert_close(xp.linalg.svdvals(x), la.svdvals(a)) + assert_close(xp.linalg.inv(make(np.eye(2, dtype=np.int64), 0)), np.eye(2)) + with gb.config.set(array_api_densify_limit=4), pytest.raises(xp.DensifyError, match="dense"): + xp.linalg.inv(x) + with pytest.raises(ValueError, match="2 dimensions"): + xp.linalg.det(xp.zeros(3)) + with pytest.raises(TypeError, match="Array"): + xp.linalg.det(a) + + +def test_nonzero_fills_follow_d6(): + # (review finding M3, DESIGN D12c) a nonzero fill of an operand that stores every + # position is taken as 0; any other nonzero fill goes dense (no rank-1 split, which turned + # inf into NaN and lost every digit when the fill dwarfs the values) + inf = make(np.array([np.inf]), 0.0) + five = xp.from_coo(([0],), [1.0], (1,), fill_value=5.0) + assert_close(xp.vecdot(inf, five), np.array(np.inf)) + m = make(np.array([[np.inf, 0.0]]), 0.0) + k = xp.from_coo(([0, 1], [0, 0]), [1.0, 5.0], (2, 1), fill_value=5.0) + assert_close(m @ k, [[np.inf]]) + big = xp.from_coo(([0, 1],), [1.0, 2.0], (2,), fill_value=1e20) + assert_close(xp.vecdot(make(np.array([1.0, -1.0]), 0.0), big), np.array(-1.0)) + big32 = xp.from_coo(([0, 1],), np.array([1.0, 2.0], np.float32), (2,), fill_value=1e8) + out = xp.vecdot(make(np.array([1.0, -1.0], np.float32), 0.0), big32) + assert_close(out, np.array(-1.0, np.float32)) + # dense under the limit on the result, not only on the operands (review finding L8) + with gb.config.set(array_api_densify_limit=1000): + c = make(np.full((100, 1), np.inf), np.inf) + d = make(np.full((1, 100), 2.0), np.inf) + with pytest.raises(xp.DensifyError): + c @ d + with pytest.raises(xp.DensifyError): + xp.tensordot(c, d, axes=0) + + +def test_tensordot_axes_out_of_bounds(): + # (review finding L2) as numpy: no wrapping modulo the number of axes + x = make(np.ones((2, 3)), 0.0) + with pytest.raises(IndexError, match="out of bounds"): + xp.tensordot(x, x.T, axes=([5], [0])) + with pytest.raises(IndexError, match="out of bounds"): + xp.tensordot(x, x.T, axes=([1], [-3])) + assert_close(xp.tensordot(x, x.T, axes=([-1], [0])), np.ones((2, 3)) @ np.ones((3, 2))) + + +def test_finite_fills_are_stored_in_full(): + # (review findings R1, R3) an operand with a finite nonzero fill that fits under the + # densify limit is stored in full, whatever the other operands store: matrices of ones + # by a large sparse matrix stay sparse, and explicit zeros in the other operand do not + # change the answer (D10) + rng = np.random.default_rng(20) + n = 400 + a = np.zeros((n, n)) + a[rng.integers(0, n, 600), rng.integers(0, n, 600)] = rng.random(600) + 0.5 + A = xp.from_coo(np.nonzero(a), a[np.nonzero(a)], (n, n)) + with gb.config.set(array_api_densify_limit=10 * n): + r1 = A @ xp.ones((n, 8)) + r2 = A @ xp.full((n, 8), 2.5) + assert_close(r1, a @ np.ones((n, 8))) + assert_close(r2, a @ np.full((n, 8), 2.5)) + b = np.ones((2, 1500)) + b[0, 0] = np.inf + B = xp.asarray(b, fill_value=1.0) + A1 = xp.from_coo(([0], [0]), np.array([1.0]), (3000, 2)) + r = np.arange(1, 2001) + A2 = xp.from_coo((np.r_[0, r], np.r_[0, np.zeros(2000, int)]), np.r_[1.0, np.zeros(2000)], + (3000, 2)) # fmt: skip + for x in [A1, A2, xp.prune(A2)]: + assert (x @ B).todense()[1, 0] == 0.0 # (D6: an unstored zero times inf is no term) diff --git a/graphblas/tests/test_array_api_manipulation.py b/graphblas/tests/test_array_api_manipulation.py new file mode 100644 index 000000000..61b3dc9d3 --- /dev/null +++ b/graphblas/tests/test_array_api_manipulation.py @@ -0,0 +1,243 @@ +"""graphblas.array_api manipulation functions against numpy.""" + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp + +from .test_array_api_elementwise import assert_array_equal, inputs, make, same_value + + +@pytest.mark.parametrize( + ("shapes", "axis"), + [ + ([(2, 3), (4, 3)], 0), + ([(2, 3), (2, 1), (2, 4)], 1), + ([(2, 3), (2, 3)], -1), + ([(2, 3, 2), (2, 1, 2)], 1), + ([(3,), (0,), (2,)], 0), + ([(2, 3), (3,)], None), + ], +) +def test_concat(shapes, axis): + for dtype in [np.dtype("float64"), np.dtype("int16"), np.dtype("bool")]: + parts = [list(inputs(dtype, shape, seed=i)) for i, shape in enumerate(shapes)] + for i in range(len(parts[0])): + dense = [p[i % len(p)][0] for p in parts] + arrays = [p[i % len(p)][1] for p in parts] + assert_array_equal(xp.concat(arrays, axis=axis), np.concatenate(dense, axis=axis)) + mixed = [make(np.array([1, 0, 2], np.int8), 0), make(np.array([1.5, 2.5]), 2.5)] + expected = np.concatenate([np.array([1, 0, 2], np.int8), np.array([1.5, 2.5])]) + assert_array_equal(xp.concat(mixed), expected) + + +def test_concat_keeps_the_most_common_fill(): + big = xp.full((10,), 7.0) + small = make(np.array([0.0, 1.0]), 0.0) + out = xp.concat([small, big]) + assert out.fill_value == 7.0 + assert out.nnz == 2 + assert_array_equal(out, np.concatenate([[0.0, 1.0], np.full(10, 7.0)])) + + +@pytest.mark.parametrize("axis", [0, 1, 2, -1, -3]) +def test_stack(axis): + for dtype in [np.dtype("float32"), np.dtype("uint8")]: + a = list(inputs(dtype, (2, 3), seed=1)) + b = list(inputs(dtype, (2, 3), seed=2)) + for (da, xa), (db, xb) in zip(a, b, strict=False): + assert_array_equal(xp.stack([xa, xb, xa], axis=axis), np.stack([da, db, da], axis=axis)) + assert_array_equal(xp.stack([xp.asarray(1), xp.asarray(2)]), np.array([1, 2])) + with pytest.raises(ValueError, match="one shape"): + xp.stack([xp.zeros(2), xp.zeros(3)]) + with pytest.raises(ValueError, match="at least one"): + xp.stack([]) + + +DTYPES3 = [np.dtype("float64"), np.dtype("int8"), np.dtype("bool")] + + +def _cases(shape, seed=0): + for dtype in DTYPES3: + yield from inputs(dtype, shape, seed=seed) + + +def test_broadcast(): + assert xp.broadcast_shapes((2, 1), (3,), ()) == (2, 3) + assert xp.broadcast_shapes() == () + with pytest.raises(ValueError, match="broadcast"): + xp.broadcast_shapes((2,), (3,)) + for dense, x in _cases((3, 1)): + for shape in [(3, 1), (3, 4), (2, 3, 4), (0, 3, 1)]: + assert_array_equal(xp.broadcast_to(x, shape), np.broadcast_to(dense, shape)) + with pytest.raises(ValueError, match="broadcast"): + xp.broadcast_to(x, (3,)) + y = make(np.array([1, 0, 2, 0], dtype=dense.dtype), np.zeros((), dense.dtype)[()]) + out = xp.broadcast_arrays(x, y, xp.asarray(1)) + expected = np.broadcast_arrays(dense, y.todense(), np.array(1)) + assert isinstance(out, tuple) + for o, e in zip(out, expected, strict=True): + assert_array_equal(o, e) + + +@pytest.mark.parametrize("axis", [0, 1, -1, 2]) +def test_unstack(axis): + for dense, x in _cases((2, 3, 4)): + out = xp.unstack(x, axis=axis) + assert isinstance(out, tuple) + expected = [np.take(dense, i, axis=axis) for i in range(dense.shape[axis])] + assert len(out) == len(expected) + for o, e in zip(out, expected, strict=True): + assert_array_equal(o, e) + with pytest.raises(ValueError, match="dimension"): + xp.unstack(xp.asarray(1.0)) + + +@pytest.mark.parametrize("axis", [None, 0, 1, -1, (0, 2), ()]) +def test_flip(axis): + for dense, x in _cases((2, 3, 4)): + assert_array_equal(xp.flip(x, axis=axis), np.flip(dense, axis=axis)) + assert_array_equal(xp.flip(xp.asarray(3)), np.array(3)) + + +@pytest.mark.parametrize( + ("source", "destination"), [(0, 2), (-1, 0), ((0, 1), (2, 0)), ((), ()), ((0, 1, 2), (2, 1, 0))] +) +def test_moveaxis(source, destination): + for dense, x in _cases((2, 3, 4)): + assert_array_equal( + xp.moveaxis(x, source, destination), np.moveaxis(dense, source, destination) + ) + with pytest.raises(ValueError, match="as many"): + xp.moveaxis(xp.zeros((2, 3)), (0, 1), 0) + + +@pytest.mark.parametrize("axis", [None, 0, 1, -1]) +def test_repeat(axis): + for dense, x in _cases((2, 3)): + n = dense.size if axis is None else dense.shape[axis] + for repeats in [0, 1, 3, np.arange(n) % 3, np.array([2])]: + if isinstance(repeats, np.ndarray): + r = xp.asarray(repeats) + expected = np.repeat(dense, np.broadcast_to(repeats, (n,)), axis=axis) + else: + r = repeats + expected = np.repeat(dense, repeats, axis=axis) + assert_array_equal(xp.repeat(x, r, axis=axis), expected) + with pytest.raises(ValueError, match="nonnegative"): + xp.repeat(xp.zeros(3), -1) + with pytest.raises(TypeError, match="integer"): + xp.repeat(xp.zeros(3), xp.asarray([1.0, 2.0, 3.0])) + + +@pytest.mark.parametrize( + ("shift", "axis"), + [(1, None), (-7, None), (2, 0), (-1, 1), ((1, 2), (0, 1)), (3, (0, 1)), (1, (1, 1))], +) +def test_roll(shift, axis): + for dense, x in _cases((3, 4)): + assert_array_equal(xp.roll(x, shift, axis=axis), np.roll(dense, shift, axis=axis)) + with pytest.raises(ValueError, match="one shift"): + xp.roll(xp.zeros((2, 2)), (1, 2, 3), axis=(0, 1)) + + +@pytest.mark.parametrize("reps", [(2,), (2, 3), (1, 1), (0, 2), (2, 1, 3), (3,)]) +def test_tile(reps): + for dense, x in _cases((2, 3)): + assert_array_equal(xp.tile(x, reps), np.tile(dense, reps)) + assert_array_equal(xp.tile(xp.asarray(2.0), (3,)), np.tile(np.array(2.0), 3)) + with pytest.raises(ValueError, match="nonnegative"): + xp.tile(xp.zeros(2), (-1,)) + + +def test_refill(): + rng = np.random.default_rng(12) + a = np.where(rng.random((6, 7)) < 0.3, rng.normal(size=(6, 7)), 0.0) + a[0, 0] = np.inf + x = make(a, 0.0, rng) + for fill in [np.inf, 1.5, 0.0, np.nan, -0.0]: + y = xp.refill(x, fill) + assert_array_equal(y, a) + assert same_value(y.fill_value, np.float64(fill)) + expected = sum(not same_value(v, y.fill_value) for v in a.reshape(-1)) + assert y.nnz == expected + # a dense distance matrix: inf becomes the fill, the zeros stay (the hint of the + # semiring error) + g = np.full((4, 4), np.inf) + np.fill_diagonal(g, 0.0) + g[0, 1] = 1.0 + G = xp.refill(xp.asarray(g), np.inf) + assert (G.nnz, G.fill_value) == (5, np.inf) + assert_array_equal(xp.matmul(G, G, semiring="min_plus"), np.min(g[:, :, None] + g[None], 1)) + # stores every position: relabeled, then pruned; other dtypes; a 0-d array keeps its value + full = xp.asarray(np.array([1, 2, 2, 3])) + out = xp.refill(full, 2) + assert_array_equal(out, np.array([1, 2, 2, 3])) + assert (out.nnz, out.fill_value) == (2, 2) + b = xp.refill(make(np.array([True, False, False]), False), True) + assert_array_equal(b, np.array([True, False, False])) + assert (b.nnz, b.fill_value) == (2, True) + assert_array_equal(xp.refill(xp.asarray(5.0), 1.0), np.array(5.0)) + # errors: not an Array, a value the dtype does not hold, past the densify limit + with pytest.raises(TypeError, match="needs an Array"): + xp.refill(a, 1.0) + with pytest.raises(ValueError, match="not a value of int64"): + xp.refill(full, np.inf) + with pytest.raises(ValueError, match="not a value of int64"): + xp.refill(full, 1.5) + with gb.config.set(array_api_densify_limit=10), pytest.raises(xp.DensifyError): + xp.refill(x, 1.0) + with pytest.raises(ValueError, match="xp.refill"): + xp.asarray(x, fill_value=1.0) + + +def test_prune(): + cases = [ + (np.array([3, 1, 3, 0]), 3), + (np.array([True, False, True]), True), + (np.array([2.5, 1.0, 2.5, np.nan, -0.0]), 2.5), + (np.array([1.0, np.nan, 2.0]), 0.0), + (np.array([0.0, -0.0, 1.0, np.nan, 0.0]), 0.0), + (np.array([-0.0, 0.0, 2.0]), -0.0), + (np.array([np.nan, 1.0, np.nan, 0.0]), np.nan), + (np.array([0j, complex(-0.0, 0.0), 1j, complex(0.0, -0.0)]), 0j), + (np.zeros(0), 0.0), + ] + for dense, fill in cases: + # every position stored, explicit fills included + idx = np.nonzero(np.ones(dense.shape, bool)) + x = xp.from_coo(idx, dense[idx], dense.shape, fill_value=fill, dtype=dense.dtype) + out = xp.prune(x) + assert_array_equal(out, dense) + assert out.fill_value == x.fill_value or np.isnan(fill) + expected = sum(not same_value(v, out.fill_value) for v in dense.reshape(-1)) + assert out.nnz == expected + assert x.nnz == dense.size + assert_array_equal(xp.prune(xp.asarray(5.0)), np.array(5.0)) + with pytest.raises(TypeError, match="Array"): + xp.prune(np.zeros(3)) + + +@pytest.mark.parametrize("axis", [0, 1, 2]) +def test_concat_any_layout(axis): + # One GxB concat of the stored Matrices, whatever layout each part is stored in (a part in + # another layout moves), parts of size 0 along the axis included + from graphblas import Tensor + + rng = np.random.default_rng(7) + shapes = [[4, 3, 5], [4, 3, 5], [4, 3, 5]] + shapes[1][axis] = 2 + shapes[2][axis] = 0 + dense = [np.where(rng.random(s) < 0.4, rng.normal(size=s), 0.0) for s in shapes] + layouts = [((0,), (1, 2)), ((0, 1), (2,)), ((2, 0), (1,)), ((1,), (2, 0))] + for la in layouts: + for lb in layouts: + parts = [] + for d, lay in zip(dense, [la, lb, la], strict=True): + idx = np.nonzero(d) + T = Tensor.from_coo(idx, d[idx], shape=d.shape).relayout(lay) + parts.append(xp.asarray(T)) + assert_array_equal(xp.concat(parts, axis=axis), np.concatenate(dense, axis=axis)) + empty = [xp.zeros((0, 3)), xp.zeros((0, 3))] + assert_array_equal(xp.concat(empty), np.zeros((0, 3))) diff --git a/graphblas/tests/test_array_api_searching.py b/graphblas/tests/test_array_api_searching.py new file mode 100644 index 000000000..5b01a3286 --- /dev/null +++ b/graphblas/tests/test_array_api_searching.py @@ -0,0 +1,236 @@ +"""graphblas.array_api searching functions against numpy.""" + +import numpy as np +import pytest + +import graphblas.array_api as xp + +from .test_array_api_elementwise import DTYPES, assert_array_equal, inputs, make, np_call + +WHERE_SHAPES = [ + ((3, 4), (3, 4), (3, 4)), + ((3, 4), (4,), (3, 1)), + ((3, 1), (1, 4), ()), + ((), (2, 3), (2, 3)), + ((2, 1, 4), (3, 1), (4,)), + ((), (), ()), +] + + +@pytest.mark.parametrize("shapes", WHERE_SHAPES, ids=str) +def test_where(shapes): + rng = np.random.default_rng(1) + for dt1, dt2 in [ + ("float64", "float64"), + ("int8", "uint8"), + ("float32", "int16"), + ("bool", "complex64"), + ]: + dt1, dt2 = np.dtype(dt1), np.dtype(dt2) + for (c, xc), (a, xa), (b, xb) in zip( + inputs(np.dtype(bool), shapes[0], seed=2), + inputs(dt1, shapes[1], seed=3), + inputs(dt2, shapes[2], seed=4), + strict=False, + ): + assert_array_equal(xp.where(xc, xa, xb), np.where(c, a, b)) + for s in [0, 1.5, True]: + assert_array_equal(xp.where(xc, xa, s), np.where(c, a, s)) + assert_array_equal(xp.where(xc, s, xb), np.where(c, s, b)) + # a condition that is not bool counts by truth value + a = np.where(rng.random((3, 4)) < 0.5, 0.0, rng.normal(size=(3, 4))) + x = make(a, 0.0) + assert_array_equal(xp.where(x, x, -1.0), np.where(a, a, -1.0)) + assert_array_equal(xp.where(x > 0, x, 0.0), np.where(a > 0, a, 0.0)) + with pytest.raises(ValueError, match="broadcast"): + xp.where(make(np.ones(3, bool), False), x, x) + + +def test_where_stores_the_condition_and_the_selected_operand(): + # Where the condition stores nothing, its fill selects one operand: the result is that + # operand with the condition's stored positions written over, and the other operand is + # read only there (no replication of a broadcast operand) + rng = np.random.default_rng(7) + n = 60 + c = rng.random((n, n)) < 0.05 + labels = np.arange(n) + big = np.iinfo(np.int64).max + G = make(c, False, rng) # (with explicit False entries) + W = xp.where(G, xp.arange(n), big) + assert_array_equal(W, np.where(c, labels, big)) + assert W.nnz == c.sum() # (G's explicit False entries select big, the fill: not stored) + assert W.fill_value == big + # where(x > t, x, 0): x's entries above t only, read there only + a = np.where(rng.random((n, n)) < 0.3, rng.normal(size=(n, n)), 0.0) + X = make(a, 0.0, rng) + R = xp.where(xp.greater(X, 0.5), X, 0.0) + assert_array_equal(R, np.where(a > 0.5, a, 0.0)) + assert R.nnz == (a > 0.5).sum() + R = xp.where(xp.less_equal(X, 0.5), 0.0, X) # (a condition with fill True) + assert_array_equal(R, np.where(a <= 0.5, 0.0, a)) + assert R.nnz == (a > 0.5).sum() + # other layouts of the condition, and a fill of True (the base is then x1) + for cond, cn in [(xp.matrix_transpose(G), c.T), (xp.logical_not(G), ~c)]: + a = rng.normal(size=(n, n)) + a[rng.random((n, n)) < 0.5] = 0 + a[0, 0] = np.nan + for x1, x2, e in [ + (make(a, 0.0), xp.arange(n) * 1.0, np.where(cn, a, labels * 1.0)), + (xp.arange(n) * 1.0, make(a, 0.0, rng), np.where(cn, labels * 1.0, a)), + (xp.arange(n)[:, None] * 1.0, -1.0, np.where(cn, labels[:, None] * 1.0, -1.0)), + ]: + assert_array_equal(xp.where(cond, x1, x2), e) + # A condition that is not bool selects by truth value (NaN is true), with a broadcast + # operand read at its entries, or over a base with entries (where(frontier, k, level)); + # a complex condition goes through numpy + f = np.where(c, np.where(rng.random((n, n)) < 0.3, np.nan, 2.5), 0.0) + f[0, 0] = -0.0 + level = make(a, 0.0) + for cond, fill in [(f, 0.0), (np.where(c, 0, 3).astype(np.int8), 3), (f + 1j * f, 0j)]: + for x1, x2 in [ + (xp.arange(n), big), + (big, xp.arange(n)[:, None]), + (labels[0], -1), + (7, level), + (xp.arange(n), level), + ]: + e = np.where(cond, np.asarray(x1), np.asarray(x2)) + assert_array_equal(xp.where(make(cond, fill), x1, x2), e) + # The selected operand broadcasts over stored entries, or the condition broadcasts: the + # union (the result really holds the operand's entries along the broadcast axis) + v = np.where(rng.random(n) < 0.5, rng.normal(size=n), 0.0) + assert_array_equal(xp.where(G, -1.0, make(v, 0.0)), np.where(c, -1.0, v)) + assert_array_equal(xp.where(make(c[0], False), make(a, 0.0), 2.0), np.where(c[0], a, 2.0)) + # An empty condition: the selected operand + E = xp.zeros((n, n), dtype=xp.bool) + assert_array_equal(xp.where(E, 1.0, make(a, 0.0)), a) + assert xp.where(E, 1.0, make(a, 0.0)).nnz == make(a, 0.0).nnz + + +def test_where_dtypes(): + for dt in DTYPES: + for dense, x in inputs(dt, (5,), seed=6): + c = dense.astype(bool) if dt.kind != "c" else dense.real.astype(bool) + assert_array_equal(xp.where(xp.asarray(c), x, x), np.where(c, dense, dense)) + assert_array_equal( + np_call(xp.where, xp.asarray(c), x, 2), np_call(np.where, c, dense, 2) + ) + + +SEARCH_DTYPES = [np.dtype(t) for t in ("bool", "int8", "uint16", "int64", "float32", "float64")] + + +@pytest.mark.parametrize("name", ["argmax", "argmin"]) +@pytest.mark.parametrize("keepdims", [False, True]) +def test_argmax_argmin(name, keepdims): + func = getattr(np, name) + for dtype in SEARCH_DTYPES: + for shape, axes in [((7,), [None, 0, -1]), ((3, 5), [None, 0, 1]), ((2, 3, 4), [0, 1, 2])]: + for dense, x in inputs(dtype, shape, seed=8): + for axis in axes: + expected = func(dense, axis=axis, keepdims=keepdims) + out = getattr(xp, name)(x, axis=axis, keepdims=keepdims) + assert_array_equal(out, np.asarray(expected, dtype=np.int64)) + with pytest.raises(ValueError, match="empty"): + getattr(xp, name)(xp.zeros((0, 3)), axis=0) + with pytest.raises(TypeError, match="real"): + getattr(xp, name)(xp.zeros(3, dtype=xp.complex64)) + assert_array_equal(getattr(xp, name)(xp.zeros((2, 0)), axis=0), np.zeros(0, np.int64)) + + +def test_argmax_ties_with_the_fill(): + # the fill ties with a stored value: the first of either position + for a, fill in [ + ([0.0, 3.0, 0.0, 3.0], 3.0), + ([3.0, 3.0, 1.0, 0.0], 3.0), + ([-0.0, 0.0, -1.0], 0.0), + ([np.nan, 1.0, 1.0, np.nan], np.nan), + ([1.0, 1.0, np.nan, 1.0], 1.0), + ]: + a = np.array(a) + for x in [make(a, fill), make(a, fill, np.random.default_rng(0), explicit=0.7)]: + assert int(xp.argmax(x)) == np.argmax(a) + assert int(xp.argmin(x)) == np.argmin(a) + + +def test_argmax_argmin_random(): + # numpy's rules (first occurrence, NaN the extreme, the fill at unstored positions) + # over fills, densities, shapes, axes and stored layouts + rng = np.random.default_rng(12) + pools = { + "float64": [0.0, -0.0, 1.0, -1.0, 2.5, np.nan, np.inf, -np.inf], + "float32": [0.0, -0.0, 1.0, -1.0, np.nan, np.inf], + "int8": [-3, -1, 0, 1, 2], + "uint16": [0, 1, 2, 3], + "bool": [False, True], + } + shapes = [(1,), (6,), (1, 1), (3, 4), (6, 1), (1, 7), (2, 3, 4), (3, 1, 2, 2)] + for _ in range(250): + dt = np.dtype(list(pools)[rng.integers(len(pools))]) + pool = np.array(pools[dt.name], dt) + shape = shapes[rng.integers(len(shapes))] + fill = pool[rng.integers(len(pool))] + values = rng.choice(pool[: rng.integers(1, len(pool) + 1)], shape) + dense = np.where(rng.random(shape) < rng.choice([0.0, 0.2, 0.6, 1.0]), values, fill) + x = make(dense, fill, rng, explicit=0.2) + if dense.ndim > 1 and rng.random() < 0.5: + # (a stored layout other than the default) + perm = tuple(int(a) for a in rng.permutation(dense.ndim)) + inverse = tuple(int(a) for a in np.argsort(perm)) + x = xp.permute_dims(xp.asarray(xp.permute_dims(x, perm), copy=True), inverse) + axis = int(rng.integers(-dense.ndim, dense.ndim)) if rng.random() < 0.8 else None + keepdims = rng.random() < 0.3 + for name in ["argmax", "argmin"]: + expected = getattr(np, name)(dense, axis=axis, keepdims=keepdims) + out = getattr(xp, name)(x, axis=axis, keepdims=keepdims) + assert_array_equal(out, np.asarray(expected, dtype=np.int64)) + + +@pytest.mark.parametrize("axis", [None, 0, 1, (0, 1), ()]) +def test_count_nonzero(axis): + for dtype in [*SEARCH_DTYPES, np.dtype("complex64")]: + for dense, x in inputs(dtype, (3, 4), seed=9): + for keepdims in [False, True]: + expected = np.count_nonzero(dense, axis=axis, keepdims=keepdims) + out = xp.count_nonzero(x, axis=axis, keepdims=keepdims) + assert_array_equal(out, np.asarray(expected, dtype=np.int64)) + + +def test_nonzero(): + for dtype in [*SEARCH_DTYPES, np.dtype("complex128")]: + for shape in [(7,), (3, 4), (2, 3, 2)]: + for dense, x in inputs(dtype, shape, seed=10): + out = xp.nonzero(x) + expected = np.nonzero(dense) + assert isinstance(out, tuple) + assert len(out) == len(expected) + for o, e in zip(out, expected, strict=True): + assert_array_equal(o, e.astype(np.int64)) + with pytest.raises(ValueError, match="dimension"): + xp.nonzero(xp.asarray(1)) + big = xp.zeros((10**6, 10**6)) + big[3, 4] = 1.0 + out = xp.nonzero(big) + assert [int(c[0]) for c in out] == [3, 4] + + +@pytest.mark.parametrize("side", ["left", "right"]) +def test_searchsorted(side): + rng = np.random.default_rng(11) + for fill in [0.0, 2.0, -np.inf, np.nan]: + for _ in range(20): + a = np.sort(np.where(rng.random(9) < 0.5, fill, np.round(rng.normal(size=9) * 2))) + x1 = make(a, fill) + v = np.array([-10.0, -1.0, 0.0, fill, 1.0, 2.0, 2.5, np.inf, np.nan]) + for x2, vv in [(make(v, 0.0), v), (make(v, fill), v), (xp.asarray(1.0), np.array(1.0))]: + out = xp.searchsorted(x1, x2, side=side) + assert_array_equal(out, np.searchsorted(a, vv, side=side).astype(np.int64)) + assert int(xp.searchsorted(x1, 2.0, side=side)) == np.searchsorted(a, 2.0, side=side) + a = np.array([3, 0, 1, 0, 5]) + sorter = np.argsort(a, kind="stable") + out = xp.searchsorted(make(a, 0), xp.asarray([0, 4]), side=side, sorter=xp.asarray(sorter)) + assert_array_equal(out, np.searchsorted(a, [0, 4], side=side, sorter=sorter).astype(np.int64)) + with pytest.raises(ValueError, match="1-d"): + xp.searchsorted(xp.zeros((2, 2)), 1.0) + with pytest.raises(ValueError, match="side"): + xp.searchsorted(xp.zeros(2), 1.0, side="middle") diff --git a/graphblas/tests/test_array_api_semiring.py b/graphblas/tests/test_array_api_semiring.py new file mode 100644 index 000000000..42dbe7976 --- /dev/null +++ b/graphblas/tests/test_array_api_semiring.py @@ -0,0 +1,451 @@ +"""graphblas.array_api: the semiring= extension of matmul, tensordot and vecdot, and einsum.""" + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp + +from .test_array_api_elementwise import check_invariants, make + +rng = np.random.default_rng(58) + + +def _rand(shape, fill, dtype=np.float64, density=0.5): + dtype = np.dtype(dtype) + if dtype.kind == "b": + a = rng.random(shape) < 0.5 + elif dtype.kind in "iu": + a = rng.integers(0, 9, shape) + else: + a = np.round(rng.normal(size=shape) * 4, 1) + return np.where(rng.random(shape) < density, a, fill).astype(dtype) + + +def assert_close(x, expected, fill=None): + expected = np.asarray(expected) + check_invariants(x) + assert x.dtype == expected.dtype, (x.dtype, expected.dtype) + assert x.shape == expected.shape, (x.shape, expected.shape) + if fill is not None and x.ndim: + assert x.fill_value == fill or (np.isnan(fill) and np.isnan(x.fill_value)) + got = x.todense() + if got.dtype.kind in "fc": + np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-6) + else: + np.testing.assert_array_equal(got, expected) + + +# (name, numpy add-monoid reduction, numpy multiply, identity) for the oracles +SEMIRINGS = { + "min_plus": (np.min, np.add, np.inf), + "max_plus": (np.max, np.add, -np.inf), + "min_max": (np.min, np.maximum, np.inf), + "max_min": (np.max, np.minimum, -np.inf), +} + + +def semiring_einsum(subscripts, *arrays, reduce, mult): + """The numpy oracle: each term's operands combined by ``mult``, then ``reduce`` over the + summed indices (every index spelled out, so each operand broadcasts over all of them). + """ + inputs, output = subscripts.replace(" ", "").split("->") + terms = inputs.split(",") + letters = sorted(set("".join(terms))) + total = None + for term, a in zip(terms, arrays, strict=True): + a = np.asarray(a) + perm = [term.index(c) for c in letters if c in term] + a = np.transpose(a, perm) + shape = [a.shape[sorted(term, key=letters.index).index(c)] if c in term else 1 + for c in letters] # fmt: skip + a = a.reshape(shape) + total = a if total is None else mult(total, a) + with np.errstate(invalid="ignore"): + summed = tuple(i for i, c in enumerate(letters) if c not in output) + out = reduce(total, axis=summed) if summed else total + kept = [c for c in letters if c in output] + return np.transpose(out, [kept.index(c) for c in output]) + + +@pytest.mark.parametrize("name", list(SEMIRINGS)) +@pytest.mark.parametrize( + "shapes", + [((3, 4), (4, 5)), ((4,), (4, 5)), ((3, 4), (4,)), ((4,), (4,)), ((2, 1, 3, 4), (3, 4, 2))], + ids=str, +) +def test_matmul_semirings(name, shapes): + reduce, mult, e = SEMIRINGS[name] + for dtype in [np.float64, np.float32]: + a, b = _rand(shapes[0], e, dtype), _rand(shapes[1], e, dtype) + x1, x2 = make(a, e, rng), make(b, e) + a2 = a[None] if a.ndim == 1 else a + b2 = b[:, None] if b.ndim == 1 else b + expected = reduce(mult(a2[..., :, :, None], b2[..., None, :, :]), axis=-2) + if a.ndim == 1: + expected = expected[..., 0, :] + if b.ndim == 1: + expected = expected[..., 0] + expected = np.asarray(expected, dtype) + for sr in [name, gb.semiring.from_string(name)]: + assert_close(xp.matmul(x1, x2, semiring=sr), expected, e) + assert_close(xp.linalg.matmul(x1, x2, semiring=sr), expected, e) + # Stored values equal to the fill do not change the answer + assert_close(xp.matmul(make(a, e, rng), make(b, e, rng), semiring=name), expected, e) + + +def test_matmul_semiring_dtypes(): + a, b = _rand((3, 4), np.inf, np.float32), _rand((4, 2), np.inf) + expected = np.min(a[:, :, None] + b[None], axis=1) + R = xp.matmul(make(a, np.inf), make(b, np.inf), semiring="min_plus") + assert_close(R, expected.astype(np.float64), np.inf) + # A typed semiring keeps its type + R = xp.matmul(make(a, np.inf), make(b, np.inf), semiring=gb.semiring.min_plus["FP32"]) + assert_close(R, expected.astype(np.float32), np.inf) + # Bools: lor_land (the default), and plus_times is the same thing + p, q = _rand((3, 4), False, bool), _rand((4, 5), False, bool) + for sr in [None, "lor_land", "plus_times", gb.semiring.lor_land]: + assert_close(xp.matmul(make(p, False), make(q, False), semiring=sr), p @ q, False) + # land_lor's identity is True + expected = np.all(p[:, :, None] | q[None], axis=1) + R = xp.matmul(make(p, True), make(q, True), semiring="land_lor") + assert_close(R, expected, True) + # Numbers through a bool semiring: nonzero values are True + f, g = _rand((3, 4), 0.0), _rand((4, 2), 0.0) + R = xp.matmul(make(f, 0.0), make(g, 0.0), semiring="lor_land") + assert_close(R, (f != 0).astype(int) @ (g != 0).astype(int) > 0, False) + # Unsigned max_times and plus_min: 0 annihilates times and min there + u, v = _rand((3, 4), 0, np.uint8), _rand((4, 5), 0, np.uint8) + expected = np.max(u[:, :, None] * v[None], axis=1) + assert_close(xp.matmul(make(u, 0), make(v, 0), semiring="max_times"), expected, 0) + expected = np.sum(np.minimum(u[:, :, None], v[None]), axis=1, dtype=np.uint8) + assert_close(xp.matmul(make(u, 0), make(v, 0), semiring="plus_min"), expected, 0) + # An explicit plus_times is the default, nonzero fills included + for fa, fb in [(0, 0), (1, 0), (2, 3)]: + a, b = _rand((3, 4), fa), _rand((4, 2), fb) + assert_close(xp.matmul(make(a, fa), make(b, fb), semiring="plus_times"), a @ b) + # Integer plus_times (any fill 0) + i, j = _rand((3, 4), 0, np.int16), _rand((4, 2), 0, np.int16) + assert_close(xp.matmul(make(i, 0), make(j, 0), semiring="plus_times"), i @ j, 0) + + +def test_contraction_skips_structurally(): + # As 0 * inf in plus_times (D6): an infinity meeting an unstored inf is no NaN + a = make(np.array([[-np.inf, np.inf]]), np.inf) + b = make(np.array([[np.inf], [-np.inf]]), np.inf) + assert a.nnz == b.nnz == 1 + assert_close(xp.matmul(a, b, semiring="min_plus"), [[np.inf]], np.inf) + # numpy: inf + -inf is NaN + with np.errstate(invalid="ignore"): + assert np.isnan(np.min(a.todense()[:, :, None] + b.todense()[None], axis=1)).all() + + +@pytest.mark.parametrize( + ("xshape", "yshape", "axes"), + [ + ((3, 4, 5), (5, 6), 1), + ((3, 4, 5), (4, 5, 2), 2), + ((3, 4, 5), (5, 6, 3), ([0, 2], [2, 0])), + ((3, 4), (2, 5), 0), + ((4,), (4,), 1), + ], + ids=str, +) +def test_tensordot_semirings(xshape, yshape, axes): + for name, (reduce, mult, e) in SEMIRINGS.items(): + a, b = _rand(xshape, e), _rand(yshape, e) + if isinstance(axes, int): + a_axes, b_axes = list(range(len(xshape) - axes, len(xshape))), list(range(axes)) + else: + a_axes, b_axes = axes + la = "abcde"[: len(xshape)] + lb = list("vwxyz"[: len(yshape)]) + for i, j in zip(a_axes, b_axes, strict=True): + lb[j] = la[i] + lb = "".join(lb) + out = "".join(c for c in la if c not in lb) + "".join(c for c in lb if c not in la) + subs = f"{la},{lb}->{out}" + expected = semiring_einsum(subs, a, b, reduce=reduce, mult=mult) + R = xp.tensordot(make(a, e), make(b, e, rng), axes=axes, semiring=name) + assert_close(R, expected, e) + assert_close(xp.linalg.tensordot(make(a, e), make(b, e), axes=axes, semiring=name), + expected, e) # fmt: skip + + +@pytest.mark.parametrize( + ("xshape", "yshape", "axis"), + [((3, 4), (3, 4), -1), ((3, 4), (4,), -1), ((2, 3, 4), (3, 4), -2), ((3, 1, 4), (5, 4), -1)], + ids=str, +) +def test_vecdot_semirings(xshape, yshape, axis): + for name, (reduce, mult, e) in SEMIRINGS.items(): + a, b = _rand(xshape, e), _rand(yshape, e) + a_b, b_b = np.broadcast_arrays(a, b) + expected = reduce(mult(a_b, b_b), axis=axis) + for func in [xp.vecdot, xp.linalg.vecdot]: + R = func(make(a, e, rng), make(b, e), axis=axis, semiring=name) + assert_close(R, expected, e) + # lor_land of bools is the default's any(x1 & x2) + p, q = _rand((3, 4), False, bool), _rand((3, 4), False, bool) + assert_close(xp.vecdot(make(p, False), make(q, False), semiring="lor_land"), + np.any(p & q, axis=-1), False) # fmt: skip + + +EINSUMS = [ + ("ij,jk->ik", [(3, 4), (4, 5)]), + ("ik,jk->ij", [(3, 4), (5, 4)]), + ("ij,jk,kl->il", [(3, 4), (4, 5), (5, 2)]), + ("ijk,jr,kr->ir", [(3, 4, 5), (4, 2), (5, 2)]), + ("ij,ij->i", [(3, 4), (3, 4)]), + ("i,j->ij", [(3,), (4,)]), + ("bij,bjk->bik", [(2, 3, 4), (2, 4, 5)]), + ("ij->ji", [(3, 4)]), + ("ij->i", [(3, 4)]), + ("ijk->", [(2, 3, 4)]), + ("ij,jk", [(3, 4), (4, 5)]), + ("ij,j->i", [(3, 4), (4,)]), +] + + +@pytest.mark.parametrize(("subscripts", "shapes"), EINSUMS, ids=[s for s, _ in EINSUMS]) +def test_einsum(subscripts, shapes): + # The default: numpy's einsum, zero fills structural, other fills densely through numpy + for dtype in [np.float64, np.int64, np.bool_, np.complex128]: + for fills in [(0,) * len(shapes), (1,) + (0,) * (len(shapes) - 1)]: + if np.dtype(dtype).kind == "b" and any(fills): + fills = (True,) * len(shapes) + arrays = [_rand(s, f, dtype) for s, f in zip(shapes, fills, strict=True)] + expected = np.einsum(subscripts, *arrays) + R = xp.einsum(subscripts, *[make(a, f) for a, f in zip(arrays, fills, strict=True)]) + assert_close(R, expected) + # Semirings + full = subscripts if "->" in subscripts else subscripts + "->ik" + for name, (reduce, mult, e) in SEMIRINGS.items(): + arrays = [_rand(s, e) for s in shapes] + expected = semiring_einsum(full, *arrays, reduce=reduce, mult=mult) + R = xp.einsum(subscripts, *[make(a, e, rng) for a in arrays], semiring=name) + assert_close(R, np.asarray(expected), e) + + +def test_einsum_operands_and_errors(): + a = _rand((3, 4), 0.0) + x = make(a, 0.0) + # numpy operands, empty operands and sizes + assert_close(xp.einsum("ij,kj->ik", a, x), a @ a.T) + assert_close(xp.einsum("ij,jk->ik", x, xp.zeros((4, 2))), np.zeros((3, 2))) + assert_close(xp.einsum("ij,jk->ik", xp.zeros((3, 0)), xp.zeros((0, 2))), np.zeros((3, 2))) + R = xp.einsum("ij,jk->ik", xp.full((3, 0), np.inf), xp.full((0, 2), np.inf), + semiring="min_plus") # fmt: skip + assert_close(R, np.full((3, 2), np.inf), np.inf) + # A 0-d operand + assert_close(xp.einsum(",ij->ij", xp.asarray(2.0), x), 2 * a) + assert_close(xp.einsum(",ij->ij", xp.asarray(0.0), x), 0 * a) + with pytest.raises(TypeError, match="subscripts"): + xp.einsum(x, x) + with pytest.raises(ValueError, match="at least one"): + xp.einsum("ij") + with pytest.raises(TypeError, match="Arrays"): + xp.einsum("ij->ij", [[1.0]]) + with pytest.raises(NotImplementedError, match="not supported"): + xp.einsum("...ij->...ji", x) + with pytest.raises(NotImplementedError, match="diagonal"): + xp.einsum("ii->i", make(np.eye(3), 0.0)) + with pytest.raises(ValueError, match="fill is inf"): + xp.einsum("ij,jk->ik", x, x.T, semiring="min_plus") + with pytest.raises(ValueError, match="operand 1"): + xp.einsum("ij,jk->ik", xp.full((3, 4), np.inf), x.T, semiring="min_plus") + # One operand only sums: the identity is all its fill needs + b = _rand((3, 4), -np.inf) + R = xp.einsum("ij->i", make(b, -np.inf), semiring="max_times") + assert_close(R, np.max(b, axis=1), -np.inf) + with gb.config.set(array_api_densify_limit=5), pytest.raises(xp.DensifyError): + xp.einsum("ij,jk->ik", make(a, 1.0), x.T) + + +def test_semiring_errors(): + a = make(_rand((3, 4), np.inf), np.inf) + b = make(_rand((4, 2), np.inf), np.inf) + z = make(_rand((4, 2), 0.0), 0.0) + # Fills that are not the identity + with pytest.raises(ValueError, match="x2 has the fill 0.0"): + xp.matmul(a, z, semiring="min_plus") + with pytest.raises(ValueError, match="x1 has the fill inf"): + xp.tensordot(a, z, axes=1, semiring="max_plus") + with pytest.raises(ValueError, match="fill is -inf"): + xp.vecdot(a, a, semiring="max_plus") + with pytest.raises(ValueError, match="fill is inf"): + xp.matmul(make(_rand((3, 4), np.nan), np.nan), b, semiring="min_plus") + with pytest.raises(ValueError, match="fill_value=xp.inf"): + xp.matmul(z.T, z, semiring="min_plus") + with pytest.raises(ValueError, match="fill_value=-xp.inf"): + xp.matmul(z.T, z, semiring="max_plus") + t = xp.zeros((4, 2), dtype=xp.bool) + with pytest.raises(ValueError, match="fill_value=True"): + xp.matmul(t.T, t, semiring="land_lor") + # Semirings whose identity does not annihilate the multiply: no fill works + for sr, dtype in [ + ("max_times", np.float64), + ("plus_pair", np.float64), + ("any_pair", np.float64), + ("min_first", np.float64), + ("plus_first", np.float64), + ("min_plus", np.int64), + ("max_plus", np.uint8), + ("min_times", np.uint16), + ("lor_eq", np.float64), + ]: + e = np.asarray(0).astype(dtype) + x = make(_rand((3, 4), e, dtype), e) + y = make(_rand((4, 2), e, dtype), e) + with pytest.raises(ValueError, match="cannot skip"): + xp.matmul(x, y, semiring=sr) + with pytest.raises(ValueError, match="cannot skip"): + xp.einsum("ij,jk->ik", x, y, semiring=sr) + with pytest.raises(ValueError, match="cannot skip"): + xp.matmul(a, b, semiring=gb.semiring.ss.min_firsti) + with pytest.raises(ValueError, match="unknown semiring"): + xp.matmul(a, b, semiring="min.nope") + with pytest.raises(TypeError, match="Semiring or its name"): + xp.matmul(a, b, semiring=gb.binary.plus) + with pytest.raises(TypeError, match="does not work with complex128"): + xp.matmul(xp.astype(a, xp.complex128), b, semiring="min_plus") + with pytest.raises(TypeError, match="does not work with bool"): + xp.vecdot(xp.zeros(3, dtype=xp.bool), xp.zeros(3, dtype=xp.bool), semiring="min_plus") + + +def test_probes(): + from graphblas.array_api._linalg import _probes + + for dt in ["bool", "int8", "uint64", "float32", "complex64"]: + p = _probes(np.dtype(dt)) + assert p.dtype == np.dtype(dt) + if p.dtype.kind in "fc": + assert np.isfinite(p).all() + assert {-128, 127, 0} <= set(_probes(np.dtype("int8")).tolist()) + + +def test_explicit_fills_do_not_change_values(): + # (review finding M2, DESIGN D12a) a stored entry equal to the fill meets an inf as an + # unstored position does: no term + a = xp.from_coo(([0], [0]), [0.0], (1, 1)) + b = make(np.array([[np.inf]]), 0.0) + assert_close(a @ b, [[0.0]]) + assert_close(xp.prune(a) @ b, [[0.0]]) + a = xp.from_coo(([0], [0]), [np.inf], (1, 1), fill_value=np.inf) + b = make(np.array([[-np.inf]]), np.inf) + for x in [a, xp.prune(a)]: + assert_close(xp.matmul(x, b, semiring="min_plus"), [[np.inf]], np.inf) + # with no inf or NaN anywhere, explicit fills stay (nothing to prune) + a = xp.from_coo(([0, 0], [0, 1]), [0.0, 2.0], (1, 2)) + b = make(np.array([[1.0], [3.0]]), 0.0) + assert_close(a @ b, [[6.0]]) + + +def test_nan_propagates_through_min_and_max(): + # (review finding M1, DESIGN D12b) a stored NaN term makes the min or max NaN, as numpy's + # do, although SuiteSparse's fmin and fmax skip it; positions it does not reach keep theirs + inf = np.inf + a = np.array([[np.nan, 5.0], [inf, 2.0]]) + b = np.array([[1.0, inf], [3.0, 4.0]]) + A, B = make(a, inf), make(b, inf) + for name, (reduce, mult, e) in [("min_plus", SEMIRINGS["min_plus"]), + ("min_max", SEMIRINGS["min_max"])]: # fmt: skip + with np.errstate(invalid="ignore"): + expected = reduce(mult(a[:, :, None], b[None]), axis=1) + # (D6: a NaN meeting an unstored inf is no term, so [0, 1] keeps its value) + expected[0, 1] = reduce(mult(a[0, 1], b[1, 1])) + assert_close(xp.matmul(A, B, semiring=name), expected, e) + assert_close(xp.tensordot(A, B, axes=1, semiring=name), expected, e) + assert_close(xp.einsum("ij,jk->ik", A, B, semiring=name), expected, e) + m = make(np.array([[np.nan, 5.0]]), -inf) + n = make(np.array([[1.0], [3.0]]), -inf) + assert_close(xp.matmul(m, n, semiring="max_plus"), [[np.nan]], -inf) + assert_close(xp.vecdot(m, xp.matrix_transpose(n), semiring="max_plus"), [np.nan], -inf) + x = make(np.array([2.0, np.nan, 1.0]), inf) + assert_close(xp.einsum("b->", x, semiring="min_max"), np.array(np.nan)) + assert_close(xp.einsum("b->", make(np.array([2.0, 1.0]), inf), semiring="min_max"), + np.array(1.0)) # fmt: skip + + +def test_fully_stored_operands_take_the_identity(): + # An operand that stores every position never shows its fill: a semiring accepts it + # whatever its fill, when its dtype holds the identity + a = np.array([[1.0, 2.0], [np.inf, 3.0]]) + b = np.array([[0.5, 1.0], [2.0, 4.0]]) + A, B = make(a, 0.0), make(b, 0.0) + assert A.nnz == A.size + expected = np.min(a[:, :, None] + b[None], axis=1) + for R in [ + xp.matmul(A, B, semiring="min_plus"), + xp.tensordot(A, B, axes=1, semiring="min_plus"), + xp.einsum("ij,jk->ik", A, B, semiring="min_plus"), + xp.vecdot(A[:, :, None], B[None], axis=1, semiring="min_plus"), + ]: + assert_close(R, expected, np.inf) + # (the inf entry is an explicit fill once relabeled; a NaN elsewhere still propagates) + a2 = a.copy() + a2[0, 0] = np.nan + R = xp.matmul(make(a2, 0.0), B, semiring="min_plus") + assert np.isnan(R.todense()[0]).all() + # An integer operand cannot hold inf: still refused + i = make(np.array([[1, 2], [3, 4]]), 0) + with pytest.raises(ValueError, match="needs operands whose fill"): + xp.matmul(i, make(b, np.inf), semiring="min_plus") + # Partly stored: refused as before + with pytest.raises(ValueError, match="needs operands whose fill"): + xp.matmul(make(np.array([[0.0, 1.0]]), 0.0), B, semiring="min_plus") + + +def test_soundness_is_cached_per_semiring_object(): + # (review finding L1) two semirings of one name: each is checked + from graphblas.core.operator import Semiring + + s1 = Semiring.register_anonymous(gb.monoid.min, gb.binary.plus, name="same_name") + s2 = Semiring.register_anonymous(gb.monoid.min, gb.binary.minus, name="same_name") + a = make(np.array([[1.0, 2.0]]), np.inf) + b = make(np.array([[1.0], [3.0]]), np.inf) + assert_close(xp.matmul(a, b, semiring=s1), [[2.0]], np.inf) + with pytest.raises(ValueError, match="cannot skip"): + xp.matmul(a, b, semiring=s2) + + +def test_vecdot_plus_times_conjugates(): + # (review finding L4) the default semiring, given or not, conjugates x1 + x = make(np.array([1j, 2.0]), 0j) + y = make(np.array([1j, 1.0]), 0j) + for semiring in [None, "plus_times"]: + assert_close(xp.vecdot(x, y, semiring=semiring), np.array(3 + 0j)) + + +def test_einsum_with_scalars_and_ones_stays_sparse(): + # (review finding L5) a 0-d operand is a scalar factor, and a small operand with a finite + # nonzero fill (a vector of ones: row sums) is stored in full, not densified with the rest + rng = np.random.default_rng(18) + n = 400 + a = np.zeros((n, n)) + a[rng.integers(0, n, 600), rng.integers(0, n, 600)] = rng.random(600) + 0.5 + A = xp.from_coo(np.nonzero(a), a[np.nonzero(a)], (n, n)) + two, ones = xp.asarray(2.0), xp.ones(n) + with gb.config.set(array_api_densify_limit=10 * n): + outs = [ + xp.einsum("ij,j->i", A, ones), + A @ ones, + xp.tensordot(A, ones, axes=1), + xp.einsum("ij,->ij", A, two), + xp.einsum(",ij,->ji", two, A, two), + xp.einsum("ij,", A, two), + ] + for out, expected in zip( + outs, [a.sum(1), a.sum(1), a.sum(1), 2 * a, 4 * a.T, 2 * a], strict=True + ): + assert_close(out, expected) + assert_close(xp.einsum(",->", two, xp.asarray(3.0)), np.array(6.0)) + # a non-finite fill that shows goes dense, as numpy computes it (inf * 0 is NaN there) + c = np.array([[1.0, np.inf], [np.inf, 2.0]]) + d = np.array([[0.0, 1.0], [1.0, 0.0]]) + with np.errstate(invalid="ignore"): + assert_close(xp.einsum("ij,->", make(d, 0.0), xp.asarray(np.inf)), + np.einsum("ij,->", d, np.inf)) # fmt: skip + with np.errstate(invalid="ignore"): + expected = np.einsum("ij,jk->ik", c, d) + assert_close(xp.einsum("ij,jk->ik", make(c, np.inf), make(d, 0.0)), expected) diff --git a/graphblas/tests/test_array_api_set.py b/graphblas/tests/test_array_api_set.py new file mode 100644 index 000000000..6bcd36084 --- /dev/null +++ b/graphblas/tests/test_array_api_set.py @@ -0,0 +1,86 @@ +"""graphblas.array_api set functions against numpy.""" + +from types import SimpleNamespace + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp + +from .test_array_api_elementwise import assert_array_equal, check_invariants, inputs, make + +DTYPES = [np.dtype(t) for t in ("bool", "int8", "uint32", "int64", "float32", "float64")] + + +def _check_values(out, expected): + """Unique values: as numpy's, except that which of 0.0 and -0.0 stands for both is free.""" + check_invariants(out) + got = out.todense() + assert got.dtype == expected.dtype + np.testing.assert_array_equal(got, expected) # (0.0 == -0.0, NaN matches NaN) + + +@pytest.mark.parametrize("shape", [(9,), (3, 4), (2, 3, 2), ()]) +def test_unique(shape): + for dtype in DTYPES: + for dense, x in inputs(dtype, shape, seed=13): + # (numpy's unique_* need not sort; np.unique does, as graphblas.array_api does) + values, indices, inverse, counts = np.unique( + dense, return_index=True, return_inverse=True, return_counts=True, equal_nan=False + ) + e = SimpleNamespace( + values=values, + indices=indices, + inverse_indices=inverse.reshape(dense.shape), + counts=counts, + ) + r = xp.unique_all(x) + _check_values(r.values, e.values) + assert_array_equal(r.counts, e.counts.astype(np.int64)) + # the first index of each value (of either zero, for 0.0 and -0.0) + assert_array_equal(r.indices, e.indices.astype(np.int64)) + assert_array_equal(r.inverse_indices, e.inverse_indices.astype(np.int64)) + _check_values(xp.unique_values(x), e.values) + c = xp.unique_counts(x) + _check_values(c.values, e.values) + assert_array_equal(c.counts, e.counts.astype(np.int64)) + i = xp.unique_inverse(x) + _check_values(i.values, e.values) + assert_array_equal(i.inverse_indices, e.inverse_indices.astype(np.int64)) + + +def test_unique_is_sparse(): + x = xp.zeros((10**6, 10**6)) + x[3, [1, 5]] = np.array([2.0, -1.0]) + r = xp.unique_all(x) + assert_array_equal(r.values, np.array([-1.0, 0.0, 2.0])) + assert_array_equal(r.indices, np.array([3 * 10**6 + 5, 0, 3 * 10**6 + 1])) + assert_array_equal(r.counts, np.array([1, 10**12 - 2, 1])) + assert r.inverse_indices.nnz == 2 + assert r.inverse_indices.fill_value == 1 + + +def test_unique_nan_fill(): + a = np.array([np.nan, 1.0, np.nan, 2.0]) + x = make(a, np.nan) + _check_values(xp.unique_values(x), np.unique(a, equal_nan=False)) + assert_array_equal(xp.unique_counts(x).counts, np.ones(4, np.int64)) + with gb.config.set(array_api_densify_limit=2), pytest.raises(xp.DensifyError): + xp.unique_values(x) + + +@pytest.mark.parametrize("invert", [False, True]) +def test_isin(invert): + for dtype in DTYPES: + for (d1, x1), (d2, x2) in zip( + inputs(dtype, (3, 4), seed=14), inputs(dtype, (5,), seed=15), strict=False + ): + expected = np.isin(d1, d2, invert=invert) + assert_array_equal(xp.isin(x1, x2, invert=invert), expected) + x = make(np.array([1, 0, 5, 2]), 0) + assert_array_equal(xp.isin(x, 2), np.isin([1, 0, 5, 2], 2)) + assert_array_equal(xp.isin(3, x), np.array(False)) + assert_array_equal(xp.isin(xp.asarray(1.0), x), np.array(True)) + with pytest.raises(TypeError, match="Array"): + xp.unique_values(np.zeros(3)) diff --git a/graphblas/tests/test_array_api_sorting.py b/graphblas/tests/test_array_api_sorting.py new file mode 100644 index 000000000..85cba8519 --- /dev/null +++ b/graphblas/tests/test_array_api_sorting.py @@ -0,0 +1,68 @@ +"""graphblas.array_api sort and argsort against numpy (stable sorts).""" + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp + +from .test_array_api_elementwise import assert_array_equal, inputs, make + +DTYPES = [np.dtype(t) for t in ("bool", "int8", "uint16", "int64", "float32", "float64")] + + +def _np_sort(a, axis, descending): + if descending: + return np.flip(np.sort(np.flip(a, axis=axis), axis=axis, kind="stable"), axis=axis) + return np.sort(a, axis=axis, kind="stable") + + +def _np_argsort(a, axis, descending): + if descending: + n = a.shape[axis] + flipped = np.argsort(np.flip(a, axis=axis), axis=axis, kind="stable") + return n - 1 - np.flip(flipped, axis=axis) + return np.argsort(a, axis=axis, kind="stable") + + +@pytest.mark.parametrize("descending", [False, True]) +def test_sort(descending): + for dtype in DTYPES: + for shape, axes in [((9,), [0, -1]), ((3, 6), [0, 1]), ((2, 3, 4), [0, 1, 2])]: + for dense, x in inputs(dtype, shape, seed=12): + for axis in axes: + out = xp.sort(x, axis=axis, descending=descending) + assert_array_equal(out, _np_sort(dense, axis, descending)) + assert out.fill_value == x.fill_value or np.isnan(out.fill_value) + out = xp.argsort(x, axis=axis, descending=descending) + assert_array_equal(out, _np_argsort(dense, axis, descending).astype(np.int64)) + + +def test_sort_keeps_sparsity(): + x = xp.zeros((1000, 1000)) + x[5, [3, 7, 9]] = np.array([2.0, -1.0, 4.0]) + x[8, 0] = -0.0 + out = xp.sort(x, axis=1) + assert out.nnz == 4 + row = out[5].todense() + assert (row[[0, 998, 999]] == [-1.0, 2.0, 4.0]).all() + assert np.signbit(out[8].todense()[0]) # -0.0 keeps its place among the zeros: first + with gb.config.set(array_api_densify_limit=10**5), pytest.raises(xp.DensifyError): + xp.argsort(x) + + +def test_sort_signed_zeros_and_nan_stable(): + a = np.array([0.0, -0.0, np.nan, 1.0, -0.0, 0.0, np.nan, -1.0]) + for fill in [0.0, -0.0, np.nan, 1.0]: + for x in [make(a, fill), make(a, fill, np.random.default_rng(1), explicit=0.6)]: + for descending in [False, True]: + assert_array_equal(xp.sort(x, descending=descending), _np_sort(a, 0, descending)) + + +def test_sort_errors(): + with pytest.raises(ValueError, match="dimension"): + xp.sort(xp.asarray(1.0)) + with pytest.raises(TypeError, match="real"): + xp.argsort(xp.zeros(2, dtype=xp.complex64)) + with pytest.raises(TypeError, match="Array"): + xp.sort(np.zeros(3)) diff --git a/graphblas/tests/test_array_api_statistical.py b/graphblas/tests/test_array_api_statistical.py new file mode 100644 index 000000000..b915d1f3b --- /dev/null +++ b/graphblas/tests/test_array_api_statistical.py @@ -0,0 +1,415 @@ +"""graphblas.array_api statistical and utility functions against numpy.""" + +import warnings + +import numpy as np +import pytest + +import graphblas as gb +import graphblas.array_api as xp + +from .test_array_api_elementwise import ( + DTYPES, + assert_array_equal, + check_invariants, + inputs, + make, + np_call, +) + +REDUCTIONS = [ + ("sum", np.sum), + ("prod", np.prod), + ("min", np.min), + ("max", np.max), + ("any", np.any), + ("all", np.all), + ("mean", np.mean), + ("var", np.var), + ("std", np.std), +] +SHAPES_AXES = [ + ((6,), [None, 0, -1, (0,), ()]), + ((3, 4), [None, 0, 1, -2, (0, 1), (1, 0), ()]), + ((2, 3, 4), [None, 0, 2, (0, 2), (1, 2), (0, 1, 2)]), + ((), [None, ()]), +] + + +def assert_close(x, expected, scale=1.0): + """``x`` equals ``expected`` up to rounding of floating-point sums (exactly otherwise); + ``scale`` is the size of the largest terms summed, which bounds the rounding. + """ + expected = np.asarray(expected) + check_invariants(x) + assert x.dtype == expected.dtype, (x.dtype, expected.dtype) + assert x.shape == expected.shape + got = x.todense() + if got.dtype.kind in "fc": + rtol = 1e-4 if got.dtype in (np.float32, np.complex64) else 1e-10 + atol = rtol * scale + np.testing.assert_allclose(got, expected, rtol=rtol, atol=atol, equal_nan=True) + # an infinity's sign, and NaN where numpy has NaN, are not rounding + assert (np.isnan(got) == np.isnan(expected)).all() + assert (np.isinf(got) == np.isinf(expected)).all() + else: + np.testing.assert_array_equal(got, expected) + + +@pytest.mark.parametrize(("name", "func"), REDUCTIONS, ids=[n for n, _ in REDUCTIONS]) +@pytest.mark.parametrize(("shape", "axes"), SHAPES_AXES, ids=str) +@pytest.mark.parametrize("keepdims", [False, True]) +def test_reductions(name, func, shape, axes, keepdims): + for dtype in DTYPES: + for dense, x in inputs(dtype, shape, seed=11): + for axis in axes: + kwargs = {"axis": axis, "keepdims": keepdims} + try: + expected = np_call(func, dense, **kwargs) + except (TypeError, ValueError) as exc: + with pytest.raises(type(exc)): + getattr(xp, name)(x, **kwargs) + continue + if name in {"min", "max"} and dtype.kind == "c": + with pytest.raises(TypeError): + getattr(xp, name)(x, **kwargs) + continue + if np.asarray(expected).dtype not in set(DTYPES): + continue + out = getattr(xp, name)(x, **kwargs) + exact = name in {"min", "max", "any", "all"} or dtype.kind in "biu" + exact = exact and name not in {"mean", "var", "std"} + try: + if exact: + assert_array_equal(out, np.asarray(expected)) + elif name == "prod" and not np.isfinite(dense).all(): + # (a product of infinities and zeros: NaN or not depends on the order) + pass + else: + finite = dense[np.isfinite(dense)] if dense.dtype.kind in "fc" else dense + scale = float(np.abs(finite.astype(np.complex128)).max(initial=1.0)) + assert_close(out, expected, scale * dense.size) + except AssertionError as exc: + raise AssertionError( + f"{name} of {dtype} fill {x.fill_value!r} axis={axis}: {exc}" + ) from None + + +@pytest.mark.parametrize("name", ["sum", "prod"]) +@pytest.mark.parametrize("dtype", ["int8", "uint8", "int32", "float32", "bool"]) +def test_sum_prod_dtype(name, dtype): + func = getattr(np, name) + for dense, x in inputs(np.dtype(dtype), (3, 5), seed=2): + for out_dtype in [None, np.dtype("int64"), np.dtype("float64"), np.dtype("uint64")]: + if out_dtype is not None and dtype == "float32" and out_dtype.kind in "iu": + continue + expected = np_call(func, dense, axis=1, dtype=out_dtype) + out = getattr(xp, name)(x, axis=1, dtype=out_dtype) + assert_close(out, expected) + + +def test_sum_wraps_like_numpy(): + a = np.array([[100, 100, 0, 0], [0, 0, 0, 50]], dtype=np.int8) + x = make(a, 0) + assert_array_equal(xp.sum(x, axis=1, dtype=xp.int8), np.sum(a, axis=1, dtype=np.int8)) + # a nonzero fill counted many times over + x = make(np.full((2, 300), 7, dtype=np.int8), 0) + 0 + y = xp.full((2, 300), 7, dtype=xp.int8) + assert_array_equal( + xp.sum(y, axis=1, dtype=xp.int8), np.full(2, 7 * 300, np.int64).astype(np.int8) + ) + assert_array_equal( + xp.prod(y, axis=1, dtype=xp.int8), + np.prod(np.full((2, 300), 7, np.int8), axis=1, dtype=np.int8), + ) + assert_array_equal(xp.prod(y, axis=1), np.prod(np.full((2, 300), 7, np.int8), axis=1)) + del x + + +def test_min_max_nan_propagates(): + a = np.array([[1.0, np.nan, 0.0, 0.0], [0.0, 0.0, 2.0, 0.0], [0, 0, 0, 0]]) + x = make(a, 0.0) + for name in ["min", "max"]: + for axis in [None, 0, 1]: + assert_array_equal(getattr(xp, name)(x, axis=axis), getattr(np, name)(a, axis=axis)) + x = make(np.array([[1.0, np.nan], [np.nan, np.nan]]), np.nan) + assert_array_equal(xp.max(x, axis=1), np.array([np.nan, np.nan])) + + +def test_empty_axis(): + x = xp.zeros((0, 3)) + assert_array_equal(xp.sum(x, axis=0), np.zeros(3)) + assert_array_equal(xp.prod(x, axis=0), np.ones(3)) + assert_array_equal(xp.any(x, axis=0), np.zeros(3, bool)) + assert_array_equal(xp.all(x, axis=0), np.ones(3, bool)) + assert_array_equal(xp.sum(x, axis=1), np.zeros(0)) + with pytest.raises(ValueError, match="empty"): + xp.max(x, axis=0) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + assert_array_equal(xp.mean(x, axis=0), np.mean(np.zeros((0, 3)), axis=0)) + assert_array_equal(xp.var(x, axis=0), np.var(np.zeros((0, 3)), axis=0)) + assert_array_equal(xp.max(x, axis=1), np.zeros(0)) + + +@pytest.mark.parametrize("correction", [0, 1, 2.5, 4, 10]) +def test_var_correction(correction): + rng = np.random.default_rng(3) + a = np.where(rng.random((3, 4)) < 0.5, 0.0, np.round(rng.normal(size=(3, 4)), 2)) + for fill in [0.0, 1.5]: + dense = np.where(a == 0, fill, a) + x = make(dense, fill) + for axis in [None, 0, 1, ()]: + for name in ["var", "std"]: + expected = np_call(getattr(np, name), dense, axis=axis, ddof=correction) + out = np_call(getattr(xp, name), x, axis=axis, correction=correction) + assert_close(out, expected) + + +def test_reductions_on_full_arrays(): + a = np.arange(1.0, 13.0).reshape(3, 4) + x = make(a, 0.0) + assert x.nnz == 12 + for name, func in REDUCTIONS: + for axis in [None, 0, 1]: + assert_close(getattr(xp, name)(x, axis=axis), func(a, axis=axis)) + + +def test_reductions_of_fill_only_arrays(): + x = xp.full((4, 5), 2.5) + assert_close(xp.sum(x, axis=1), np.full(4, 12.5)) + assert xp.sum(x, axis=1).nnz == 0 + assert_close(xp.prod(x), np.prod(np.full((4, 5), 2.5))) + assert_close(xp.var(x, axis=0), np.zeros(5)) + x = xp.full((3, 3), np.inf) + assert_array_equal(np_call(xp.var, x, axis=0), np_call(np.var, np.full((3, 3), np.inf), axis=0)) + + +@pytest.mark.parametrize("name", ["cumulative_sum", "cumulative_prod"]) +@pytest.mark.parametrize("include_initial", [False, True]) +def test_cumulative(name, include_initial): + func = np.cumsum if name == "cumulative_sum" else np.cumprod + for dtype in [np.dtype(t) for t in ("bool", "int8", "uint16", "int64", "float32", "float64")]: + for shape, axes in [((7,), [None, 0, -1]), ((3, 4), [0, 1, -1]), ((2, 3, 2), [0, 2])]: + for dense, x in inputs(dtype, shape, seed=5): + for axis in axes: + np_axis = 0 if axis is None else axis + expected = np_call(func, dense, axis=np_axis) + if include_initial: + initial = np.zeros if name == "cumulative_sum" else np.ones + shp = list(expected.shape) + shp[np_axis] = 1 + expected = np.concatenate( + [initial(shp, expected.dtype), expected], axis=np_axis + ) + out = getattr(xp, name)(x, axis=axis, include_initial=include_initial) + assert_close(out, expected) + x = make(np.ones((2, 3)), 0.0) + with pytest.raises(ValueError, match="axis"): + getattr(xp, name)(x) + out = getattr(xp, name)(make(np.array([1, 2, 3], np.int8), 0), dtype=xp.float32) + assert out.dtype == xp.float32 + + +def test_cumulative_is_guarded(): + with gb.config.set(array_api_densify_limit=10), pytest.raises(xp.DensifyError): + xp.cumulative_sum(xp.full(100, 1.0)) + # with the identity as fill only the positions from each fiber's first entry are stored + # (review finding L8) + y = xp.zeros((20000, 20000)) + y[[1, 5, 9], [3, 19000, 7]] = 1.0 + assert xp.cumulative_sum(y, axis=1).nnz == 19997 + 1000 + 19993 + assert xp.cumulative_sum(xp.zeros(10**12)).nnz == 0 + with gb.config.set(array_api_densify_limit=100), pytest.raises(xp.DensifyError): + xp.cumulative_sum(y, axis=1) + + +def test_cumulative_sparse_uneven_fibers(): + # one long fiber among many short ones: the fibers accumulate one by one + a = np.zeros((200, 300)) + a[0] = np.arange(1.0, 301.0) + a[np.arange(1, 200), np.arange(1, 200)] = -0.5 + x = make(a, 0.0) + assert_array_equal(xp.cumulative_sum(x, axis=1), np.cumsum(a, axis=1)) + c = xp.zeros(1, dtype=xp.complex64) + assert_array_equal(xp.cumulative_sum(c, include_initial=True), np.zeros(2, np.complex64)) + + +@pytest.mark.parametrize("name", ["cumulative_sum", "cumulative_prod"]) +def test_cumulative_sparse_exact(name): + # runs of unstored positions add (multiply by) the fill: 0.0 + -0.0 is 0.0, as in numpy + func = np.cumsum if name == "cumulative_sum" else np.cumprod + identity = 0 if name == "cumulative_sum" else 1 + rng = np.random.default_rng(17) + specials = np.array([-0.0, 1.5, -2.0, 3.0, np.inf, np.nan, 0.5]) + for dtype in [np.float64, np.float32, np.int16, np.uint8, np.bool_]: + for shape in [(9,), (3, 7), (2, 3, 5), (1, 40)]: + if np.dtype(dtype).kind == "f": + vals = rng.choice(specials, shape) + else: + vals = rng.integers(0, 4, shape) + a = np.where(rng.random(shape) < 0.6, identity, vals).astype(dtype) + x = make(a, identity, rng) + for axis in range(len(shape)): + for include_initial in [False, True]: + expected = np_call(func, a, axis=axis) + if include_initial: + s = list(expected.shape) + s[axis] = 1 + initial = np.full(s, identity, expected.dtype) + expected = np.concatenate([initial, expected], axis=axis) + out = np_call(getattr(xp, name), x, axis=axis, include_initial=include_initial) + assert_array_equal(out, expected) + + +@pytest.mark.parametrize("n", [0, 1, 2, 5]) +def test_diff(n): + for dtype in [np.dtype(t) for t in ("bool", "int16", "uint8", "float64", "complex64")]: + for shape, axes in [((6,), [0, -1]), ((3, 4), [0, 1])]: + for dense, x in inputs(dtype, shape, seed=6): + for axis in axes: + expected = np_call(np.diff, dense, n=n, axis=axis) + assert_array_equal(xp.diff(x, n=n, axis=axis), expected) + + +def test_diff_prepend_append(): + a = np.array([[1.0, 0.0, 3.0], [0.0, 0.0, -1.0]]) + x = make(a, 0.0) + p = np.array([[5.0], [0.0]]) + q = np.array([[0.0, 2.0, 0.0]]) + assert_array_equal(xp.diff(x, axis=1, prepend=make(p, 0.0)), np.diff(a, axis=1, prepend=p)) + assert_array_equal(xp.diff(x, axis=0, append=make(q, 1.0)), np.diff(a, axis=0, append=q)) + assert_array_equal(xp.diff(x, axis=1, append=2.0), np.diff(a, axis=1, append=2.0)) + assert_array_equal( + xp.diff(x, axis=1, n=2, prepend=make(p, 0.0), append=make(p, 7.0)), + np.diff(a, axis=1, n=2, prepend=p, append=p), + ) + + +def test_axis_errors(): + x = make(np.ones((2, 3)), 0.0) + with pytest.raises(IndexError): + xp.sum(x, axis=2) + with pytest.raises(ValueError, match="repeated"): + xp.sum(x, axis=(0, -2)) + + +def test_var_of_a_huge_sparse_array(): + # The means are not stored densely: a 1e9 x 1e9 array with 4 entries + n = 10**9 + coords = (np.array([0, 5, 5, 10**8]), np.array([3, 3, 7, 3])) + x = xp.from_coo(coords, np.array([1.0, 2.0, 4.0, 8.0]), (n, n)) + v = xp.var(x, axis=0) + assert v.nnz == 2 + assert v.fill_value == 0.0 + for col, values in [(3, [1.0, 2.0, 8.0]), (7, [4.0])]: + values = np.array(values) + mean = values.sum() / n + expected = (((values - mean) ** 2).sum() + (n - values.size) * mean**2) / n + assert float(v[col]) == pytest.approx(expected, rel=1e-12) + + +def test_sum_starts_from_positive_zero(): + # numpy's sum starts from +0.0: a sum of -0.0 is 0.0 (review finding M1) + for a in [ + np.array([-0.0, -0.0]), + np.array([[-0.0, 0.0], [-0.0, 1.0]]), + np.array([-0.0], np.float32), + ]: + x = make(a, 0.0) + for axis in [None, -1]: + assert_array_equal(xp.sum(x, axis=axis), np.sum(a, axis=axis)) + assert_array_equal(xp.mean(x, axis=axis), np.mean(a, axis=axis)) + c = np.array([complex(-0.0, -0.0)]) + assert_array_equal(xp.sum(make(c, 0j)), np.sum(c)) + x = make(np.array([1.0, -0.0]), -0.0) + assert_array_equal(xp.sum(x, axis=0), np.sum(np.array([1.0, -0.0]))) + assert_array_equal(xp.sum(xp.full((3,), -0.0)), np.sum(np.full(3, -0.0))) + + +def test_min_max_of_signed_zeros(): + # GraphBLAS's fmin and fmax take either zero, and numpy orders -0.0 below 0.0: the Array + # API leaves that order unspecified (review finding M2 closed by the contract); NaN + # propagates + cases = [ + (np.array([-0.0, 0.0]), 1.0), + (np.array([0.0, -0.0]), 1.0), + (np.array([[0.0, 2.0], [-0.0, 3.0]]), 5.0), + (np.array([[np.nan, -0.0], [np.nan, 0.0]]), np.nan), + (np.array([[0.0, 1.0], [2.0, 0.0]]), -0.0), + (np.array([[-0.0, 1.0], [2.0, -0.0]]), 0.0), + ] + for a, fill in cases: + for x in [make(a, fill), make(np.where(a == 0, fill, a), fill)]: + dense = x.todense() + for name in ["min", "max"]: + for axis in [None, 0, -1]: + assert_array_equal( + getattr(xp, name)(x, axis=axis), + getattr(np, name)(dense, axis=axis), + zero_signs=False, + ) + + +def test_complex_sum_of_infinite_fill(): + # k copies of inf+0j sum to inf+0j, not inf+nanj (review finding M3) + for fill in [complex(np.inf, 0), complex(np.nan, 0), complex(0, -np.inf)]: + a = np.array([[1 + 1j, fill, fill], [fill, fill, fill]]) + x = make(a, fill) + for axis in [None, 0, 1]: + for name in ["sum", "mean"]: + assert_array_equal( + np_call(getattr(xp, name), x, axis=axis), + np_call(getattr(np, name), a, axis=axis), + ) + + +def test_sum_in_bool(): + # a sum of booleans in bool is their logical or (review finding M5) + a = np.array([[1, 0, 3], [0, 0, 2]]) + for fill in [0, 2]: + x = make(np.where(a == 0, fill, a), fill) + dense = x.todense() + for axis in [None, 0, 1]: + assert_array_equal( + xp.sum(x, axis=axis, dtype=xp.bool), np.sum(dense, axis=axis, dtype=bool) + ) + assert_array_equal( + xp.prod(x, axis=axis, dtype=xp.bool), np.prod(dense, axis=axis, dtype=bool) + ) + + +def test_mean_of_empty_complex(): + # 0j / 0 is nan+nanj, as numpy's mean of nothing (review finding L1) + for dt in [np.complex128, np.complex64, np.float32]: + assert_array_equal( + np_call(xp.mean, xp.zeros((0,), dtype=dt)), np_call(np.mean, np.zeros(0, dt)) + ) + assert_array_equal( + np_call(xp.mean, xp.zeros((2, 0), dtype=dt), axis=1), + np_call(np.mean, np.zeros((2, 0), dt), axis=1), + ) + + +def test_diff_n0_ignores_prepend_append(): + # (review finding L4) + a = np.array([1.0, 0.0, 2.0]) + x = make(a, 0.0) + assert_array_equal( + xp.diff(x, n=0, prepend=make(np.array([5.0]), 0.0)), np.diff(a, n=0, prepend=5.0) + ) + assert_array_equal(xp.diff(x, n=0, append=7.0), np.diff(a, n=0, append=7.0)) + + +def test_reductions_of_fibers_past_2_63_positions(): + # (review finding L6) a fiber of 2**63 positions: counts do not overflow + shape = (2**21,) * 3 + y = xp.full(shape, 1.0) + y[0, 0, 0] = 3.0 + assert float(xp.sum(y)) == float(2**63 + 2) + assert float(xp.sum(y, axis=(1, 2))[0]) == 2.0**42 + 2 + assert float(xp.max(y)) == 3.0 + z = xp.full(shape, 1, dtype=xp.int64) + z[0, 0, 0] = 3 + assert int(xp.sum(z)) == 3 + (2**63 - 1) - 2**64 # (wraps, as integer sums do) + assert int(xp.prod(z)) == 3 diff --git a/graphblas/tests/test_tensor.py b/graphblas/tests/test_tensor.py index 4713b119c..bad951547 100644 --- a/graphblas/tests/test_tensor.py +++ b/graphblas/tests/test_tensor.py @@ -383,6 +383,10 @@ def test_dup_and_clear(): D.clear() assert D.nvals == 0 check(T, values, present) + # A new dtype keeps an n x 1 stored Matrix held by row (Matrix.dup holds it by column) + for shape in [(2, 1), (3, 1, 1)]: + T, values, present = random_tensor(shape, rng, density=1.0) + check(T.dup(dtypes.FP32), values.astype(np.float32), present) def test_vanilla_check(monkeypatch): @@ -2812,6 +2816,72 @@ def test_tensordot_semirings_masks_and_errors(): assert ("ttv", "entries") in {(r["point"], r["choice"]) for r in recs} +def test_private_map_values(): + # Tensor._map_values (for graphblas.array_api): numpy over the stored values, in place + rng = np.random.default_rng(33) + X, xv, xp = random_tensor((4, 5, 3), rng, density=0.4) + Y = Tensor(dtypes.FP64, (4, 5, 3)) + Y(X.S) << X.apply(binary.plus, right=1.0) # (X's pattern) + Y = Y.relayout(((2,), (0, 1))) # (another layout) + R = X._map_values(np.hypot, [Y]) + check(R, np.where(xp, np.hypot(xv, xv + 1), 0), xp, approx=True) + assert R.layout == X.layout + # one Matrix shared by two Tensors, a dtype change, an iso Tensor, and no entries + R = X._map_values(lambda a, b: a * b > 1, [X.dup()], dtype=dtypes.BOOL) + check(R, np.where(xp, xv * xv > 1, False), xp) + iso = Tensor.from_coo(([0, 2], [1, 3]), 2.5, shape=(3, 4)) + assert iso.ss.is_iso + check(iso._map_values(np.sqrt), np.where(iso.to_dense(0) != 0, np.sqrt(2.5), 0), + iso.to_dense(0) != 0, approx=True) # fmt: skip + E = Tensor(dtypes.INT64, (2, 3)) + assert E._map_values(np.negative, dtype=dtypes.FP32).dtype == dtypes.FP32 + assert E._map_values(np.negative).nvals == 0 + # The fallback through coordinates + from graphblas.core.tensor import _map_values_coo + + T = Tensor.from_coo([[0, 1], [1, 2]], [1.0, 2.0], shape=(2, 3)) + assert _map_values_coo(np.add, [T, T], dtypes.FP64).to_coo()[1].tolist() == [2.0, 4.0] + assert _map_values_coo(np.negative, [T], dtypes.FP64).to_coo()[1].tolist() == [-1.0, -2.0] + out = _map_values_coo(lambda v: np.float64(7.0), [T], dtypes.FP64) + assert out.to_coo()[1].tolist() == [7.0, 7.0] + + +def test_masked_contraction_takes_a_transpose_descriptor(): + # Under a mask, an operand stored as the transpose of what the product needs is not + # copied (SuiteSparse's masked dot product, an SDDMM); the values are the masked product + from graphblas.core import tensor as tensor_mod + from graphblas.core import tensor_census as census + + rng = np.random.default_rng(32) + X, xv, xp = random_tensor((6, 4), rng, density=0.7) + Y, yv, yp = random_tensor((5, 4), rng, density=0.7) + Mk, _, mp = random_tensor((6, 5), rng, dtype=dtypes.BOOL, density=0.4) + ev, ep = np_matmul(xv, xp, yv.T, yp.T) + rv, rp = model_update(np.zeros((6, 5)), np.zeros((6, 5), bool), ev, ep, mask=mp) + XT = X.relayout(((1,), (0,))) # (X's transpose in storage) + for expr in [ + X.matmul(Y.T), + X.tensordot(Y, ([1], [1])), + tensor_mod.einsum("ik,jk->ij", X, Y), + XT.T.T.matmul(Y.T), + ]: + census.start() + try: + R = expr.new(mask=Mk.S) + recs = census.stop() + finally: + census.stop() + census.records.clear() + check(R, rv, rp, approx=True) + moved = [r["moved"] for r in recs if r["point"] == "contract" and "moved" in r] + assert moved == [[False, False]], moved + # Unmasked, the planner's relayout stays + check(X.matmul(Y.T).new(), ev, ep, approx=True) + C, cv, cp = random_tensor((6, 5), rng) + C(Mk.S, binary.plus) << X.matmul(Y.T) + check(C, *model_update(cv, cp, ev, ep, mask=mp, accum=np.add), approx=True) + + # -------------------------------------------------------------------------------------------- # Core hooks (Matrix, Vector, Scalar and operators meet a Tensor) @@ -2973,6 +3043,21 @@ def test_broadcast(xshape, yshape): check((X > Y).new(), *np_bcast("mult", xv, xp, yv, yp, np.greater)) +@pytest.mark.parametrize( + ("xshape", "yshape"), [((2, 1), (1, 0)), ((1, 0), (2, 1)), ((0, 1), (1, 3)), ((2, 0, 1), (4,))] +) +def test_broadcast_to_no_positions(xshape, yshape): + # A broadcast result with no positions: ewise_mult's diagonal scaling multiplied + # mismatched shapes (GrB_DIMENSION_MISMATCH) + X = Tensor.from_dense(np.ones(xshape)) + Y = Tensor.from_dense(np.ones(yshape)) + shape = np.broadcast_shapes(xshape, yshape) + for method in ["ewise_mult", "ewise_add"]: + R = getattr(X, method)(Y).new() + assert R.shape == shape + assert R.nvals == 0 + + def test_broadcast_operands_and_dtypes(): rng = np.random.default_rng(42) X, xv, xp = random_tensor((3, 4, 5), rng, dtype=dtypes.INT64) diff --git a/pyproject.toml b/pyproject.toml index c46d012eb..21dd78bf8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -115,6 +115,7 @@ all = [ packages = [ "graphblas", "graphblas.agg", + "graphblas.array_api", "graphblas.binary", "graphblas.core", "graphblas.core.operator", @@ -229,7 +230,7 @@ skip_empty = true exclude_lines = ["pragma: no cover", "raise AssertionError", "raise NotImplementedError"] [tool.codespell] -ignore-words-list = "bu,coo,ba" +ignore-words-list = "bu,coo,ba,fro" [tool.ruff] # https://github.com/charliermarsh/ruff/