From c90d13288cec74d99d320fba9c184c8f69319768 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Tue, 14 Oct 2025 21:01:49 -0500 Subject: [PATCH 1/3] Reorganize the `sparse` module --- pytensor/link/jax/dispatch/sparse.py | 2 +- pytensor/sparse/__init__.py | 1 + pytensor/sparse/basic.py | 2756 ++------------------------ pytensor/sparse/linalg.py | 93 + pytensor/sparse/math.py | 2092 +++++++++++++++++++ pytensor/sparse/rewriting.py | 70 +- pytensor/sparse/sharedvar.py | 2 +- pytensor/sparse/variable.py | 476 +++++ tests/sparse/test_basic.py | 1967 +----------------- tests/sparse/test_linalg.py | 26 + tests/sparse/test_math.py | 1462 ++++++++++++++ tests/sparse/test_rewriting.py | 17 +- tests/sparse/test_var.py | 238 --- tests/sparse/test_variable.py | 235 +++ 14 files changed, 4671 insertions(+), 4766 deletions(-) create mode 100644 pytensor/sparse/linalg.py create mode 100644 pytensor/sparse/math.py create mode 100644 pytensor/sparse/variable.py create mode 100644 tests/sparse/test_linalg.py create mode 100644 tests/sparse/test_math.py delete mode 100644 tests/sparse/test_var.py create mode 100644 tests/sparse/test_variable.py diff --git a/pytensor/link/jax/dispatch/sparse.py b/pytensor/link/jax/dispatch/sparse.py index 4e58199ae0..b127f48c12 100644 --- a/pytensor/link/jax/dispatch/sparse.py +++ b/pytensor/link/jax/dispatch/sparse.py @@ -3,7 +3,7 @@ from pytensor.graph.basic import Constant from pytensor.link.jax.dispatch import jax_funcify, jax_typify -from pytensor.sparse.basic import Dot, StructuredDot +from pytensor.sparse.math import Dot, StructuredDot from pytensor.sparse.type import SparseTensorType diff --git a/pytensor/sparse/__init__.py b/pytensor/sparse/__init__.py index eca4e14483..25c9bc0812 100644 --- a/pytensor/sparse/__init__.py +++ b/pytensor/sparse/__init__.py @@ -1,5 +1,6 @@ from pytensor.sparse import rewriting, sharedvar from pytensor.sparse.basic import * +from pytensor.sparse.math import * from pytensor.sparse.sharedvar import sparse_constructor as shared from pytensor.sparse.type import SparseTensorType, _is_sparse diff --git a/pytensor/sparse/basic.py b/pytensor/sparse/basic.py index f22cd10454..60ac79f149 100644 --- a/pytensor/sparse/basic.py +++ b/pytensor/sparse/basic.py @@ -8,7 +8,6 @@ """ -from typing import Literal from warnings import warn import numpy as np @@ -19,55 +18,18 @@ from pytensor import _as_symbolic, as_symbolic from pytensor import scalar as ps from pytensor.configdefaults import config -from pytensor.gradient import DisconnectedType, grad_not_implemented, grad_undefined +from pytensor.gradient import DisconnectedType, grad_undefined from pytensor.graph.basic import Apply, Constant, Variable from pytensor.graph.op import Op -from pytensor.link.c.op import COp from pytensor.link.c.type import generic from pytensor.sparse.type import SparseTensorType, _is_sparse -from pytensor.sparse.utils import hash_from_sparse from pytensor.tensor import basic as ptb from pytensor.tensor.basic import Split -from pytensor.tensor.math import ( - _conj, - arcsin, - arcsinh, - arctan, - arctanh, - ceil, - deg2rad, - exp, - expm1, - floor, - log, - log1p, - maximum, - minimum, - rad2deg, - round_half_to_even, - sigmoid, - sign, - sin, - sinh, - sqr, - sqrt, - tan, - tanh, - trunc, -) -from pytensor.tensor.math import add as pt_add -from pytensor.tensor.math import dot as pt_dot -from pytensor.tensor.math import pow as pt_pow -from pytensor.tensor.shape import shape, specify_broadcastable -from pytensor.tensor.slinalg import BaseBlockDiagonal, _largest_common_dtype -from pytensor.tensor.type import TensorType, iscalar, ivector, scalar, tensor, vector +from pytensor.tensor.math import minimum +from pytensor.tensor.shape import specify_broadcastable +from pytensor.tensor.type import TensorType, ivector, scalar, tensor, vector from pytensor.tensor.type import continuous_dtypes as tensor_continuous_dtypes from pytensor.tensor.type import discrete_dtypes as tensor_discrete_dtypes -from pytensor.tensor.variable import ( - TensorConstant, - TensorVariable, - _tensor_py_operators, -) sparse_formats = ["csc", "csr"] @@ -186,6 +148,8 @@ def as_sparse_variable(x, name=None, ndim=None, **kwargs): ) return x try: + from pytensor.sparse.variable import constant + return constant(x, name=name) except TypeError: raise TypeError(f"Cannot convert {x} to SparseTensorType", type(x)) @@ -196,17 +160,6 @@ def as_sparse_variable(x, name=None, ndim=None, **kwargs): as_sparse_or_tensor_variable = as_symbolic -def constant(x, name=None): - if not isinstance(x, scipy.sparse.spmatrix): - raise TypeError("sparse.constant must be called on a scipy.sparse.spmatrix") - try: - return SparseConstant( - SparseTensorType(format=x.format, dtype=x.dtype), x.copy(), name=name - ) - except TypeError: - raise TypeError(f"Could not convert {x} to SparseTensorType", type(x)) - - def sp_ones_like(x): """ Construct a sparse matrix of ones with the same sparsity pattern. @@ -253,252 +206,6 @@ def sp_zeros_like(x): ) -def override_dense(*methods): - def decorate(cls): - def native(method): - original = getattr(cls.__base__, method) - - def to_dense(self, *args, **kwargs): - self = self.toarray() - new_args = [ - arg.toarray() - if hasattr(arg, "type") and isinstance(arg.type, SparseTensorType) - else arg - for arg in args - ] - warn( - f"Method {method} is not implemented for sparse variables. The variable will be converted to dense." - ) - return original(self, *new_args, **kwargs) - - return to_dense - - for method in methods: - setattr(cls, method, native(method)) - return cls - - return decorate - - -@override_dense( - "__abs__", - "__ceil__", - "__floor__", - "__trunc__", - "transpose", - "any", - "all", - "flatten", - "ravel", - "arccos", - "arcsin", - "arctan", - "arccosh", - "arcsinh", - "arctanh", - "ceil", - "cos", - "cosh", - "deg2rad", - "exp", - "exp2", - "expm1", - "floor", - "log", - "log10", - "log1p", - "log2", - "rad2deg", - "sin", - "sinh", - "sqrt", - "tan", - "tanh", - "copy", - "prod", - "mean", - "var", - "std", - "min", - "max", - "argmin", - "argmax", - "round", - "trace", - "cumsum", - "cumprod", - "ptp", - "squeeze", - "diagonal", - "__and__", - "__or__", - "__xor__", - "__pow__", - "__mod__", - "__divmod__", - "__truediv__", - "__floordiv__", - "reshape", - "dimshuffle", -) -class _sparse_py_operators(_tensor_py_operators): - T = property( - lambda self: transpose(self), doc="Return aliased transpose of self (read-only)" - ) - - def astype(self, dtype): - return cast(self, dtype) - - def __neg__(self): - return neg(self) - - def __add__(left, right): - return add(left, right) - - def __radd__(right, left): - return add(left, right) - - def __sub__(left, right): - return sub(left, right) - - def __rsub__(right, left): - return sub(left, right) - - def __mul__(left, right): - return mul(left, right) - - def __rmul__(left, right): - return mul(left, right) - - # comparison operators - - def __lt__(self, other): - return lt(self, other) - - def __le__(self, other): - return le(self, other) - - def __gt__(self, other): - return gt(self, other) - - def __ge__(self, other): - return ge(self, other) - - def __dot__(left, right): - return structured_dot(left, right) - - def __rdot__(right, left): - return structured_dot(left, right) - - def sum(self, axis=None, sparse_grad=False): - return sp_sum(self, axis=axis, sparse_grad=sparse_grad) - - dot = __dot__ - - def toarray(self): - return dense_from_sparse(self) - - @property - def shape(self): - # TODO: The plan is that the ShapeFeature in ptb.opt will do shape - # propagation and remove the dense_from_sparse from the graph. This - # will *NOT* actually expand your sparse matrix just to get the shape. - return shape(dense_from_sparse(self)) - - ndim = property(lambda self: self.type.ndim) - dtype = property(lambda self: self.type.dtype) - - # Note that the `size` attribute of sparse matrices behaves differently - # from dense matrices: it is the number of elements stored in the matrix - # rather than the total number of elements that may be stored. Note also - # that stored zeros *do* count in the size. - size = property(lambda self: csm_data(self).size) - - def zeros_like(model): - return sp_zeros_like(model) - - def __getitem__(self, args): - if not isinstance(args, tuple): - args = (args,) - - if len(args) == 2: - scalar_arg_1 = ( - np.isscalar(args[0]) or getattr(args[0], "type", None) == iscalar - ) - scalar_arg_2 = ( - np.isscalar(args[1]) or getattr(args[1], "type", None) == iscalar - ) - if scalar_arg_1 and scalar_arg_2: - ret = get_item_scalar(self, args) - elif isinstance(args[0], list): - ret = get_item_2lists(self, args[0], args[1]) - else: - ret = get_item_2d(self, args) - elif isinstance(args[0], list): - ret = get_item_list(self, args[0]) - else: - ret = get_item_2d(self, args) - return ret - - def conj(self): - return conjugate(self) - - -class SparseVariable(_sparse_py_operators, TensorVariable): - format = property(lambda self: self.type.format) - - def __str__(self): - return f"{self.__class__.__name__}{{{self.format},{self.dtype}}}" - - def __repr__(self): - return str(self) - - -class SparseConstantSignature(tuple): - def __eq__(self, other): - (a, b), (x, y) = self, other - return ( - a == x - and (b.dtype == y.dtype) - and (type(b) is type(y)) - and (b.shape == y.shape) - and (abs(b - y).sum() < 1e-6 * b.nnz) - ) - - def __ne__(self, other): - return not self == other - - def __hash__(self): - (a, b) = self - return hash(type(self)) ^ hash(a) ^ hash(type(b)) - - def pytensor_hash(self): - (_, d) = self - return hash_from_sparse(d) - - -class SparseConstant(SparseVariable, TensorConstant): - format = property(lambda self: self.type.format) - - def signature(self): - assert self.data is not None - return SparseConstantSignature((self.type, self.data)) - - def __str__(self): - return f"{self.__class__.__name__}{{{self.format},{self.dtype},shape={self.data.shape},nnz={self.data.nnz}}}" - - def __repr__(self): - return str(self) - - @property - def unique_value(self): - return None - - -SparseTensorType.variable_type = SparseVariable -SparseTensorType.constant_type = SparseConstant - - # for more dtypes, call SparseTensorType(format, dtype) def matrix(format, name=None, dtype=None): if dtype is None: @@ -1610,6 +1317,8 @@ def perform(self, node, inputs, outputs): z[0] = y def grad(self, inputs, gout): + from pytensor.sparse.math import sp_sum + (x, s) = inputs (gz,) = gout return [col_scale(gz, s), sp_sum(x * gz, axis=0)] @@ -1659,6 +1368,8 @@ def perform(self, node, inputs, outputs): z[0] = scipy.sparse.csc_matrix((y_data, indices, indptr), (M, N)) def grad(self, inputs, gout): + from pytensor.sparse.math import sp_sum + (x, s) = inputs (gz,) = gout return [row_scale(gz, s), sp_sum(x * gz, axis=1)] @@ -1724,126 +1435,6 @@ def row_scale(x, s): return col_scale(x.T, s).T -class SpSum(Op): - """ - - WARNING: judgement call... - We are not using the structured in the comparison or hashing - because it doesn't change the perform method therefore, we - *do* want Sums with different structured values to be merged - by the merge optimization and this requires them to compare equal. - """ - - __props__ = ("axis",) - - def __init__(self, axis=None, sparse_grad=True): - super().__init__() - self.axis = axis - self.structured = sparse_grad - if self.axis not in (None, 0, 1): - raise ValueError("Illegal value for self.axis.") - - def make_node(self, x): - x = as_sparse_variable(x) - assert x.format in ("csr", "csc") - - if self.axis is not None: - out_shape = (None,) - else: - out_shape = () - - z = TensorType(dtype=x.dtype, shape=out_shape)() - return Apply(self, [x], [z]) - - def perform(self, node, inputs, outputs): - (x,) = inputs - (z,) = outputs - if self.axis is None: - z[0] = np.asarray(x.sum()) - else: - z[0] = np.asarray(x.sum(self.axis)).ravel() - - def grad(self, inputs, gout): - (x,) = inputs - (gz,) = gout - if x.dtype not in continuous_dtypes: - return [x.zeros_like(dtype=config.floatX)] - if self.structured: - if self.axis is None: - r = gz * pytensor.sparse.sp_ones_like(x) - elif self.axis == 0: - r = col_scale(pytensor.sparse.sp_ones_like(x), gz) - elif self.axis == 1: - r = row_scale(pytensor.sparse.sp_ones_like(x), gz) - else: - raise ValueError("Illegal value for self.axis.") - else: - o_format = x.format - x = dense_from_sparse(x) - if _is_sparse_variable(gz): - gz = dense_from_sparse(gz) - if self.axis is None: - r = ptb.second(x, gz) - else: - ones = ptb.ones_like(x) - if self.axis == 0: - r = specify_broadcastable(gz.dimshuffle("x", 0), 0) * ones - elif self.axis == 1: - r = specify_broadcastable(gz.dimshuffle(0, "x"), 1) * ones - else: - raise ValueError("Illegal value for self.axis.") - r = SparseFromDense(o_format)(r) - return [r] - - def infer_shape(self, fgraph, node, shapes): - r = None - if self.axis is None: - r = [()] - elif self.axis == 0: - r = [(shapes[0][1],)] - else: - r = [(shapes[0][0],)] - return r - - def __str__(self): - return f"{self.__class__.__name__}{{axis={self.axis}}}" - - -def sp_sum(x, axis=None, sparse_grad=False): - """ - Calculate the sum of a sparse matrix along the specified axis. - - It operates a reduction along the specified axis. When `axis` is `None`, - it is applied along all axes. - - Parameters - ---------- - x - Sparse matrix. - axis - Axis along which the sum is applied. Integer or `None`. - sparse_grad : bool - `True` to have a structured grad. - - Returns - ------- - object - The sum of `x` in a dense format. - - Notes - ----- - The grad implementation is controlled with the `sparse_grad` parameter. - `True` will provide a structured grad and `False` will provide a regular - grad. For both choices, the grad returns a sparse matrix having the same - format as `x`. - - This op does not return a sparse matrix, but a dense tensor matrix. - - """ - - return SpSum(axis, sparse_grad)(x) - - class Diag(Op): """Extract the diagonal of a square sparse matrix as a dense vector. @@ -2023,2168 +1614,236 @@ def clean(x): return ensure_sorted_indices(remove0(x)) -class AddSS(Op): - # add(sparse, sparse). - # see the doc of add() for more detail. - __props__ = () +class Stack(Op): + __props__ = ("format", "dtype") + + def __init__(self, format=None, dtype=None): + if format is None: + self.format = "csc" + else: + self.format = format + + if dtype is None: + raise ValueError("The output dtype must be specified.") + self.dtype = dtype + + def make_node(self, *mat): + if not mat: + raise ValueError("Cannot join an empty list of sparses.") + var = [as_sparse_variable(x) for x in mat] + + for x in var: + assert x.format in ("csr", "csc") - def make_node(self, x, y): - x, y = map(as_sparse_variable, [x, y]) - assert x.format in ("csr", "csc") - assert y.format in ("csr", "csc") - out_dtype = ps.upcast(x.type.dtype, y.type.dtype) return Apply( - self, [x, y], [SparseTensorType(dtype=out_dtype, format=x.type.format)()] + self, var, [SparseTensorType(dtype=self.dtype, format=self.format)()] ) - def perform(self, node, inputs, outputs): - (x, y) = inputs + def __str__(self): + return f"{self.__class__.__name__}({self.format},{self.dtype})" + + +class HStack(Stack): + def perform(self, node, block, outputs): (out,) = outputs - assert _is_sparse(x) and _is_sparse(y) - assert x.shape == y.shape - out[0] = x + y + for b in block: + assert _is_sparse(b) + out[0] = scipy.sparse.hstack(block, format=self.format, dtype=self.dtype) + # Some version of scipy (at least 0.14.0.dev-c4314b0) + # Do not cast to the wanted dtype. + if out[0].dtype != self.dtype: + out[0] = out[0].astype(self.dtype) def grad(self, inputs, gout): - (x, y) = inputs (gz,) = gout - assert _is_sparse_variable(x) and _is_sparse_variable(y) - assert _is_sparse_variable(gz) - return gz, gz + is_continuous = [ + (inputs[i].dtype in tensor_continuous_dtypes) for i in range(len(inputs)) + ] - def infer_shape(self, fgraph, node, shapes): - return [shapes[0]] + if _is_sparse_variable(gz): + gz = dense_from_sparse(gz) + + split = Split(len(inputs))(gz, 1, ptb.stack([x.shape[1] for x in inputs])) + if not isinstance(split, list): + split = [split] + derivative = [SparseFromDense(self.format)(s) for s in split] -add_s_s = AddSS() + def choose(continuous, derivative): + if continuous: + return derivative + else: + return None + return [choose(c, d) for c, d in zip(is_continuous, derivative, strict=True)] -class AddSSData(Op): - """Add two sparse matrices assuming they have the same sparsity pattern. + def infer_shape(self, fgraph, node, ins_shapes): + d = sum(shape[1] for shape in ins_shapes) + return [(ins_shapes[0][0], d)] - Notes - ----- - The grad implemented is structured. +def hstack(blocks, format=None, dtype=None): """ + Stack sparse matrices horizontally (column wise). - __props__ = () + This wrap the method hstack from scipy. - def make_node(self, x, y): - """ + Parameters + ---------- + blocks + List of sparse array of compatible shape. + format + String representing the output format. Default is csc. + dtype + Output dtype. - Parameters - ---------- - x - Sparse matrix. - y - Sparse matrix. + Returns + ------- + array + The concatenation of the sparse array column wise. - Notes - ----- - `x` and `y` are assumed to have the same sparsity pattern. + Notes + ----- + The number of line of the sparse matrix must agree. - """ - x, y = map(as_sparse_variable, [x, y]) - assert x.format in ("csr", "csc") - assert y.format in ("csr", "csc") - if x.type.dtype != y.type.dtype: - raise NotImplementedError() - if x.type.format != y.type.format: - raise NotImplementedError() - return Apply( - self, [x, y], [SparseTensorType(dtype=x.type.dtype, format=x.type.format)()] - ) + The grad implemented is regular, i.e. not structured. - def perform(self, node, inputs, outputs): - (x, y) = inputs + """ + + blocks = [as_sparse_variable(i) for i in blocks] + if dtype is None: + dtype = ps.upcast(*[i.dtype for i in blocks]) + return HStack(format=format, dtype=dtype)(*blocks) + + +class VStack(Stack): + def perform(self, node, block, outputs): (out,) = outputs - assert _is_sparse(x) and _is_sparse(y) - assert x.shape == y.shape - assert x.data.shape == y.data.shape - out[0] = x.copy() - out[0].data += y.data + for b in block: + assert _is_sparse(b) + out[0] = scipy.sparse.vstack(block, format=self.format, dtype=self.dtype) + # Some version of scipy (at least 0.14.0.dev-c4314b0) + # Do not cast to the wanted dtype. + if out[0].dtype != self.dtype: + out[0] = out[0].astype(self.dtype) def grad(self, inputs, gout): (gz,) = gout - is_continuous = [(i.dtype in continuous_dtypes) for i in inputs] - derivative = {True: gz, False: None} - return [derivative[b] for b in is_continuous] - - def infer_shape(self, fgraph, node, ins_shapes): - return [ins_shapes[0]] + is_continuous = [ + (inputs[i].dtype in tensor_continuous_dtypes) for i in range(len(inputs)) + ] + if _is_sparse_variable(gz): + gz = dense_from_sparse(gz) -add_s_s_data = AddSSData() + split = Split(len(inputs))(gz, 0, ptb.stack([x.shape[0] for x in inputs])) + if not isinstance(split, list): + split = [split] + derivative = [SparseFromDense(self.format)(s) for s in split] -class AddSD(Op): - # add(sparse, sparse). - # see the doc of add() for more detail. - __props__ = () + def choose(continuous, derivative): + if continuous: + return derivative + else: + return None - def make_node(self, x, y): - x, y = as_sparse_variable(x), ptb.as_tensor_variable(y) - assert x.format in ("csr", "csc") - out_dtype = ps.upcast(x.type.dtype, y.type.dtype) + return [choose(c, d) for c, d in zip(is_continuous, derivative, strict=True)] - # The magic number two here arises because L{scipy.sparse} - # objects must be matrices (have dimension 2) - assert y.type.ndim == 2 - return Apply( - self, - [x, y], - [TensorType(dtype=out_dtype, shape=y.type.shape)()], - ) + def infer_shape(self, fgraph, node, ins_shapes): + d = sum(shape[0] for shape in ins_shapes) + return [(d, ins_shapes[0][1])] - def perform(self, node, inputs, outputs): - (x, y) = inputs - (out,) = outputs - assert _is_dense(y) - # The asarray is needed as in some case, this return a - # numpy.matrixlib.defmatrix.matrix object and not an ndarray. - out[0] = np.asarray(x + y, dtype=node.outputs[0].type.dtype) +def vstack(blocks, format=None, dtype=None): + """ + Stack sparse matrices vertically (row wise). - def grad(self, inputs, gout): - (x, y) = inputs - (gz,) = gout - assert _is_sparse_variable(x) and _is_dense_variable(y) - assert _is_dense_variable(gz) - return sp_ones_like(x) * gz, gz - - def infer_shape(self, fgraph, node, shapes): - return [shapes[1]] - - -add_s_d = AddSD() - - -class StructuredAddSV(Op): - """Structured addition of a sparse matrix and a dense vector. - - The elements of the vector are only added to the corresponding - non-zero elements of the sparse matrix. Therefore, this operation - outputs another sparse matrix. - - Notes - ----- - The grad implemented is structured since the op is structured. - - """ - - __props__ = () - - def make_node(self, x, y): - """ - Parameters - ---------- - x - Sparse matrix. - y - Tensor type vector. - - """ - x = as_sparse_variable(x) - assert x.format in ("csr", "csc") - y = ptb.as_tensor_variable(y) - - assert y.type.ndim == 1 - - if x.type.dtype != y.type.dtype: - raise NotImplementedError() - return Apply( - self, [x, y], [SparseTensorType(dtype=x.type.dtype, format=x.type.format)()] - ) - - def perform(self, node, inputs, outputs): - (x, y) = inputs - (out,) = outputs - assert _is_sparse(x) and not _is_sparse(y) - assert x.shape[1] == y.shape[0] - out[0] = x.__class__(x + (x.toarray() != 0) * y) - - def grad(self, inputs, gout): - (x, y) = inputs - (gz,) = gout - assert _is_sparse_variable(x) and not _is_sparse_variable(y) - assert _is_sparse_variable(gz) - return gz, sp_sum(gz, axis=0, sparse_grad=True) - - def infer_shape(self, fgraph, node, ins_shapes): - return [ins_shapes[0]] - - -structured_add_s_v = StructuredAddSV() - - -def add(x, y): - """ - Add two matrices, at least one of which is sparse. - - This method will provide the right op according - to the inputs. - - Parameters - ---------- - x - A matrix variable. - y - A matrix variable. - - Returns - ------- - A sparse matrix - `x` + `y` - - Notes - ----- - At least one of `x` and `y` must be a sparse matrix. - - The grad will be structured only when one of the variable will be a dense - matrix. - - """ - - if hasattr(x, "getnnz"): - x = as_sparse_variable(x) - if hasattr(y, "getnnz"): - y = as_sparse_variable(y) - if not isinstance(x, Variable): - x = ptb.as_tensor_variable(x) - if not isinstance(y, Variable): - y = ptb.as_tensor_variable(y) - - x_is_sparse_variable = _is_sparse_variable(x) - y_is_sparse_variable = _is_sparse_variable(y) - - assert x_is_sparse_variable or y_is_sparse_variable - if x_is_sparse_variable and y_is_sparse_variable: - return add_s_s(x, y) - elif x_is_sparse_variable and not y_is_sparse_variable: - return add_s_d(x, y) - elif y_is_sparse_variable and not x_is_sparse_variable: - return add_s_d(y, x) - else: - raise NotImplementedError() - - -def subtract( - x: SparseVariable | TensorVariable, y: SparseVariable | TensorVariable -) -> SparseVariable: - """ - Subtract two matrices, at least one of which is sparse. - - This method will provide the right op according to the inputs. - - Parameters - ---------- - x : SparseVariable or TensorVariable - A matrix variable. - y : SparseVariable or TensorVariable - A matrix variable. - - Returns - ------- - result: SparseVariable - Result of `x - y`, as a sparse matrix. - - Notes - ----- - At least one of `x` and `y` must be a sparse matrix. - - The grad will be structured only when one of the variable will be a dense matrix. - """ - return x + (-y) - - -def sub(x, y): - warn( - "pytensor.sparse.sub is deprecated and will be removed in a future version. Use " - "pytensor.sparse.subtract instead.", - category=DeprecationWarning, - stacklevel=2, - ) - - return subtract(x, y) - - -sub.__doc__ = subtract.__doc__ - - -class MulSS(Op): - # mul(sparse, sparse) - # See the doc of mul() for more detail - __props__ = () - - def make_node(self, x, y): - x, y = as_sparse_variable(x), as_sparse_variable(y) - assert x.format in ("csr", "csc") - assert y.format in ("csr", "csc") - out_dtype = ps.upcast(x.type.dtype, y.type.dtype) - return Apply( - self, [x, y], [SparseTensorType(dtype=out_dtype, format=x.type.format)()] - ) - - def perform(self, node, inputs, outputs): - (x, y) = inputs - (out,) = outputs - assert _is_sparse(x) and _is_sparse(y) - assert len(x.shape) == 2 - assert y.shape == x.shape - # This calls the element-wise multiple - # x * y calls dot... - out[0] = x.multiply(y) - - def grad(self, inputs, gout): - (x, y) = inputs - (gz,) = gout - return y * gz, x * gz - - def infer_shape(self, fgraph, node, shapes): - return [shapes[0]] - - -mul_s_s = MulSS() - - -class MulSD(Op): - # mul(sparse, dense) - # See the doc of mul() for more detail - __props__ = () - - def make_node(self, x, y): - x, y = as_sparse_variable(x), ptb.as_tensor_variable(y) - - assert x.format in ("csr", "csc") - - # upcast the tensor. Is the cast of sparse done implemented? - dtype = ps.upcast(x.type.dtype, y.type.dtype) - - # The magic number two here arises because L{scipy.sparse} - # objects must be matrices (have dimension 2) - # Broadcasting of the sparse matrix is not supported. - # We support nd == 0 used by grad of SpSum() - assert y.type.ndim in (0, 2) - out = SparseTensorType(dtype=dtype, format=x.type.format)() - return Apply(self, [x, y], [out]) - - def perform(self, node, inputs, outputs): - (x, y) = inputs - (out,) = outputs - assert _is_sparse(x) and _is_dense(y) - if len(y.shape) == 0: - out_dtype = node.outputs[0].dtype - if x.dtype == out_dtype: - z = x.copy() - else: - z = x.astype(out_dtype) - out[0] = z - out[0].data *= y - elif len(y.shape) == 1: - raise NotImplementedError() # RowScale / ColScale - elif len(y.shape) == 2: - # if we have enough memory to fit y, maybe we can fit x.asarray() - # too? - # TODO: change runtime from O(M*N) to O(nonzeros) - M, N = x.shape - assert x.shape == y.shape - out_dtype = node.outputs[0].dtype - - if x.format == "csc": - indices = x.indices - indptr = x.indptr - if x.dtype == out_dtype: - z = x.copy() - else: - z = x.astype(out_dtype) - z_data = z.data - - for j in range(0, N): - for i_idx in range(indptr[j], indptr[j + 1]): - i = indices[i_idx] - z_data[i_idx] *= y[i, j] - out[0] = z - elif x.format == "csr": - indices = x.indices - indptr = x.indptr - if x.dtype == out_dtype: - z = x.copy() - else: - z = x.astype(out_dtype) - z_data = z.data - - for i in range(0, M): - for j_idx in range(indptr[i], indptr[i + 1]): - j = indices[j_idx] - z_data[j_idx] *= y[i, j] - out[0] = z - else: - warn( - "This implementation of MulSD is deficient: {x.format}", - ) - out[0] = type(x)(x.toarray() * y) - - def grad(self, inputs, gout): - (x, y) = inputs - (gz,) = gout - assert _is_sparse_variable(x) and _is_dense_variable(y) - assert _is_sparse_variable(gz) - return y * gz, dense_from_sparse(x * gz) - - def infer_shape(self, fgraph, node, shapes): - return [shapes[0]] - - -mul_s_d = MulSD() - - -class MulSV(Op): - """Element-wise multiplication of sparse matrix by a broadcasted dense vector element wise. - - Notes - ----- - The grad implemented is regular, i.e. not structured. - - """ - - __props__ = () - - def make_node(self, x, y): - """ - Parameters - ---------- - x - Sparse matrix to multiply. - y - Tensor broadcastable vector. - - """ - x = as_sparse_variable(x) - assert x.format in ("csr", "csc") - y = ptb.as_tensor_variable(y) - - assert y.type.ndim == 1 - - if x.type.dtype != y.type.dtype: - raise NotImplementedError( - "MulSV not implemented for differing dtypes." - f"Got {x.type.dtype} and {y.type.dtype}." - ) - return Apply( - self, [x, y], [SparseTensorType(dtype=x.type.dtype, format=x.type.format)()] - ) - - def perform(self, node, inputs, outputs): - (x, y) = inputs - (out,) = outputs - assert _is_sparse(x) and not _is_sparse(y) - assert x.shape[1] == y.shape[0] - out[0] = x.__class__(x.toarray() * y) - - def grad(self, inputs, gout): - (x, y) = inputs - (gz,) = gout - assert _is_sparse_variable(x) and _is_dense_variable(y) - assert _is_sparse_variable(gz) - - # mul_s_v is not implemented if the types vary - - if gz.dtype == "float64" and y.dtype == "float32": - y = y.astype("float64") - - if gz.dtype == "float32" and y.dtype == "float64": - gz = gz.astype("float64") - - return mul_s_v(gz, y), sp_sum(x * gz, axis=0, sparse_grad=True) - - def infer_shape(self, fgraph, node, ins_shapes): - return [ins_shapes[0]] - - -mul_s_v = MulSV() - - -def multiply( - x: SparseTensorType | TensorType, y: SparseTensorType | TensorType -) -> SparseVariable: - """ - Multiply elementwise two matrices, at least one of which is sparse. - - This method will provide the right op according to the inputs. - - Parameters - ---------- - x : SparseVariable - A matrix variable. - y : SparseVariable - A matrix variable. - - Returns - ------- - result: SparseVariable - The elementwise multiplication of `x` and `y`. - - Notes - ----- - At least one of `x` and `y` must be a sparse matrix. - - The gradient is regular, i.e. not structured. - """ - - x = as_sparse_or_tensor_variable(x) - y = as_sparse_or_tensor_variable(y) - - x_is_sparse_variable = _is_sparse_variable(x) - y_is_sparse_variable = _is_sparse_variable(y) - - assert x_is_sparse_variable or y_is_sparse_variable - if x_is_sparse_variable and y_is_sparse_variable: - # mul_s_s is not implemented if the types differ - if y.dtype == "float64" and x.dtype == "float32": - x = x.astype("float64") - - return mul_s_s(x, y) - elif x_is_sparse_variable and not y_is_sparse_variable: - # mul is unimplemented if the dtypes differ - if y.dtype == "float64" and x.dtype == "float32": - x = x.astype("float64") - - return mul_s_d(x, y) - elif y_is_sparse_variable and not x_is_sparse_variable: - return mul_s_d(y, x) - else: - raise NotImplementedError() - - -def mul(x, y): - warn( - "pytensor.sparse.mul is deprecated and will be removed in a future version. Use " - "pytensor.sparse.multiply instead.", - category=DeprecationWarning, - stacklevel=2, - ) - - return multiply(x, y) - - -mul.__doc__ = multiply.__doc__ - - -class __ComparisonOpSS(Op): - """ - Used as a superclass for all comparisons between two sparses matrices. - - Parameters - ---------- - x - First compared sparse matrix. - y - Second compared sparse matrix - - Returns - ------- - object - Comparison(x,y) - - """ - - __props__ = () - - # Function to override - def comparison(self, x, y): - raise NotImplementedError() - - def make_node(self, x, y): - x = as_sparse_variable(x) - y = as_sparse_variable(y) - - if x.type.format != y.type.format: - raise NotImplementedError() - return Apply( - self, [x, y], [SparseTensorType(dtype="uint8", format=x.type.format)()] - ) - - def perform(self, node, inputs, outputs): - (x, y) = inputs - (out,) = outputs - assert _is_sparse(x) and _is_sparse(y) - assert x.shape == y.shape - out[0] = self.comparison(x, y).astype("uint8") - - def infer_shape(self, fgraph, node, ins_shapes): - return [ins_shapes[0]] - - -class __ComparisonOpSD(Op): - """ - Used as a superclass for all comparisons between sparse and dense matrix. - - Parameters - ---------- - x - Sparse matrix. - y - Dense matrix. - - Returns - ------- - object - Comparison(x,y) - - """ - - __props__ = () - - # Function to override - def comparison(self, x, y): - raise NotImplementedError() - - def make_node(self, x, y): - x, y = as_sparse_variable(x), ptb.as_tensor_variable(y) - - assert y.type.ndim == 2 - out = TensorType(dtype="uint8", shape=(None, None))() - return Apply(self, [x, y], [out]) - - def perform(self, node, inputs, outputs): - (x, y) = inputs - (out,) = outputs - assert _is_sparse(x) - assert x.shape == y.shape - assert _is_dense(y) - o = self.comparison(x, y).astype("uint8") - o = np.asarray(o) - out[0] = o - - def infer_shape(self, fgraph, node, ins_shapes): - return [ins_shapes[0]] - - -def __ComparisonSwitch(SS, SD, DS): - """ - - Parameters - ---------- - SS - Function to apply between two sparses matrices. - SD - Function to apply between a sparse and a dense matrix. - DS - Function to apply between a dense and a sparse matrix. - - Returns - ------- - function - Switch function taking two matrices as input. - - Notes - ----- - At least one of `x` and `y` must be a sparse matrix. - - DS swap input as a dense matrix cannot be a left operand. - - """ - - def helper(x, y): - scipy_ver = [int(n) for n in scipy.__version__.split(".")[:2]] - - assert scipy_ver >= [0, 13] - - if hasattr(x, "getnnz"): - x = as_sparse_variable(x) - if hasattr(y, "getnnz"): - y = as_sparse_variable(y) - if not isinstance(x, Variable): - x = ptb.as_tensor_variable(x) - if not isinstance(y, Variable): - y = ptb.as_tensor_variable(y) - - x_is_sparse_variable = _is_sparse_variable(x) - y_is_sparse_variable = _is_sparse_variable(y) - - assert x_is_sparse_variable or y_is_sparse_variable - if x_is_sparse_variable and y_is_sparse_variable: - return SS(x, y) - elif x_is_sparse_variable and not y_is_sparse_variable: - return SD(x, y) - elif y_is_sparse_variable and not x_is_sparse_variable: - return DS(y, x) - else: - raise NotImplementedError() - - return helper - - -class EqualSS(__ComparisonOpSS): - def comparison(self, x, y): - return x == y - - -equal_s_s = EqualSS() - - -class EqualSD(__ComparisonOpSD): - def comparison(self, x, y): - return x == y - - -equal_s_d = EqualSD() - - -class NotEqualSS(__ComparisonOpSS): - def comparison(self, x, y): - return x != y - - -not_equal_s_s = NotEqualSS() - - -class NotEqualSD(__ComparisonOpSD): - def comparison(self, x, y): - return x != y - - -not_equal_s_d = NotEqualSD() - - -class LessThanSS(__ComparisonOpSS): - def comparison(self, x, y): - return x < y - - -less_than_s_s = LessThanSS() - - -class LessThanSD(__ComparisonOpSD): - def comparison(self, x, y): - return x < y - - -less_than_s_d = LessThanSD() - - -class GreaterThanSS(__ComparisonOpSS): - def comparison(self, x, y): - return x > y - - -greater_than_s_s = GreaterThanSS() - - -class GreaterThanSD(__ComparisonOpSD): - def comparison(self, x, y): - return x > y - - -greater_than_s_d = GreaterThanSD() - - -class LessEqualSS(__ComparisonOpSS): - def comparison(self, x, y): - return x <= y - - -less_equal_s_s = LessEqualSS() - - -class LessEqualSD(__ComparisonOpSD): - def comparison(self, x, y): - return x <= y - - -less_equal_s_d = LessEqualSD() - - -class GreaterEqualSS(__ComparisonOpSS): - def comparison(self, x, y): - return x >= y - - -greater_equal_s_s = GreaterEqualSS() - - -class GreaterEqualSD(__ComparisonOpSD): - def comparison(self, x, y): - return x >= y - - -greater_equal_s_d = GreaterEqualSD() - -eq = __ComparisonSwitch(equal_s_s, equal_s_d, equal_s_d) - -neq = __ComparisonSwitch(not_equal_s_s, not_equal_s_d, not_equal_s_d) - -lt = __ComparisonSwitch(less_than_s_s, less_than_s_d, greater_than_s_d) - -gt = __ComparisonSwitch(greater_than_s_s, greater_than_s_d, less_than_s_d) - -le = __ComparisonSwitch(less_equal_s_s, less_equal_s_d, greater_equal_s_d) - -ge = __ComparisonSwitch(greater_equal_s_s, greater_equal_s_d, less_equal_s_d) - - -class Stack(Op): - __props__ = ("format", "dtype") - - def __init__(self, format=None, dtype=None): - if format is None: - self.format = "csc" - else: - self.format = format - - if dtype is None: - raise ValueError("The output dtype must be specified.") - self.dtype = dtype - - def make_node(self, *mat): - if not mat: - raise ValueError("Cannot join an empty list of sparses.") - var = [as_sparse_variable(x) for x in mat] - - for x in var: - assert x.format in ("csr", "csc") - - return Apply( - self, var, [SparseTensorType(dtype=self.dtype, format=self.format)()] - ) - - def __str__(self): - return f"{self.__class__.__name__}({self.format},{self.dtype})" - - -class HStack(Stack): - def perform(self, node, block, outputs): - (out,) = outputs - for b in block: - assert _is_sparse(b) - out[0] = scipy.sparse.hstack(block, format=self.format, dtype=self.dtype) - # Some version of scipy (at least 0.14.0.dev-c4314b0) - # Do not cast to the wanted dtype. - if out[0].dtype != self.dtype: - out[0] = out[0].astype(self.dtype) - - def grad(self, inputs, gout): - (gz,) = gout - is_continuous = [ - (inputs[i].dtype in tensor_continuous_dtypes) for i in range(len(inputs)) - ] - - if _is_sparse_variable(gz): - gz = dense_from_sparse(gz) - - split = Split(len(inputs))(gz, 1, ptb.stack([x.shape[1] for x in inputs])) - if not isinstance(split, list): - split = [split] - - derivative = [SparseFromDense(self.format)(s) for s in split] - - def choose(continuous, derivative): - if continuous: - return derivative - else: - return None - - return [choose(c, d) for c, d in zip(is_continuous, derivative, strict=True)] - - def infer_shape(self, fgraph, node, ins_shapes): - d = sum(shape[1] for shape in ins_shapes) - return [(ins_shapes[0][0], d)] - - -def hstack(blocks, format=None, dtype=None): - """ - Stack sparse matrices horizontally (column wise). - - This wrap the method hstack from scipy. - - Parameters - ---------- - blocks - List of sparse array of compatible shape. - format - String representing the output format. Default is csc. - dtype - Output dtype. - - Returns - ------- - array - The concatenation of the sparse array column wise. - - Notes - ----- - The number of line of the sparse matrix must agree. - - The grad implemented is regular, i.e. not structured. - - """ - - blocks = [as_sparse_variable(i) for i in blocks] - if dtype is None: - dtype = ps.upcast(*[i.dtype for i in blocks]) - return HStack(format=format, dtype=dtype)(*blocks) - - -class VStack(Stack): - def perform(self, node, block, outputs): - (out,) = outputs - for b in block: - assert _is_sparse(b) - out[0] = scipy.sparse.vstack(block, format=self.format, dtype=self.dtype) - # Some version of scipy (at least 0.14.0.dev-c4314b0) - # Do not cast to the wanted dtype. - if out[0].dtype != self.dtype: - out[0] = out[0].astype(self.dtype) - - def grad(self, inputs, gout): - (gz,) = gout - is_continuous = [ - (inputs[i].dtype in tensor_continuous_dtypes) for i in range(len(inputs)) - ] - - if _is_sparse_variable(gz): - gz = dense_from_sparse(gz) - - split = Split(len(inputs))(gz, 0, ptb.stack([x.shape[0] for x in inputs])) - if not isinstance(split, list): - split = [split] - - derivative = [SparseFromDense(self.format)(s) for s in split] - - def choose(continuous, derivative): - if continuous: - return derivative - else: - return None - - return [choose(c, d) for c, d in zip(is_continuous, derivative, strict=True)] - - def infer_shape(self, fgraph, node, ins_shapes): - d = sum(shape[0] for shape in ins_shapes) - return [(d, ins_shapes[0][1])] - - -def vstack(blocks, format=None, dtype=None): - """ - Stack sparse matrices vertically (row wise). - - This wrap the method vstack from scipy. - - Parameters - ---------- - blocks - List of sparse array of compatible shape. - format - String representing the output format. Default is csc. - dtype - Output dtype. - - Returns - ------- - array - The concatenation of the sparse array row wise. - - Notes - ----- - The number of column of the sparse matrix must agree. - - The grad implemented is regular, i.e. not structured. - - """ - - blocks = [as_sparse_variable(i) for i in blocks] - if dtype is None: - dtype = ps.upcast(*[i.dtype for i in blocks]) - return VStack(format=format, dtype=dtype)(*blocks) - - -class Remove0(Op): - """Remove explicit zeros from a sparse matrix. - - Notes - ----- - The grad implemented is regular, i.e. not structured. - - """ - - __props__ = ("inplace",) - - def __init__(self, inplace=False): - self.inplace = inplace - if self.inplace: - self.destroy_map = {0: [0]} - - def __str__(self): - l = [] - if self.inplace: - l.append("inplace") - return f"{self.__class__.__name__}{{{', '.join(l)}}}" - - def make_node(self, x): - """ - - Parameters - ---------- - x - Sparse matrix. - - """ - x = as_sparse_variable(x) - assert x.format in ("csr", "csc") - return Apply(self, [x], [x.type()]) - - def perform(self, node, inputs, outputs): - (x,) = inputs - (z,) = outputs - if self.inplace: - c = x - else: - c = x.copy() - c.eliminate_zeros() - z[0] = c - - def grad(self, inputs, gout): - (_x,) = inputs - (gz,) = gout - return [gz] - - def infer_shape(self, fgraph, node, i0_shapes): - return i0_shapes - - -remove0 = Remove0() - - -def structured_monoid(tensor_op): - # Generic operation to perform many kinds of monoid element-wise - # operations on the non-zeros of a sparse matrix. - - # The first parameter must always be a sparse matrix. The other parameters - # must be scalars which will be passed as argument to the tensor_op. - - def decorator(f): - def wrapper(*args): - x = as_sparse_variable(args[0]) - assert x.format in ("csr", "csc") - - xs = [ps.as_scalar(arg) for arg in args[1:]] - - data, ind, ptr, _shape = csm_properties(x) - - data = tensor_op(data, *xs) - - return CSM(x.format)(data, ind, ptr, _shape) - - wrapper.__name__ = str(tensor_op.scalar_op) - return wrapper - - return decorator - - -@structured_monoid(sigmoid) -def structured_sigmoid(x): - """ - Structured elemwise sigmoid. - - """ - - -@structured_monoid(exp) -def structured_exp(x): - """ - Structured elemwise exponential. - - """ - - -@structured_monoid(log) -def structured_log(x): - """ - Structured elemwise logarithm. - - """ - - -@structured_monoid(pt_pow) -def structured_pow(x, y): - """ - Structured elemwise power of sparse matrix x by scalar y. - - """ - - -@structured_monoid(minimum) -def structured_minimum(x, y): - """ - Structured elemwise minimum of sparse matrix x by scalar y. - - """ - - -@structured_monoid(maximum) -def structured_maximum(x, y): - """ - Structured elemwise maximum of sparse matrix x by scalar y. - - """ - - -@structured_monoid(pt_add) -def structured_add(x): - """ - Structured addition of sparse matrix x and scalar y. - - """ - - -@structured_monoid(sin) # type: ignore[no-redef] -def sin(x): - """ - Elemwise sinus of `x`. - - """ - - -@structured_monoid(tan) # type: ignore[no-redef] -def tan(x): - """ - Elemwise tan of `x`. - - """ - - -@structured_monoid(arcsin) # type: ignore[no-redef] -def arcsin(x): - """ - Elemwise arcsinus of `x`. - - """ - - -@structured_monoid(arctan) # type: ignore[no-redef] -def arctan(x): - """ - Elemwise arctan of `x`. - - """ - - -@structured_monoid(sinh) # type: ignore[no-redef] -def sinh(x): - """ - Elemwise sinh of `x`. - - """ - - -@structured_monoid(arcsinh) # type: ignore[no-redef] -def arcsinh(x): - """ - Elemwise arcsinh of `x`. - - """ - - -@structured_monoid(tanh) # type: ignore[no-redef] -def tanh(x): - """ - Elemwise tanh of `x`. - - """ - - -@structured_monoid(arctanh) # type: ignore[no-redef] -def arctanh(x): - """ - Elemwise arctanh of `x`. - - """ - - -@structured_monoid(round_half_to_even) -def rint(x): - """ - Elemwise round half to even of `x`. - - """ - - -# Give it a simple name instead of the complex one that would automatically -# be derived from `round_half_to_even`. -rint.__name__ = "rint" - - -@structured_monoid(sign) # type: ignore[no-redef] -def sign(x): - """ - Elemwise signe of `x`. - - """ - - -@structured_monoid(ceil) # type: ignore[no-redef] -def ceil(x): - """ - Elemwise ceiling of `x`. - - """ - - -@structured_monoid(floor) # type: ignore[no-redef] -def floor(x): - """ - Elemwise floor of `x`. - - """ - - -@structured_monoid(log1p) # type: ignore[no-redef] -def log1p(x): - """ - Elemwise log(1 + `x`). - - """ - - -@structured_monoid(expm1) # type: ignore[no-redef] -def expm1(x): - """ - Elemwise e^`x` - 1. - - """ - - -@structured_monoid(deg2rad) # type: ignore[no-redef] -def deg2rad(x): - """ - Elemwise degree to radian. - - """ - - -@structured_monoid(rad2deg) # type: ignore[no-redef] -def rad2deg(x): - """ - Elemwise radian to degree. - - """ - - -@structured_monoid(trunc) # type: ignore[no-redef] -def trunc(x): - """ - Elemwise truncation. - - """ - - -@structured_monoid(sqr) # type: ignore[no-redef] -def sqr(x): - """ - Elemwise `x` * `x`. - - """ - - -@structured_monoid(sqrt) # type: ignore[no-redef] -def sqrt(x): - """ - Elemwise square root of `x`. - - """ - - -@structured_monoid(_conj) # type: ignore[no-redef] -def _conj(x): - """ - Elemwise complex conjugate of `x`. - - """ - - -def conjugate(x): - _x = as_sparse_variable(x) - if _x.type.dtype not in complex_dtypes: - return _x - return _conj(_x) - - -conj = conjugate - - -class TrueDot(Op): - # TODO - # Simplify code by splitting into DotSS and DotSD. - - __props__ = () - - # The grad_preserves_dense attribute doesn't change the - # execution behavior. To let the optimizer merge nodes with - # different values of this attribute we shouldn't compare it - # here. - - def __init__(self, grad_preserves_dense=True): - self.grad_preserves_dense = grad_preserves_dense - - def make_node(self, x, y): - # NOTE - # Because of trickiness of implementing, - # we assume that the left argument x is a - # SparseVariable (not dense) - - if x.type.dtype != y.type.dtype: - raise NotImplementedError() - - if not _is_sparse_variable(x): - raise TypeError(x) - - # These are the conversions performed by scipy.sparse.dot - if x.type.format == "csc" or x.type.format == "coo": - myformat = "csc" - elif x.type.format == "csr": - myformat = "csr" - else: - raise NotImplementedError() - - inputs = [x, y] # Need to convert? e.g. assparse - outputs = [SparseTensorType(dtype=x.type.dtype, format=myformat)()] - return Apply(self, inputs, outputs) - - def perform(self, node, inp, out_): - # TODO - # -Verify that output is sufficiently sparse, - # and raise a warning if it is not. - # -Also determine that we are storing the - # output in the best storage format? - - x, y = inp - (out,) = out_ - rval = x.dot(y) - if not scipy.sparse.issparse(rval): - rval = getattr(scipy.sparse, x.format + "_matrix")(rval) - # x.dot call tocsr() that will "upcast" to ['int8', 'uint8', 'short', - # 'ushort', 'intc', 'uintc', 'longlong', 'ulonglong', 'single', - # 'double', 'longdouble', 'csingle', 'cdouble', 'clongdouble'] - # But ulonglong is uint64 on x86-64, but with a different typenum! - if rval.dtype.num != np.dtype(str(rval.dtype)).num: - assert str(rval.dtype) == node.outputs[0].dtype - # Create a view with the expected typenum. - format = node.outputs[0].type.format - data = rval.data.view(dtype=node.outputs[0].dtype) - indices = rval.indices - indptr = rval.indptr - _shape = rval.shape - # No need to copy indices and indptr as in CSM.perform(), - # as there is only one user of them. - if format == "csc": - rval = scipy.sparse.csc_matrix( - (data, indices, indptr), _shape, copy=False - ) - else: - assert format == "csr" - rval = scipy.sparse.csr_matrix( - (data, indices, indptr), _shape, copy=False - ) - out[0] = rval - - def grad(self, inputs, gout): - (x, y) = inputs - (gz,) = gout - assert _is_sparse_variable(gz) - assert _is_sparse_variable(x) - - rval = [true_dot(gz, y.T), true_dot(x.T, gz)] - if _is_dense_variable(y): - if self.grad_preserves_dense: - rval[1] = dense_from_sparse(rval[1]) - return rval - - def infer_shape(self, fgraph, node, shapes): - return [(shapes[0][0], shapes[1][1])] - - -def true_dot(x, y, grad_preserves_dense=True): - """ - Operation for efficiently calculating the dot product when - one or all operands are sparse. Supported formats are CSC and CSR. - The output of the operation is sparse. - - Parameters - ---------- - x - Sparse matrix. - y - Sparse matrix or 2d tensor variable. - grad_preserves_dense : bool - If True (default), makes the grad of dense inputs dense. - Otherwise the grad is always sparse. - - Returns - ------- - The dot product `x`.`y` in a sparse format. - - Notex - ----- - The grad implemented is regular, i.e. not structured. - - """ - # TODO - # Maybe the triple-transposition formulation - # (when x is dense) is slow. See if there is a - # direct way to do this. - - if hasattr(x, "getnnz"): - x = as_sparse_variable(x) - assert x.format in ("csr", "csc") - if hasattr(y, "getnnz"): - y = as_sparse_variable(y) - assert y.format in ("csr", "csc") - - x_is_sparse_variable = _is_sparse_variable(x) - y_is_sparse_variable = _is_sparse_variable(y) - - if not x_is_sparse_variable and not y_is_sparse_variable: - raise TypeError() - if x_is_sparse_variable: - return TrueDot(grad_preserves_dense)(x, y) - else: - assert y_is_sparse_variable - return transpose(TrueDot(grad_preserves_dense)(y.T, x.T)) - - -class StructuredDot(Op): - __props__ = () - - def make_node(self, a, b): - a = as_sparse_variable(a) - assert a.format in ("csr", "csc", "bsr") - - if not _is_sparse_variable(a): - raise TypeError( - "First argument must be of type SparseVariable or SparseConstant" - ) - dtype_out = ps.upcast(a.type.dtype, b.type.dtype) - if b.type.ndim != 2: - raise NotImplementedError("non-matrix b") - - if _is_sparse_variable(b): - return Apply(self, [a, b], [SparseTensorType(a.type.format, dtype_out)()]) - else: - return Apply( - self, - [a, b], - [ - tensor( - dtype=dtype_out, - shape=(None, 1 if b.type.shape[1] == 1 else None), - ) - ], - ) - - def perform(self, node, inputs, outputs): - (a, b) = inputs - (out,) = outputs - if a.shape[1] != b.shape[0]: - raise ValueError( - "shape mismatch in StructuredDot.perform", (a.shape, b.shape) - ) - - variable = a * b - if isinstance(node.outputs[0].type, SparseTensorType): - assert _is_sparse(variable) - out[0] = variable - return - - assert _is_dense(variable) # scipy 0.7 automatically converts to dense - - # dot of an NxM sparse matrix, with a Mx1 dense matrix, returns vector - # not matrix - if variable.ndim == 1: - variable = np.expand_dims(variable, 1) - elif variable.ndim != 2: - raise Exception("Output of structured dot should be a matrix (ndim=2)") - - assert variable.ndim == 2 - - if variable.shape != (a.shape[0], b.shape[1]): - if b.shape[0] == 1: - raise Exception( - f"a.shape={a.shape}, b.shape={b.shape}, " - f"variable.shape={variable.shape}? This is probably " - "because scipy.csc_matrix.dot has a bug " - "with singleton dimensions (i.e. " - "b.shape[0]=1) in SciPy 0.6. Use SciPy " - f"0.7. (You have SciPy version {scipy.__version__}.)" - ) - else: - raise Exception( - f"a.shape={a.shape}, b.shape={b.shape}, variable.shape={variable.shape}?" - ) - - # The cast is needed as otherwise we hit the bug mentioned into - # _asarray function documentation. - out[0] = np.asarray(variable, str(variable.dtype)) - - def grad(self, inputs, gout): - # a is sparse, b is dense, g_out is dense - # ga = g_out x b.T - # gb = a.T x g_out - (a, b) = inputs - (g_out,) = gout - return [structured_dot_grad(a, b, g_out), structured_dot(a.T, g_out)] - - def infer_shape(self, fgraph, node, shapes): - return [(shapes[0][0], shapes[1][1])] - - -_structured_dot = StructuredDot() - - -def structured_dot(x, y): - """ - Structured Dot is like dot, except that only the - gradient wrt non-zero elements of the sparse matrix - `a` are calculated and propagated. - - The output is presumed to be a dense matrix, and is represented by a - TensorType instance. + This wrap the method vstack from scipy. Parameters ---------- - a - A sparse matrix. - b - A sparse or dense matrix. + blocks + List of sparse array of compatible shape. + format + String representing the output format. Default is csc. + dtype + Output dtype. Returns ------- - A sparse matrix - The dot product of `a` and `b`. - - Notes - ----- - The grad implemented is structured. - - """ - - # @todo: Maybe the triple-transposition formulation (when x is dense) - # is slow. See if there is a direct way to do this. - # (JB 20090528: Transposing tensors and sparse matrices is constant-time, - # inplace, and fast.) - - if hasattr(x, "getnnz"): - x = as_sparse_variable(x) - assert x.format in ("csr", "csc") - if hasattr(y, "getnnz"): - y = as_sparse_variable(y) - assert y.format in ("csr", "csc") - - x_is_sparse_variable = _is_sparse_variable(x) - y_is_sparse_variable = _is_sparse_variable(y) - if not x_is_sparse_variable and not y_is_sparse_variable: - raise TypeError("structured_dot requires at least one sparse argument") - - if x_is_sparse_variable: - return _structured_dot(x, y) - else: - assert y_is_sparse_variable - return _structured_dot(y.T, x.T).T - - -class StructuredDotGradCSC(COp): - # Op that produces the grad of StructuredDot. - - # :param a_indices: Matrix indices - # :param a_indptr: Matrix indptr - # :param b: Right operand - # :param g_ab: Accumulated gradient. - - # :return: The grad of `a`.`b` for `a` accumulated - # with g_ab. - - # :note: The grad implemented is structured. - # :note: a_* are the corresponding properties of a sparse - # matrix in csc format. - __props__ = () - - def make_node(self, a_indices, a_indptr, b, g_ab): - return Apply( - self, - [a_indices, a_indptr, b, g_ab], - [tensor(dtype=g_ab.dtype, shape=(None,))], - ) - - def perform(self, node, inputs, outputs): - (a_indices, a_indptr, b, g_ab) = inputs - (out,) = outputs - g_a_data = np.zeros(a_indices.shape, dtype=g_ab.dtype) - for j in range(len(a_indptr) - 1): - ind0 = a_indptr[j] - ind1 = a_indptr[j + 1] - for i_idx in range(ind0, ind1): - i = a_indices[i_idx] - # Depending on the type of g_ab and b (sparse or dense), - # the following dot product can result in a scalar or - # a (1, 1) sparse matrix. - dot_val = np.dot(g_ab[i], b[j].T) - if isinstance(dot_val, scipy.sparse.spmatrix): - dot_val = dot_val[0, 0] - g_a_data[i_idx] = dot_val - out[0] = g_a_data - - def c_code_cache_version(self): - return (2,) - - def c_code(self, node, name, inputs, outputs, sub): - (_indices, _indptr, _d, _g) = inputs - (_zout,) = outputs - if node.inputs[2].type.dtype in ("complex64", "complex128"): - raise NotImplementedError("Complex types are not supported for b") - if node.inputs[3].type.dtype in ("complex64", "complex128"): - raise NotImplementedError("Complex types are not supported for g_ab") - - fail = sub["fail"] - return f""" - if (PyArray_NDIM({_d}) != 2) {{PyErr_SetString(PyExc_NotImplementedError, "rank(d) != 2"); {fail};}} - if (PyArray_NDIM({_g}) != 2) {{PyErr_SetString(PyExc_NotImplementedError, "rank(g) != 2"); {fail};}} - if (PyArray_NDIM({_indices}) != 1) {{PyErr_SetString(PyExc_NotImplementedError, "rank(indices) != 1"); {fail};}} - if (PyArray_NDIM({_indptr}) != 1) {{PyErr_SetString(PyExc_NotImplementedError, "rank(indptr) != 1"); {fail};}} - - if( PyArray_TYPE({_indices}) != NPY_INT32) {{ - PyErr_SetString(PyExc_NotImplementedError, "C"); {fail};}} - - if( PyArray_TYPE({_indptr}) != NPY_INT32) - {{PyErr_SetString(PyExc_NotImplementedError, "D"); {fail};}} - - if( PyArray_DIMS({_d})[1] != PyArray_DIMS({_g})[1]) - {{PyErr_SetString(PyExc_NotImplementedError, "d and g have different numbers of columns"); {fail};}} - - if (!{_zout} - || (PyArray_DIMS({_zout})[0] != PyArray_DIMS({_indices})[0])) - {{ - Py_XDECREF({_zout}); - {_zout} = (PyArrayObject*) PyArray_SimpleNew(1, PyArray_DIMS({_indices}), PyArray_TYPE({_g})); - }} - - {{ //makes it compile even though labels jump over variable definitions. - npy_intp nnz = PyArray_DIMS({_indices})[0]; - npy_intp N = PyArray_DIMS({_indptr})[0]-1; //TODO: error checking with this - - npy_intp Sindices = PyArray_STRIDES({_indices})[0]/PyArray_ITEMSIZE({_indices}); - npy_intp Sindptr = PyArray_STRIDES({_indptr})[0]/PyArray_ITEMSIZE({_indptr}); - - const npy_intp Sd1 = PyArray_STRIDES({_d})[1]/PyArray_ITEMSIZE({_d}); - const npy_intp Sg1 = PyArray_STRIDES({_g})[1]/PyArray_ITEMSIZE({_g}); - - const npy_intp K = PyArray_DIMS({_d})[1]; - - const npy_int32 * __restrict__ indptr = (npy_int32 *)PyArray_DATA({_indptr}); - const npy_int32 * __restrict__ indices = (npy_int32 *)PyArray_DATA({_indices}); - - // loop over columns - for (npy_int32 j = 0; j < N; ++j) - {{ - // extract j-th row of dense matrix - const dtype_{_d}* __restrict__ d_row = (dtype_{_d}*)(PyArray_BYTES({_d}) + PyArray_STRIDES({_d})[0] * j); - if(j >= PyArray_DIMS({_d})[0]) {{PyErr_SetString(PyExc_NotImplementedError, "G"); {fail};}} - - // for each non-null value in the sparse column - for (npy_int32 i_idx = indptr[j * Sindptr]; i_idx < indptr[(j+1) * Sindptr]; ++i_idx) - {{ - // extract row index of non-null value - npy_int32 i = indices[i_idx * Sindices]; - - // extract corresponding row in gradient - const dtype_{_g}* __restrict__ g_row = (dtype_{_g}*)(PyArray_BYTES({_g}) + PyArray_STRIDES({_g})[0] * i); - double ip = 0.0; - - // make sure that row index is not bigger than actual number of rows - // Note: wouldn't the above operation fail if that were the case ? - // when would this ever be true anyway ? - if (i >= PyArray_DIMS({_g})[0]) - {{PyErr_SetString(PyExc_NotImplementedError, "H"); {fail};}} - - // perform dot product of dense and sparse rows - for(int k = 0; k < K; ++k) - {{ - ip += d_row[k * Sd1] * g_row[k*Sg1]; - }} - - // write resulting gradient to sparse output - ((dtype_{_zout}* __restrict__)(PyArray_BYTES({_zout}) + i_idx * PyArray_STRIDES({_zout})[0]))[0] = ip; - }} - }} - }} - - """ - - def infer_shape(self, fgraph, node, shapes): - return [shapes[0]] - - -sdg_csc = StructuredDotGradCSC() - - -class StructuredDotGradCSR(COp): - # Op that produces the grad of StructuredDot. - - # :param a_indices: Matrix indices - # :param a_indptr: Matrix indptr - # :param b: Right operand - # :param g_ab: Accumulated gradient. - - # :return: The grad of `a`.`b` for `a` accumulated - # with g_ab. - - # :note: The grad implemented is structured. - # :note: a_* are the corresponding properties of a sparse - # matrix in csr format. - __props__ = () - - def make_node(self, a_indices, a_indptr, b, g_ab): - return Apply( - self, [a_indices, a_indptr, b, g_ab], [tensor(dtype=b.dtype, shape=(None,))] - ) - - def perform(self, node, inputs, outputs): - (a_indices, a_indptr, b, g_ab) = inputs - (out,) = outputs - g_a_data = np.zeros(a_indices.shape, dtype=g_ab.dtype) - for i in range(len(a_indptr) - 1): # loop over rows - ind0 = a_indptr[i] - ind1 = a_indptr[i + 1] - # loop over values in that row (columns) - for j_idx in range(ind0, ind1): - j = a_indices[j_idx] - # grad is dot product of i-th row of gradient with j-th row of b - # Depending on the type of g_ab and b (sparse or dense), - # the following dot product can result in a scalar or - # a (1, 1) sparse matrix. - dot_val = np.dot(g_ab[i], b[j].T) - if isinstance(dot_val, scipy.sparse.spmatrix): - dot_val = dot_val[0, 0] - g_a_data[j_idx] = dot_val - out[0] = g_a_data - - def c_code_cache_version(self): - return (2,) - - def c_code(self, node, name, inputs, outputs, sub): - (_indices, _indptr, _d, _g) = inputs - (_zout,) = outputs - if node.inputs[2].type.dtype in ("complex64", "complex128"): - raise NotImplementedError("Complex types are not supported for b") - if node.inputs[3].type.dtype in ("complex64", "complex128"): - raise NotImplementedError("Complex types are not supported for g_ab") - - fail = sub["fail"] - return f""" - if (PyArray_NDIM({_d}) != 2) {{PyErr_SetString(PyExc_NotImplementedError, "rank(d) != 2"); {fail};}} - if (PyArray_NDIM({_g}) != 2) {{PyErr_SetString(PyExc_NotImplementedError, "rank(g) != 2"); {fail};}} - if (PyArray_NDIM({_indices}) != 1) {{PyErr_SetString(PyExc_NotImplementedError, "rank(indices) != 1"); {fail};}} - if (PyArray_NDIM({_indptr}) != 1) {{PyErr_SetString(PyExc_NotImplementedError, "rank(indptr) != 1"); {fail};}} - - if( PyArray_TYPE({_indices}) != NPY_INT32) {{ - PyErr_SetString(PyExc_NotImplementedError, "C"); {fail};}} - - if( PyArray_TYPE({_indptr}) != NPY_INT32) - {{PyErr_SetString(PyExc_NotImplementedError, "D"); {fail};}} - - if( PyArray_DIMS({_d})[1] != PyArray_DIMS({_g})[1]) - {{PyErr_SetString(PyExc_NotImplementedError, "d and g have different numbers of columns"); {fail};}} - - if (!{_zout} - || (PyArray_DIMS({_zout})[0] != PyArray_DIMS({_indices})[0])) - {{ - Py_XDECREF({_zout}); - {_zout} = (PyArrayObject*) PyArray_SimpleNew(1, PyArray_DIMS({_indices}), PyArray_TYPE({_g})); - }} - - {{ //makes it compile even though labels jump over variable definitions. - npy_intp nnz = PyArray_DIMS({_indices})[0]; - // extract number of rows - npy_intp N = PyArray_DIMS({_indptr})[0]-1; //TODO: error checking with this - - npy_intp Sindices = PyArray_STRIDES({_indices})[0]/PyArray_ITEMSIZE({_indices}); - npy_intp Sindptr = PyArray_STRIDES({_indptr})[0]/PyArray_ITEMSIZE({_indptr}); - - const npy_intp Sd1 = PyArray_STRIDES({_d})[1]/PyArray_ITEMSIZE({_d}); - const npy_intp Sg1 = PyArray_STRIDES({_g})[1]/PyArray_ITEMSIZE({_g}); - - const npy_intp K = PyArray_DIMS({_d})[1]; - - const npy_int32 * __restrict__ indptr = (npy_int32 *)PyArray_DATA({_indptr}); - const npy_int32 * __restrict__ indices = (npy_int32 *)PyArray_DATA({_indices}); - - // loop over columns of sparse matrix - for (npy_int32 i = 0; i < N; ++i) - {{ - // for each non-null value in the sparse row - for (npy_int32 j_idx = indptr[i * Sindptr]; j_idx < indptr[(i+1) * Sindptr]; ++j_idx) - {{ - // extract column index of non-null value - npy_int32 j = indices[j_idx * Sindices]; - - // extract j-th row of dense matrix - const dtype_{_d}* __restrict__ d_row = (dtype_{_d}*)(PyArray_BYTES({_d}) + PyArray_STRIDES({_d})[0] * j); - if(j >= PyArray_DIMS({_d})[0]) {{PyErr_SetString(PyExc_NotImplementedError, "G"); {fail};}} - - // extract corresponding row in gradient - const dtype_{_g}* __restrict__ g_row = (dtype_{_g}*)(PyArray_BYTES({_g}) + PyArray_STRIDES({_g})[0] * i); - double ip = 0.0; - - // make sure that row index is not bigger than actual number of rows - // Note: wouldn't the above operation fail if that were the case ? - // when would this ever be true anyway ? - if (i >= PyArray_DIMS({_g})[0]) - {{PyErr_SetString(PyExc_NotImplementedError, "H"); {fail};}} - - // perform dot product of dense and sparse rows - for(int k = 0; k < K; ++k) - {{ - ip += d_row[k * Sd1] * g_row[k*Sg1]; - }} - - // write resulting gradient to sparse output - ((dtype_{_zout}* __restrict__)(PyArray_BYTES({_zout}) + j_idx * PyArray_STRIDES({_zout})[0]))[0] = ip; - }} - }} - }} - - """ - - def infer_shape(self, fgraph, node, shapes): - return [shapes[0]] - - -sdg_csr = StructuredDotGradCSR() - - -def structured_dot_grad(sparse_A, dense_B, ga): - if sparse_A.type.format in ("csc", "csr"): - if sparse_A.type.format == "csc": - sdgcsx = sdg_csc - CSx = CSC - else: - sdgcsx = sdg_csr - CSx = CSR - - g_A_data = sdgcsx(csm_indices(sparse_A), csm_indptr(sparse_A), dense_B, ga) - return CSx( - g_A_data, csm_indices(sparse_A), csm_indptr(sparse_A), csm_shape(sparse_A) - ) - else: - raise NotImplementedError() - - -class SamplingDot(Op): - """Compute the dot product ``dot(x, y.T) = z`` for only a subset of `z`. - - This is equivalent to ``p * (x . y.T)`` where ``*`` is the element-wise - product, ``x`` and ``y`` operands of the dot product and ``p`` is a matrix that - contains 1 when the corresponding element of ``z`` should be calculated - and ``0`` when it shouldn't. Note that `SamplingDot` has a different interface - than `dot` because it requires ``x`` to be a ``m x k`` matrix while - ``y`` is a ``n x k`` matrix instead of the usual ``k x n`` matrix. + array + The concatenation of the sparse array row wise. Notes ----- - It will work if the pattern is not binary value, but if the - pattern doesn't have a high sparsity proportion it will be slower - then a more optimized dot followed by a normal elemwise - multiplication. + The number of column of the sparse matrix must agree. The grad implemented is regular, i.e. not structured. """ - __props__ = () - - def make_node(self, x, y, p): - """ - Parameters - ---------- - x - Tensor matrix. - y - Tensor matrix. - p - Sparse matrix in csr format. - - """ - x = ptb.as_tensor_variable(x) - y = ptb.as_tensor_variable(y) - p = as_sparse_variable(p) - assert p.format in ("csr", "csc") - - if not _is_sparse_variable(p): - raise TypeError(p) - - # TODO: use it. - # dtype_out = ps.upcast(x.type.dtype, y.type.dtype, p.type.dtype) - - return Apply(self, [x, y, p], [p.type()]) - - def perform(self, node, inputs, outputs): - (x, y, p) = inputs - (out,) = outputs - if _is_sparse(x): - raise TypeError(x) - - if _is_sparse(y): - raise TypeError(y) - - if not _is_sparse(p): - raise TypeError(p) - - out[0] = p.__class__(p.multiply(np.dot(x, y.T))) - - def grad(self, inputs, gout): - (x, y, p) = inputs - (gz,) = gout - rval = [dot(p * gz, y), dot((p * gz).T, x), grad_not_implemented(self, 2, p)] - - return rval - - def infer_shape(self, fgraph, node, ins_shapes): - return [ins_shapes[2]] - - -sampling_dot = SamplingDot() - - -class Dot(Op): - __props__ = () - - def __str__(self): - return "Sparse" + self.__class__.__name__ - - def infer_shape(self, fgraph, node, shapes): - xshp, yshp = shapes - x, y = node.inputs - if x.ndim == 2 and y.ndim == 2: - return [(xshp[0], yshp[1])] - if x.ndim == 1 and y.ndim == 2: - return [(yshp[1],)] - if x.ndim == 2 and y.ndim == 1: - return [(xshp[0],)] - if x.ndim == 1 and y.ndim == 1: - return [()] - raise NotImplementedError() - - def make_node(self, x, y): - dtype_out = ps.upcast(x.dtype, y.dtype) - - # Sparse dot product should have at least one sparse variable - # as input. If the other one is not sparse, it has to be converted - # into a tensor. - if isinstance(x, scipy.sparse.spmatrix): - x = as_sparse_variable(x) - if isinstance(y, scipy.sparse.spmatrix): - y = as_sparse_variable(y) - - x_is_sparse_var = _is_sparse_variable(x) - y_is_sparse_var = _is_sparse_variable(y) - - if not x_is_sparse_var and not y_is_sparse_var: - raise TypeError( - "Sparse dot product should have at least one " - "sparse variable as inputs, but the inputs are " - f"{x} ({x.type}) and {y} ({y.type})." - ) - - if x_is_sparse_var: - shape_x = (None,) * x.type.ndim - else: - x = ptb.as_tensor_variable(x) - shape_x = x.type.shape - assert y.format in ("csr", "csc") - if x.ndim not in (1, 2): - raise TypeError( - "Input 0 (0-indexed) must have ndim of " - f"1 or 2, {int(x.type.ndim)} given." - ) - - if y_is_sparse_var: - shape_y = (None,) * y.type.ndim - else: - y = ptb.as_tensor_variable(y) - shape_y = y.type.shape - assert x.format in ("csr", "csc") - if y.ndim not in (1, 2): - raise TypeError( - "Input 1 (1-indexed) must have ndim of " - f"1 or 2, {int(y.type.ndim)} given." - ) - - if len(shape_y) == 2: - shape_out = shape_x[:-1] + shape_y[1:] - elif len(shape_y) == 1: - shape_out = shape_x[:-1] - - return Apply(self, [x, y], [tensor(dtype=dtype_out, shape=shape_out)]) - - def perform(self, node, inputs, out): - x, y = inputs - out = out[0] - x_is_sparse = _is_sparse(x) - y_is_sparse = _is_sparse(y) - - if not x_is_sparse and not y_is_sparse: - raise TypeError(x) - - rval = x * y - - if x_is_sparse and y_is_sparse: - rval = rval.toarray() - - out[0] = np.asarray(rval, dtype=node.outputs[0].dtype) - - def grad(self, inputs, gout): - (x, y) = inputs - (gz,) = gout - assert _is_sparse_variable(x) or _is_sparse_variable(y) - rval = [] - - if _is_dense_variable(y): - rval.append(pt_dot(gz, y.T)) - else: - rval.append(dot(gz, y.T)) - if _is_dense_variable(x): - rval.append(pt_dot(x.T, gz)) - else: - rval.append(dot(x.T, gz)) - - return rval - - -_dot = Dot() - - -def dot(x, y): - """Efficiently compute the dot product when one or all operands are sparse. - - Supported formats are CSC and CSR. The output of the operation is dense. + blocks = [as_sparse_variable(i) for i in blocks] + if dtype is None: + dtype = ps.upcast(*[i.dtype for i in blocks]) + return VStack(format=format, dtype=dtype)(*blocks) - Parameters - ---------- - x - Sparse or dense matrix variable. - y - Sparse or dense matrix variable. - Returns - ------- - The dot product ``x @ y`` in a dense format. +class Remove0(Op): + """Remove explicit zeros from a sparse matrix. Notes ----- The grad implemented is regular, i.e. not structured. - At least one of `x` or `y` must be a sparse matrix. - - When the operation has the form ``dot(csr_matrix, dense)`` - the gradient of this operation can be performed inplace - by `UsmmCscDense`. This leads to significant speed-ups. - """ - if hasattr(x, "getnnz"): - x = as_sparse_variable(x) - if hasattr(y, "getnnz"): - y = as_sparse_variable(y) - - x_is_sparse_variable = _is_sparse_variable(x) - y_is_sparse_variable = _is_sparse_variable(y) - - if not x_is_sparse_variable and not y_is_sparse_variable: - raise TypeError() - - return _dot(x, y) - - -class Usmm(Op): - """Computes the dense matrix resulting from ``alpha * x @ y + z``. - - Notes - ----- - At least one of `x` or `y` must be a sparse matrix. - - """ + __props__ = ("inplace",) - __props__ = () + def __init__(self, inplace=False): + self.inplace = inplace + if self.inplace: + self.destroy_map = {0: [0]} def __str__(self): - return "Usmm{no_inplace}" + l = [] + if self.inplace: + l.append("inplace") + return f"{self.__class__.__name__}{{{', '.join(l)}}}" - def make_node(self, alpha, x, y, z): + def make_node(self, x): """ Parameters ---------- - alpha - A scalar. x - Matrix variable. - y - Matrix variable. - z - Dense matrix. + Sparse matrix. """ - if not _is_sparse_variable(x) and not _is_sparse_variable(y): - # If x and y are tensor, we don't want to use this class - # We should use Dot22 and Gemm in that case. - raise TypeError(x) - - dtype_out = ps.upcast( - alpha.type.dtype, x.type.dtype, y.type.dtype, z.type.dtype - ) - alpha = ptb.as_tensor_variable(alpha) - z = ptb.as_tensor_variable(z) - - assert z.type.ndim == 2 - assert alpha.type.shape == (1,) * alpha.type.ndim - if not _is_sparse_variable(x): - x = ptb.as_tensor_variable(x) - assert y.format in ("csr", "csc") - assert x.type.ndim == 2 - if not _is_sparse_variable(y): - y = ptb.as_tensor_variable(y) - assert x.format in ("csr", "csc") - assert y.type.ndim == 2 - - return Apply( - self, - [alpha, x, y, z], - [tensor(dtype=dtype_out, shape=(None, None))], - ) + x = as_sparse_variable(x) + assert x.format in ("csr", "csc") + return Apply(self, [x], [x.type()]) def perform(self, node, inputs, outputs): - (alpha, x, y, z) = inputs - (out,) = outputs - x_is_sparse = _is_sparse(x) - y_is_sparse = _is_sparse(y) - - if not x_is_sparse and not y_is_sparse: - raise TypeError(x) - - rval = x * y - if isinstance(rval, scipy.sparse.spmatrix): - rval = rval.toarray() - if rval.dtype == alpha.dtype: - rval *= alpha # Faster because operation is inplace - else: - rval = rval * alpha - if rval.dtype == z.dtype: - rval += z # Faster because operation is inplace + (x,) = inputs + (z,) = outputs + if self.inplace: + c = x else: - rval = rval + z + c = x.copy() + c.eliminate_zeros() + z[0] = c + + def grad(self, inputs, gout): + (_x,) = inputs + (gz,) = gout + return [gz] - out[0] = rval + def infer_shape(self, fgraph, node, i0_shapes): + return i0_shapes -usmm = Usmm() +remove0 = Remove0() class ConstructSparseFromList(Op): @@ -4285,88 +1944,3 @@ def grad(self, inputs, grads): construct_sparse_from_list = ConstructSparseFromList() - - -class SparseBlockDiagonal(BaseBlockDiagonal): - __props__ = ( - "n_inputs", - "format", - ) - - def __init__(self, n_inputs: int, format: Literal["csc", "csr"] = "csc"): - super().__init__(n_inputs) - self.format = format - - def make_node(self, *matrices): - matrices = self._validate_and_prepare_inputs( - matrices, as_sparse_or_tensor_variable - ) - dtype = _largest_common_dtype(matrices) - out_type = matrix(format=self.format, dtype=dtype) - - return Apply(self, matrices, [out_type]) - - def perform(self, node, inputs, output_storage, params=None): - dtype = node.outputs[0].type.dtype - output_storage[0][0] = scipy.sparse.block_diag( - inputs, format=self.format - ).astype(dtype) - - -def block_diag(*matrices: TensorVariable, format: Literal["csc", "csr"] = "csc"): - r""" - Construct a block diagonal matrix from a sequence of input matrices. - - Given the inputs `A`, `B` and `C`, the output will have these arrays arranged on the diagonal: - - [[A, 0, 0], - [0, B, 0], - [0, 0, C]] - - Parameters - ---------- - A, B, C ... : tensors - Input tensors to form the block diagonal matrix. last two dimensions of the inputs will be used, and all - inputs should have at least 2 dimensins. - - Note that the input matrices need not be sparse themselves, and will be automatically converted to the - requested format if they are not. - - format: str, optional - The format of the output sparse matrix. One of 'csr' or 'csc'. Default is 'csr'. Ignored if sparse=False. - - Returns - ------- - out: sparse matrix tensor - Symbolic sparse matrix in the specified format. - - Examples - -------- - Create a sparse block diagonal matrix from two sparse 2x2 matrices: - - .. testcode:: - import numpy as np - from pytensor.sparse import block_diag - from scipy.sparse import csr_matrix - - A = csr_matrix([[1, 2], [3, 4]]) - B = csr_matrix([[5, 6], [7, 8]]) - result_sparse = block_diag(A, B, format='csr') - - print(result_sparse) - print(result_sparse.toarray().eval()) - - .. testoutput:: - - SparseVariable{csr,int64} - [[1 2 0 0] - [3 4 0 0] - [0 0 5 6] - [0 0 7 8]] - - """ - if len(matrices) == 1: - return matrices - - _sparse_block_diagonal = SparseBlockDiagonal(n_inputs=len(matrices), format=format) - return _sparse_block_diagonal(*matrices) diff --git a/pytensor/sparse/linalg.py b/pytensor/sparse/linalg.py new file mode 100644 index 0000000000..6f04ad04d3 --- /dev/null +++ b/pytensor/sparse/linalg.py @@ -0,0 +1,93 @@ +from typing import Literal + +import scipy.sparse + +from pytensor.graph import Apply +from pytensor.sparse import as_sparse_or_tensor_variable, matrix +from pytensor.tensor import TensorVariable +from pytensor.tensor.slinalg import BaseBlockDiagonal, _largest_common_dtype + + +class SparseBlockDiagonal(BaseBlockDiagonal): + __props__ = ( + "n_inputs", + "format", + ) + + def __init__(self, n_inputs: int, format: Literal["csc", "csr"] = "csc"): + super().__init__(n_inputs) + self.format = format + + def make_node(self, *matrices): + matrices = self._validate_and_prepare_inputs( + matrices, as_sparse_or_tensor_variable + ) + dtype = _largest_common_dtype(matrices) + out_type = matrix(format=self.format, dtype=dtype) + + return Apply(self, matrices, [out_type]) + + def perform(self, node, inputs, output_storage, params=None): + dtype = node.outputs[0].type.dtype + output_storage[0][0] = scipy.sparse.block_diag( + inputs, format=self.format + ).astype(dtype) + + +def block_diag(*matrices: TensorVariable, format: Literal["csc", "csr"] = "csc"): + r""" + Construct a block diagonal matrix from a sequence of input matrices. + + Given the inputs `A`, `B` and `C`, the output will have these arrays arranged on the diagonal: + + [[A, 0, 0], + [0, B, 0], + [0, 0, C]] + + Parameters + ---------- + A, B, C ... : tensors + Input tensors to form the block diagonal matrix. last two dimensions of the inputs will be used, and all + inputs should have at least 2 dimensins. + + Note that the input matrices need not be sparse themselves, and will be automatically converted to the + requested format if they are not. + + format: str, optional + The format of the output sparse matrix. One of 'csr' or 'csc'. Default is 'csr'. Ignored if sparse=False. + + Returns + ------- + out: sparse matrix tensor + Symbolic sparse matrix in the specified format. + + Examples + -------- + Create a sparse block diagonal matrix from two sparse 2x2 matrices: + + .. testcode:: + import numpy as np + from pytensor.sparse.linalg import block_diag + from scipy.sparse import csr_matrix + + A = csr_matrix([[1, 2], [3, 4]]) + B = csr_matrix([[5, 6], [7, 8]]) + result_sparse = block_diag(A, B, format='csr') + + print(result_sparse) + print(result_sparse.toarray().eval()) + + .. testoutput:: + + SparseVariable{csr,int64} + [[1 2 0 0] + [3 4 0 0] + [0 0 5 6] + [0 0 7 8]] + + """ + if len(matrices) == 1: + return matrices + + _sparse_block_diagonal = SparseBlockDiagonal(n_inputs=len(matrices), format=format) + return _sparse_block_diagonal(*matrices) diff --git a/pytensor/sparse/math.py b/pytensor/sparse/math.py new file mode 100644 index 0000000000..c2244341ae --- /dev/null +++ b/pytensor/sparse/math.py @@ -0,0 +1,2092 @@ +from functools import wraps +from warnings import warn + +import numpy as np +import scipy.sparse as scipy_sparse + +import pytensor.scalar as ps +import pytensor.sparse.basic as psb +import pytensor.tensor.basic as ptb +import pytensor.tensor.math as ptm +from pytensor import config +from pytensor.gradient import grad_not_implemented +from pytensor.graph import Apply, Op +from pytensor.link.c.op import COp +from pytensor.sparse.basic import SparseTensorType +from pytensor.tensor import TensorType, Variable, specify_broadcastable, tensor +from pytensor.tensor.type import complex_dtypes + + +def structured_elemwise(tensor_op): + """ + A decorator to create structured element-wise operations on sparse matrices. + + An operation is called "structured" if it operates only on the non-zeros elements of a sparse matrix. + """ + + def decorator(f): + @wraps(f) + def wrapper(*args): + x = psb.as_sparse_variable(args[0]) + assert x.format in ("csr", "csc") + + xs = [ps.as_scalar(arg) for arg in args[1:]] + data, ind, ptr, _shape = psb.csm_properties(x) + data = tensor_op(data, *xs) + + return psb.CSM(x.format)(data, ind, ptr, _shape) + + return wrapper + + return decorator + + +@structured_elemwise(ptm.abs) +def structured_abs(x): + """ + Compute abs(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.sigmoid) +def structured_sigmoid(x): + """ + Compute sigmoid(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.exp) +def structured_exp(x): + """ + Compute exp(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.exp2) +def structured_exp2(x): + """ + Compute exp2(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.log) +def structured_log(x): + """ + Compute log(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.log2) +def structured_log2(x): + """ + Compute log2(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.log10) +def structured_log10(x): + """ + Compute log10(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.pow) +def structured_pow(x, y): + """ + Compute x**y for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.minimum) +def structured_minimum(x, y): + """ + Compute min(x, y) for all non-zero elements of x, where y is a scalar. + """ + + +@structured_elemwise(ptm.maximum) +def structured_maximum(x, y): + """ + Compute max(x, y) for all non-zero elements of x, where y is a scalar. + """ + + +@structured_elemwise(ptm.add) +def structured_add(x, y): + """ + Compute x + y for all non-zero elements of x, where y is a scalar. + """ + + +@structured_elemwise(ptm.sin) +def structured_sin(x): + """ + Compute sin(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.sinh) +def structured_sinh(x): + """ + Compute sinh(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.arcsin) +def structured_arcsin(x): + """ + Compute arcsin(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.arcsinh) +def structured_arcsinh(x): + """ + Compute arcsinh(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.cos) +def structured_cos(x): + """ + Compute cos(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.cosh) +def structured_cosh(x): + """ + Compute cosh(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.arccos) +def structured_arccos(x): + """ + Compute arccos(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.arccosh) +def structured_arccosh(x): + """ + Compute arccosh(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.tan) +def structured_tan(x): + """ + Compute tan(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.tanh) +def structured_tanh(x): + """ + Compute tanh(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.arctan) +def structured_arctan(x): + """ + Compute arctan(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.arctanh) +def structured_arctanh(x): + """ + Compute arctanh(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.round_half_to_even) +def structured_rint(x): + """ + Compute round_half_to_even(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.sign) +def structured_sign(x): + """ + Compute sign(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.ceil) +def structured_ceil(x): + """ + Compute ceil(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.floor) +def structured_floor(x): + """ + Compute floor(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.log1p) +def structured_log1p(x): + """ + Compute log(1 + x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.expm1) +def structured_expm1(x): + """ + Compute exp(x) - 1 for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.deg2rad) +def structured_deg2rad(x): + """ + Convert degrees to radians for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.rad2deg) +def structured_rad2deg(x): + """ + Convert radians to degrees for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.trunc) +def structured_trunc(x): + """ + Truncate the decimal part of x for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.sqr) +def structured_sqr(x): + """ + Compute sqr(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.sqrt) +def structured_sqrt(x): + """ + Compute sqrt(x) for all non-zero elements of x. + """ + + +@structured_elemwise(ptm.conjugate) +def _conj(x): + """ + Compute the complex conjugate of x for all non-zero elements of x. + """ + + +def conjugate(x): + _x = psb.as_sparse_variable(x) + if _x.type.dtype not in complex_dtypes: + return _x + return _conj(_x) + + +structured_conjugate = conjugate + + +class SpSum(Op): + """ + + WARNING: judgement call... + We are not using the structured in the comparison or hashing + because it doesn't change the perform method therefore, we + *do* want Sums with different structured values to be merged + by the merge optimization and this requires them to compare equal. + """ + + __props__ = ("axis",) + + def __init__(self, axis=None, sparse_grad=True): + super().__init__() + self.axis = axis + self.structured = sparse_grad + if self.axis not in (None, 0, 1): + raise ValueError("Illegal value for self.axis.") + + def make_node(self, x): + x = psb.as_sparse_variable(x) + assert x.format in ("csr", "csc") + + if self.axis is not None: + out_shape = (None,) + else: + out_shape = () + + z = TensorType(dtype=x.dtype, shape=out_shape)() + return Apply(self, [x], [z]) + + def perform(self, node, inputs, outputs): + (x,) = inputs + (z,) = outputs + if self.axis is None: + z[0] = np.asarray(x.sum()) + else: + z[0] = np.asarray(x.sum(self.axis)).ravel() + + def grad(self, inputs, gout): + (x,) = inputs + (gz,) = gout + if x.dtype not in psb.continuous_dtypes: + return [x.zeros_like(dtype=config.floatX)] + if self.structured: + if self.axis is None: + r = gz * psb.sp_ones_like(x) + elif self.axis == 0: + r = psb.col_scale(psb.sp_ones_like(x), gz) + elif self.axis == 1: + r = psb.row_scale(psb.sp_ones_like(x), gz) + else: + raise ValueError("Illegal value for self.axis.") + else: + o_format = x.format + x = psb.dense_from_sparse(x) + if psb._is_sparse_variable(gz): + gz = psb.dense_from_sparse(gz) + if self.axis is None: + r = ptb.second(x, gz) + else: + ones = ptb.ones_like(x) + if self.axis == 0: + r = specify_broadcastable(gz.dimshuffle("x", 0), 0) * ones + elif self.axis == 1: + r = specify_broadcastable(gz.dimshuffle(0, "x"), 1) * ones + else: + raise ValueError("Illegal value for self.axis.") + r = psb.SparseFromDense(o_format)(r) + return [r] + + def infer_shape(self, fgraph, node, shapes): + r = None + if self.axis is None: + r = [()] + elif self.axis == 0: + r = [(shapes[0][1],)] + else: + r = [(shapes[0][0],)] + return r + + def __str__(self): + return f"{self.__class__.__name__}{{axis={self.axis}}}" + + +def sp_sum(x, axis=None, sparse_grad=False): + """ + Calculate the sum of a sparse matrix along the specified axis. + + It operates a reduction along the specified axis. When `axis` is `None`, + it is applied along all axes. + + Parameters + ---------- + x + Sparse matrix. + axis + Axis along which the sum is applied. Integer or `None`. + sparse_grad : bool + `True` to have a structured grad. + + Returns + ------- + object + The sum of `x` in a dense format. + + Notes + ----- + The grad implementation is controlled with the `sparse_grad` parameter. + `True` will provide a structured grad and `False` will provide a regular + grad. For both choices, the grad returns a sparse matrix having the same + format as `x`. + + This op does not return a sparse matrix, but a dense tensor matrix. + + """ + + return SpSum(axis, sparse_grad)(x) + + +class AddSS(Op): + # add(sparse, sparse). + # see the doc of add() for more detail. + __props__ = () + + def make_node(self, x, y): + x, y = map(psb.as_sparse_variable, [x, y]) + assert x.format in ("csr", "csc") + assert y.format in ("csr", "csc") + out_dtype = ps.upcast(x.type.dtype, y.type.dtype) + return Apply( + self, + [x, y], + [psb.SparseTensorType(dtype=out_dtype, format=x.type.format)()], + ) + + def perform(self, node, inputs, outputs): + (x, y) = inputs + (out,) = outputs + assert psb._is_sparse(x) and psb._is_sparse(y) + assert x.shape == y.shape + out[0] = x + y + + def grad(self, inputs, gout): + (x, y) = inputs + (gz,) = gout + assert psb._is_sparse_variable(x) and psb._is_sparse_variable(y) + assert psb._is_sparse_variable(gz) + return gz, gz + + def infer_shape(self, fgraph, node, shapes): + return [shapes[0]] + + +add_s_s = AddSS() + + +class AddSSData(Op): + """Add two sparse matrices assuming they have the same sparsity pattern. + + Notes + ----- + The grad implemented is structured. + + """ + + __props__ = () + + def make_node(self, x, y): + """ + + Parameters + ---------- + x + Sparse matrix. + y + Sparse matrix. + + Notes + ----- + `x` and `y` are assumed to have the same sparsity pattern. + + """ + x, y = map(psb.as_sparse_variable, [x, y]) + assert x.format in ("csr", "csc") + assert y.format in ("csr", "csc") + if x.type.dtype != y.type.dtype: + raise NotImplementedError() + if x.type.format != y.type.format: + raise NotImplementedError() + return Apply( + self, + [x, y], + [psb.SparseTensorType(dtype=x.type.dtype, format=x.type.format)()], + ) + + def perform(self, node, inputs, outputs): + (x, y) = inputs + (out,) = outputs + assert psb._is_sparse(x) and psb._is_sparse(y) + assert x.shape == y.shape + assert x.data.shape == y.data.shape + out[0] = x.copy() + out[0].data += y.data + + def grad(self, inputs, gout): + (gz,) = gout + is_continuous = [(i.dtype in psb.continuous_dtypes) for i in inputs] + derivative = {True: gz, False: None} + return [derivative[b] for b in is_continuous] + + def infer_shape(self, fgraph, node, ins_shapes): + return [ins_shapes[0]] + + +add_s_s_data = AddSSData() + + +class AddSD(Op): + # add(sparse, sparse). + # see the doc of add() for more detail. + __props__ = () + + def make_node(self, x, y): + x, y = psb.as_sparse_variable(x), ptb.as_tensor_variable(y) + assert x.format in ("csr", "csc") + out_dtype = ps.upcast(x.type.dtype, y.type.dtype) + + # The magic number two here arises because L{scipy.sparse} + # objects must be matrices (have dimension 2) + assert y.type.ndim == 2 + return Apply( + self, + [x, y], + [TensorType(dtype=out_dtype, shape=y.type.shape)()], + ) + + def perform(self, node, inputs, outputs): + (x, y) = inputs + (out,) = outputs + assert psb._is_dense(y) + + # The asarray is needed as in some case, this return a + # numpy.matrixlib.defmatrix.matrix object and not an ndarray. + out[0] = np.asarray(x + y, dtype=node.outputs[0].type.dtype) + + def grad(self, inputs, gout): + (x, y) = inputs + (gz,) = gout + assert psb._is_sparse_variable(x) and psb._is_dense_variable(y) + assert psb._is_dense_variable(gz) + return psb.sp_ones_like(x) * gz, gz + + def infer_shape(self, fgraph, node, shapes): + return [shapes[1]] + + +add_s_d = AddSD() + + +class StructuredAddSV(Op): + """Structured addition of a sparse matrix and a dense vector. + + The elements of the vector are only added to the corresponding + non-zero elements of the sparse matrix. Therefore, this operation + outputs another sparse matrix. + + Notes + ----- + The grad implemented is structured since the op is structured. + + """ + + __props__ = () + + def make_node(self, x, y): + """ + Parameters + ---------- + x + Sparse matrix. + y + Tensor type vector. + + """ + x = psb.as_sparse_variable(x) + assert x.format in ("csr", "csc") + y = ptb.as_tensor_variable(y) + + assert y.type.ndim == 1 + + if x.type.dtype != y.type.dtype: + raise NotImplementedError() + return Apply( + self, + [x, y], + [psb.SparseTensorType(dtype=x.type.dtype, format=x.type.format)()], + ) + + def perform(self, node, inputs, outputs): + (x, y) = inputs + (out,) = outputs + assert psb._is_sparse(x) and not psb._is_sparse(y) + assert x.shape[1] == y.shape[0] + out[0] = x.__class__(x + (x.toarray() != 0) * y) + + def grad(self, inputs, gout): + (x, y) = inputs + (gz,) = gout + assert psb._is_sparse_variable(x) and not psb._is_sparse_variable(y) + assert psb._is_sparse_variable(gz) + return gz, sp_sum(gz, axis=0, sparse_grad=True) + + def infer_shape(self, fgraph, node, ins_shapes): + return [ins_shapes[0]] + + +structured_add_s_v = StructuredAddSV() + + +def add(x, y): + """ + Add two matrices, at least one of which is sparse. + + This method will provide the right op according + to the inputs. + + Parameters + ---------- + x + A matrix variable. + y + A matrix variable. + + Returns + ------- + A sparse matrix + `x` + `y` + + Notes + ----- + At least one of `x` and `y` must be a sparse matrix. + + The grad will be structured only when one of the variable will be a dense + matrix. + + """ + + if hasattr(x, "getnnz"): + x = psb.as_sparse_variable(x) + if hasattr(y, "getnnz"): + y = psb.as_sparse_variable(y) + if not isinstance(x, Variable): + x = ptb.as_tensor_variable(x) + if not isinstance(y, Variable): + y = ptb.as_tensor_variable(y) + + x_is_sparse_variable = psb._is_sparse_variable(x) + y_is_sparse_variable = psb._is_sparse_variable(y) + + assert x_is_sparse_variable or y_is_sparse_variable + if x_is_sparse_variable and y_is_sparse_variable: + return add_s_s(x, y) + elif x_is_sparse_variable and not y_is_sparse_variable: + return add_s_d(x, y) + elif y_is_sparse_variable and not x_is_sparse_variable: + return add_s_d(y, x) + else: + raise NotImplementedError() + + +def subtract( + x: SparseTensorType | TensorType, y: SparseTensorType | TensorType +) -> SparseTensorType: + """ + Subtract two matrices, at least one of which is sparse. + + This method will provide the right op according to the inputs. + + Parameters + ---------- + x : SparseVariable or TensorVariable + A matrix variable. + y : SparseVariable or TensorVariable + A matrix variable. + + Returns + ------- + result: SparseVariable + Result of `x - y`, as a sparse matrix. + + Notes + ----- + At least one of `x` and `y` must be a sparse matrix. + + The grad will be structured only when one of the variable will be a dense matrix. + """ + return x + (-y) + + +def sub(x, y): + warn( + "pytensor.sparse.sub is deprecated and will be removed in a future version. Use " + "pytensor.sparse.subtract instead.", + category=DeprecationWarning, + stacklevel=2, + ) + + return subtract(x, y) + + +sub.__doc__ = subtract.__doc__ + + +class MulSS(Op): + # mul(sparse, sparse) + # See the doc of mul() for more detail + __props__ = () + + def make_node(self, x, y): + x, y = psb.as_sparse_variable(x), psb.as_sparse_variable(y) + assert x.format in ("csr", "csc") + assert y.format in ("csr", "csc") + out_dtype = ps.upcast(x.type.dtype, y.type.dtype) + return Apply( + self, + [x, y], + [psb.SparseTensorType(dtype=out_dtype, format=x.type.format)()], + ) + + def perform(self, node, inputs, outputs): + (x, y) = inputs + (out,) = outputs + assert psb._is_sparse(x) and psb._is_sparse(y) + assert len(x.shape) == 2 + assert y.shape == x.shape + # This calls the element-wise multiple + # x * y calls dot... + out[0] = x.multiply(y) + + def grad(self, inputs, gout): + (x, y) = inputs + (gz,) = gout + return y * gz, x * gz + + def infer_shape(self, fgraph, node, shapes): + return [shapes[0]] + + +mul_s_s = MulSS() + + +class MulSD(Op): + # mul(sparse, dense) + # See the doc of mul() for more detail + __props__ = () + + def make_node(self, x, y): + x, y = psb.as_sparse_variable(x), ptb.as_tensor_variable(y) + + assert x.format in ("csr", "csc") + + # upcast the tensor. Is the cast of sparse done implemented? + dtype = ps.upcast(x.type.dtype, y.type.dtype) + + # The magic number two here arises because L{scipy.sparse} + # objects must be matrices (have dimension 2) + # Broadcasting of the sparse matrix is not supported. + # We support nd == 0 used by grad of SpSum() + assert y.type.ndim in (0, 2) + out = psb.SparseTensorType(dtype=dtype, format=x.type.format)() + return Apply(self, [x, y], [out]) + + def perform(self, node, inputs, outputs): + (x, y) = inputs + (out,) = outputs + assert psb._is_sparse(x) and psb._is_dense(y) + if len(y.shape) == 0: + out_dtype = node.outputs[0].dtype + if x.dtype == out_dtype: + z = x.copy() + else: + z = x.astype(out_dtype) + out[0] = z + out[0].data *= y + elif len(y.shape) == 1: + raise NotImplementedError() # RowScale / ColScale + elif len(y.shape) == 2: + # if we have enough memory to fit y, maybe we can fit x.asarray() + # too? + # TODO: change runtime from O(M*N) to O(nonzeros) + M, N = x.shape + assert x.shape == y.shape + out_dtype = node.outputs[0].dtype + + if x.format == "csc": + indices = x.indices + indptr = x.indptr + if x.dtype == out_dtype: + z = x.copy() + else: + z = x.astype(out_dtype) + z_data = z.data + + for j in range(0, N): + for i_idx in range(indptr[j], indptr[j + 1]): + i = indices[i_idx] + z_data[i_idx] *= y[i, j] + out[0] = z + elif x.format == "csr": + indices = x.indices + indptr = x.indptr + if x.dtype == out_dtype: + z = x.copy() + else: + z = x.astype(out_dtype) + z_data = z.data + + for i in range(0, M): + for j_idx in range(indptr[i], indptr[i + 1]): + j = indices[j_idx] + z_data[j_idx] *= y[i, j] + out[0] = z + else: + warn( + "This implementation of MulSD is deficient: {x.format}", + ) + out[0] = type(x)(x.toarray() * y) + + def grad(self, inputs, gout): + (x, y) = inputs + (gz,) = gout + assert psb._is_sparse_variable(x) and psb._is_dense_variable(y) + assert psb._is_sparse_variable(gz) + return y * gz, psb.dense_from_sparse(x * gz) + + def infer_shape(self, fgraph, node, shapes): + return [shapes[0]] + + +mul_s_d = MulSD() + + +class MulSV(Op): + """Element-wise multiplication of sparse matrix by a broadcasted dense vector element wise. + + Notes + ----- + The grad implemented is regular, i.e. not structured. + + """ + + __props__ = () + + def make_node(self, x, y): + """ + Parameters + ---------- + x + Sparse matrix to multiply. + y + Tensor broadcastable vector. + + """ + x = psb.as_sparse_variable(x) + assert x.format in ("csr", "csc") + y = ptb.as_tensor_variable(y) + + assert y.type.ndim == 1 + + if x.type.dtype != y.type.dtype: + raise NotImplementedError( + "MulSV not implemented for differing dtypes." + f"Got {x.type.dtype} and {y.type.dtype}." + ) + return Apply( + self, + [x, y], + [psb.SparseTensorType(dtype=x.type.dtype, format=x.type.format)()], + ) + + def perform(self, node, inputs, outputs): + (x, y) = inputs + (out,) = outputs + assert psb._is_sparse(x) and not psb._is_sparse(y) + assert x.shape[1] == y.shape[0] + out[0] = x.__class__(x.toarray() * y) + + def grad(self, inputs, gout): + (x, y) = inputs + (gz,) = gout + assert psb._is_sparse_variable(x) and psb._is_dense_variable(y) + assert psb._is_sparse_variable(gz) + + # mul_s_v is not implemented if the types vary + + if gz.dtype == "float64" and y.dtype == "float32": + y = y.astype("float64") + + if gz.dtype == "float32" and y.dtype == "float64": + gz = gz.astype("float64") + + return mul_s_v(gz, y), sp_sum(x * gz, axis=0, sparse_grad=True) + + def infer_shape(self, fgraph, node, ins_shapes): + return [ins_shapes[0]] + + +mul_s_v = MulSV() + + +def multiply( + x: SparseTensorType | TensorType, y: SparseTensorType | TensorType +) -> SparseTensorType: + """ + Multiply elementwise two matrices, at least one of which is sparse. + + This method will provide the right op according to the inputs. + + Parameters + ---------- + x : SparseTensorType + A matrix variable. + y : SparseTensorType + A matrix variable. + + Returns + ------- + result: SparseTensorType + The elementwise multiplication of `x` and `y`. + + Notes + ----- + At least one of `x` and `y` must be a sparse matrix. + + The gradient is regular, i.e. not structured. + """ + + x = psb.as_sparse_or_tensor_variable(x) + y = psb.as_sparse_or_tensor_variable(y) + + x_is_sparse_variable = psb._is_sparse_variable(x) + y_is_sparse_variable = psb._is_sparse_variable(y) + + assert x_is_sparse_variable or y_is_sparse_variable + if x_is_sparse_variable and y_is_sparse_variable: + # mul_s_s is not implemented if the types differ + if y.dtype == "float64" and x.dtype == "float32": + x = x.astype("float64") + + return mul_s_s(x, y) + elif x_is_sparse_variable and not y_is_sparse_variable: + # mul is unimplemented if the dtypes differ + if y.dtype == "float64" and x.dtype == "float32": + x = x.astype("float64") + + return mul_s_d(x, y) + elif y_is_sparse_variable and not x_is_sparse_variable: + return mul_s_d(y, x) + else: + raise NotImplementedError() + + +def mul(x, y): + warn( + "pytensor.sparse.mul is deprecated and will be removed in a future version. Use " + "pytensor.sparse.multiply instead.", + category=DeprecationWarning, + stacklevel=2, + ) + + return multiply(x, y) + + +mul.__doc__ = multiply.__doc__ + + +class __ComparisonOpSS(Op): + """ + Used as a superclass for all comparisons between two sparses matrices. + + Parameters + ---------- + x + First compared sparse matrix. + y + Second compared sparse matrix + + Returns + ------- + object + Comparison(x,y) + + """ + + __props__ = () + + # Function to override + def comparison(self, x, y): + raise NotImplementedError() + + def make_node(self, x, y): + x = psb.as_sparse_variable(x) + y = psb.as_sparse_variable(y) + + if x.type.format != y.type.format: + raise NotImplementedError() + return Apply( + self, [x, y], [psb.SparseTensorType(dtype="uint8", format=x.type.format)()] + ) + + def perform(self, node, inputs, outputs): + (x, y) = inputs + (out,) = outputs + assert psb._is_sparse(x) and psb._is_sparse(y) + assert x.shape == y.shape + out[0] = self.comparison(x, y).astype("uint8") + + def infer_shape(self, fgraph, node, ins_shapes): + return [ins_shapes[0]] + + +class __ComparisonOpSD(Op): + """ + Used as a superclass for all comparisons between sparse and dense matrix. + + Parameters + ---------- + x + Sparse matrix. + y + Dense matrix. + + Returns + ------- + object + Comparison(x,y) + + """ + + __props__ = () + + # Function to override + def comparison(self, x, y): + raise NotImplementedError() + + def make_node(self, x, y): + x, y = psb.as_sparse_variable(x), ptb.as_tensor_variable(y) + + assert y.type.ndim == 2 + out = TensorType(dtype="uint8", shape=(None, None))() + return Apply(self, [x, y], [out]) + + def perform(self, node, inputs, outputs): + (x, y) = inputs + (out,) = outputs + assert psb._is_sparse(x) + assert x.shape == y.shape + assert psb._is_dense(y) + o = self.comparison(x, y).astype("uint8") + o = np.asarray(o) + out[0] = o + + def infer_shape(self, fgraph, node, ins_shapes): + return [ins_shapes[0]] + + +def __ComparisonSwitch(SS, SD, DS): + """ + + Parameters + ---------- + SS + Function to apply between two sparses matrices. + SD + Function to apply between a sparse and a dense matrix. + DS + Function to apply between a dense and a sparse matrix. + + Returns + ------- + function + Switch function taking two matrices as input. + + Notes + ----- + At least one of `x` and `y` must be a sparse matrix. + + DS swap input as a dense matrix cannot be a left operand. + + """ + + def helper(x, y): + if hasattr(x, "getnnz"): + x = psb.as_sparse_variable(x) + if hasattr(y, "getnnz"): + y = psb.as_sparse_variable(y) + if not isinstance(x, Variable): + x = ptb.as_tensor_variable(x) + if not isinstance(y, Variable): + y = ptb.as_tensor_variable(y) + + x_is_sparse_variable = psb._is_sparse_variable(x) + y_is_sparse_variable = psb._is_sparse_variable(y) + + assert x_is_sparse_variable or y_is_sparse_variable + if x_is_sparse_variable and y_is_sparse_variable: + return SS(x, y) + elif x_is_sparse_variable and not y_is_sparse_variable: + return SD(x, y) + elif y_is_sparse_variable and not x_is_sparse_variable: + return DS(y, x) + else: + raise NotImplementedError() + + return helper + + +class EqualSS(__ComparisonOpSS): + def comparison(self, x, y): + return x == y + + +equal_s_s = EqualSS() + + +class EqualSD(__ComparisonOpSD): + def comparison(self, x, y): + return x == y + + +equal_s_d = EqualSD() + + +class NotEqualSS(__ComparisonOpSS): + def comparison(self, x, y): + return x != y + + +not_equal_s_s = NotEqualSS() + + +class NotEqualSD(__ComparisonOpSD): + def comparison(self, x, y): + return x != y + + +not_equal_s_d = NotEqualSD() + + +class LessThanSS(__ComparisonOpSS): + def comparison(self, x, y): + return x < y + + +less_than_s_s = LessThanSS() + + +class LessThanSD(__ComparisonOpSD): + def comparison(self, x, y): + return x < y + + +less_than_s_d = LessThanSD() + + +class GreaterThanSS(__ComparisonOpSS): + def comparison(self, x, y): + return x > y + + +greater_than_s_s = GreaterThanSS() + + +class GreaterThanSD(__ComparisonOpSD): + def comparison(self, x, y): + return x > y + + +greater_than_s_d = GreaterThanSD() + + +class LessEqualSS(__ComparisonOpSS): + def comparison(self, x, y): + return x <= y + + +less_equal_s_s = LessEqualSS() + + +class LessEqualSD(__ComparisonOpSD): + def comparison(self, x, y): + return x <= y + + +less_equal_s_d = LessEqualSD() + + +class GreaterEqualSS(__ComparisonOpSS): + def comparison(self, x, y): + return x >= y + + +greater_equal_s_s = GreaterEqualSS() + + +class GreaterEqualSD(__ComparisonOpSD): + def comparison(self, x, y): + return x >= y + + +greater_equal_s_d = GreaterEqualSD() + +eq = __ComparisonSwitch(equal_s_s, equal_s_d, equal_s_d) + +neq = __ComparisonSwitch(not_equal_s_s, not_equal_s_d, not_equal_s_d) + +lt = __ComparisonSwitch(less_than_s_s, less_than_s_d, greater_than_s_d) + +gt = __ComparisonSwitch(greater_than_s_s, greater_than_s_d, less_than_s_d) + +le = __ComparisonSwitch(less_equal_s_s, less_equal_s_d, greater_equal_s_d) + +ge = __ComparisonSwitch(greater_equal_s_s, greater_equal_s_d, less_equal_s_d) + + +class TrueDot(Op): + # TODO + # Simplify code by splitting into DotSS and DotSD. + + __props__ = () + + # The grad_preserves_dense attribute doesn't change the + # execution behavior. To let the optimizer merge nodes with + # different values of this attribute we shouldn't compare it + # here. + + def __init__(self, grad_preserves_dense=True): + self.grad_preserves_dense = grad_preserves_dense + + def make_node(self, x, y): + # NOTE + # Because of trickiness of implementing, + # we assume that the left argument x is a + # SparseVariable (not dense) + + if x.type.dtype != y.type.dtype: + raise NotImplementedError() + + if not psb._is_sparse_variable(x): + raise TypeError(x) + + # These are the conversions performed by scipy.sparse.dot + if x.type.format == "csc" or x.type.format == "coo": + myformat = "csc" + elif x.type.format == "csr": + myformat = "csr" + else: + raise NotImplementedError() + + inputs = [x, y] # Need to convert? e.g. assparse + outputs = [psb.SparseTensorType(dtype=x.type.dtype, format=myformat)()] + return Apply(self, inputs, outputs) + + def perform(self, node, inp, out_): + # TODO + # -Verify that output is sufficiently sparse, + # and raise a warning if it is not. + # -Also determine that we are storing the + # output in the best storage format? + + x, y = inp + (out,) = out_ + rval = x.dot(y) + if not scipy_sparse.issparse(rval): + rval = getattr(scipy_sparse, x.format + "_matrix")(rval) + # x.dot call tocsr() that will "upcast" to ['int8', 'uint8', 'short', + # 'ushort', 'intc', 'uintc', 'longlong', 'ulonglong', 'single', + # 'double', 'longdouble', 'csingle', 'cdouble', 'clongdouble'] + # But ulonglong is uint64 on x86-64, but with a different typenum! + if rval.dtype.num != np.dtype(str(rval.dtype)).num: + assert str(rval.dtype) == node.outputs[0].dtype + # Create a view with the expected typenum. + format = node.outputs[0].type.format + data = rval.data.view(dtype=node.outputs[0].dtype) + indices = rval.indices + indptr = rval.indptr + _shape = rval.shape + # No need to copy indices and indptr as in CSM.perform(), + # as there is only one user of them. + if format == "csc": + rval = scipy_sparse.csc_matrix( + (data, indices, indptr), _shape, copy=False + ) + else: + assert format == "csr" + rval = scipy_sparse.csr_matrix( + (data, indices, indptr), _shape, copy=False + ) + out[0] = rval + + def grad(self, inputs, gout): + (x, y) = inputs + (gz,) = gout + assert psb._is_sparse_variable(gz) + assert psb._is_sparse_variable(x) + + rval = [true_dot(gz, y.T), true_dot(x.T, gz)] + if psb._is_dense_variable(y): + if self.grad_preserves_dense: + rval[1] = psb.dense_from_sparse(rval[1]) + return rval + + def infer_shape(self, fgraph, node, shapes): + return [(shapes[0][0], shapes[1][1])] + + +def true_dot(x, y, grad_preserves_dense=True): + """ + Operation for efficiently calculating the dot product when + one or all operands are sparse. Supported formats are CSC and CSR. + The output of the operation is sparse. + + Parameters + ---------- + x + Sparse matrix. + y + Sparse matrix or 2d tensor variable. + grad_preserves_dense : bool + If True (default), makes the grad of dense inputs dense. + Otherwise the grad is always sparse. + + Returns + ------- + The dot product `x`.`y` in a sparse format. + + Notex + ----- + The grad implemented is regular, i.e. not structured. + + """ + # TODO + # Maybe the triple-transposition formulation + # (when x is dense) is slow. See if there is a + # direct way to do this. + + if hasattr(x, "getnnz"): + x = psb.as_sparse_variable(x) + assert x.format in ("csr", "csc") + if hasattr(y, "getnnz"): + y = psb.as_sparse_variable(y) + assert y.format in ("csr", "csc") + + x_is_sparse_variable = psb._is_sparse_variable(x) + y_is_sparse_variable = psb._is_sparse_variable(y) + + if not x_is_sparse_variable and not y_is_sparse_variable: + raise TypeError() + if x_is_sparse_variable: + return TrueDot(grad_preserves_dense)(x, y) + else: + assert y_is_sparse_variable + return psb.transpose(TrueDot(grad_preserves_dense)(y.T, x.T)) + + +class StructuredDot(Op): + __props__ = () + + def make_node(self, a, b): + a = psb.as_sparse_variable(a) + assert a.format in ("csr", "csc", "bsr") + + if not psb._is_sparse_variable(a): + raise TypeError( + "First argument must be of type SparseVariable or SparseConstant" + ) + dtype_out = ps.upcast(a.type.dtype, b.type.dtype) + if b.type.ndim != 2: + raise NotImplementedError("non-matrix b") + + if psb._is_sparse_variable(b): + return Apply( + self, [a, b], [psb.SparseTensorType(a.type.format, dtype_out)()] + ) + else: + return Apply( + self, + [a, b], + [ + tensor( + dtype=dtype_out, + shape=(None, 1 if b.type.shape[1] == 1 else None), + ) + ], + ) + + def perform(self, node, inputs, outputs): + (a, b) = inputs + (out,) = outputs + if a.shape[1] != b.shape[0]: + raise ValueError( + "shape mismatch in StructuredDot.perform", (a.shape, b.shape) + ) + + variable = a * b + if isinstance(node.outputs[0].type, psb.SparseTensorType): + assert psb._is_sparse(variable) + out[0] = variable + return + + assert psb._is_dense(variable) # scipy 0.7 automatically converts to dense + + # dot of an NxM sparse matrix, with a Mx1 dense matrix, returns vector + # not matrix + if variable.ndim == 1: + variable = np.expand_dims(variable, 1) + elif variable.ndim != 2: + raise Exception("Output of structured dot should be a matrix (ndim=2)") + + assert variable.ndim == 2 + + if variable.shape != (a.shape[0], b.shape[1]): + raise Exception( + f"a.shape={a.shape}, b.shape={b.shape}, variable.shape={variable.shape}?" + ) + + # The cast is needed as otherwise we hit the bug mentioned into + # _asarray function documentation. + out[0] = np.asarray(variable, str(variable.dtype)) + + def grad(self, inputs, gout): + # a is sparse, b is dense, g_out is dense + # ga = g_out x b.T + # gb = a.T x g_out + (a, b) = inputs + (g_out,) = gout + return [structured_dot_grad(a, b, g_out), structured_dot(a.T, g_out)] + + def infer_shape(self, fgraph, node, shapes): + return [(shapes[0][0], shapes[1][1])] + + +_structured_dot = StructuredDot() + + +def structured_dot(x, y): + """ + Structured Dot is like dot, except that only the gradient wrt non-zero elements of the sparse matrix + `a` are calculated and propagated. + + The output is presumed to be a dense matrix, and is represented by a TensorType instance. + + Parameters + ---------- + a + A sparse matrix. + b + A sparse or dense matrix. + + Returns + ------- + A sparse matrix + The dot product of `a` and `b`. + + Notes + ----- + The grad implemented is structured. + + """ + + # @todo: Maybe the triple-transposition formulation (when x is dense) + # is slow. See if there is a direct way to do this. + # (JB 20090528: Transposing tensors and sparse matrices is constant-time, + # inplace, and fast.) + + if hasattr(x, "getnnz"): + x = psb.as_sparse_variable(x) + assert x.format in ("csr", "csc") + if hasattr(y, "getnnz"): + y = psb.as_sparse_variable(y) + assert y.format in ("csr", "csc") + + x_is_sparse_variable = psb._is_sparse_variable(x) + y_is_sparse_variable = psb._is_sparse_variable(y) + if not x_is_sparse_variable and not y_is_sparse_variable: + raise TypeError("structured_dot requires at least one sparse argument") + + if x_is_sparse_variable: + return _structured_dot(x, y) + else: + assert y_is_sparse_variable + return _structured_dot(y.T, x.T).T + + +class StructuredDotGradCSC(COp): + # Op that produces the grad of StructuredDot. + + # :param a_indices: Matrix indices + # :param a_indptr: Matrix indptr + # :param b: Right operand + # :param g_ab: Accumulated gradient. + + # :return: The grad of `a`.`b` for `a` accumulated + # with g_ab. + + # :note: The grad implemented is structured. + # :note: a_* are the corresponding properties of a sparse + # matrix in csc format. + __props__ = () + + def make_node(self, a_indices, a_indptr, b, g_ab): + return Apply( + self, + [a_indices, a_indptr, b, g_ab], + [tensor(dtype=g_ab.dtype, shape=(None,))], + ) + + def perform(self, node, inputs, outputs): + (a_indices, a_indptr, b, g_ab) = inputs + (out,) = outputs + g_a_data = np.zeros(a_indices.shape, dtype=g_ab.dtype) + for j in range(len(a_indptr) - 1): + ind0 = a_indptr[j] + ind1 = a_indptr[j + 1] + for i_idx in range(ind0, ind1): + i = a_indices[i_idx] + # Depending on the type of g_ab and b (sparse or dense), + # the following dot product can result in a scalar or + # a (1, 1) sparse matrix. + dot_val = np.dot(g_ab[i], b[j].T) + if isinstance(dot_val, scipy_sparse.spmatrix): + dot_val = dot_val[0, 0] + g_a_data[i_idx] = dot_val + out[0] = g_a_data + + def c_code_cache_version(self): + return (2,) + + def c_code(self, node, name, inputs, outputs, sub): + (_indices, _indptr, _d, _g) = inputs + (_zout,) = outputs + if node.inputs[2].type.dtype in ("complex64", "complex128"): + raise NotImplementedError("Complex types are not supported for b") + if node.inputs[3].type.dtype in ("complex64", "complex128"): + raise NotImplementedError("Complex types are not supported for g_ab") + + fail = sub["fail"] + return f""" + if (PyArray_NDIM({_d}) != 2) {{PyErr_SetString(PyExc_NotImplementedError, "rank(d) != 2"); {fail};}} + if (PyArray_NDIM({_g}) != 2) {{PyErr_SetString(PyExc_NotImplementedError, "rank(g) != 2"); {fail};}} + if (PyArray_NDIM({_indices}) != 1) {{PyErr_SetString(PyExc_NotImplementedError, "rank(indices) != 1"); {fail};}} + if (PyArray_NDIM({_indptr}) != 1) {{PyErr_SetString(PyExc_NotImplementedError, "rank(indptr) != 1"); {fail};}} + + if( PyArray_TYPE({_indices}) != NPY_INT32) {{ + PyErr_SetString(PyExc_NotImplementedError, "C"); {fail};}} + + if( PyArray_TYPE({_indptr}) != NPY_INT32) + {{PyErr_SetString(PyExc_NotImplementedError, "D"); {fail};}} + + if( PyArray_DIMS({_d})[1] != PyArray_DIMS({_g})[1]) + {{PyErr_SetString(PyExc_NotImplementedError, "d and g have different numbers of columns"); {fail};}} + + if (!{_zout} + || (PyArray_DIMS({_zout})[0] != PyArray_DIMS({_indices})[0])) + {{ + Py_XDECREF({_zout}); + {_zout} = (PyArrayObject*) PyArray_SimpleNew(1, PyArray_DIMS({_indices}), PyArray_TYPE({_g})); + }} + + {{ //makes it compile even though labels jump over variable definitions. + npy_intp nnz = PyArray_DIMS({_indices})[0]; + npy_intp N = PyArray_DIMS({_indptr})[0]-1; //TODO: error checking with this + + npy_intp Sindices = PyArray_STRIDES({_indices})[0]/PyArray_ITEMSIZE({_indices}); + npy_intp Sindptr = PyArray_STRIDES({_indptr})[0]/PyArray_ITEMSIZE({_indptr}); + + const npy_intp Sd1 = PyArray_STRIDES({_d})[1]/PyArray_ITEMSIZE({_d}); + const npy_intp Sg1 = PyArray_STRIDES({_g})[1]/PyArray_ITEMSIZE({_g}); + + const npy_intp K = PyArray_DIMS({_d})[1]; + + const npy_int32 * __restrict__ indptr = (npy_int32 *)PyArray_DATA({_indptr}); + const npy_int32 * __restrict__ indices = (npy_int32 *)PyArray_DATA({_indices}); + + // loop over columns + for (npy_int32 j = 0; j < N; ++j) + {{ + // extract j-th row of dense matrix + const dtype_{_d}* __restrict__ d_row = (dtype_{_d}*)(PyArray_BYTES({_d}) + PyArray_STRIDES({_d})[0] * j); + if(j >= PyArray_DIMS({_d})[0]) {{PyErr_SetString(PyExc_NotImplementedError, "G"); {fail};}} + + // for each non-null value in the sparse column + for (npy_int32 i_idx = indptr[j * Sindptr]; i_idx < indptr[(j+1) * Sindptr]; ++i_idx) + {{ + // extract row index of non-null value + npy_int32 i = indices[i_idx * Sindices]; + + // extract corresponding row in gradient + const dtype_{_g}* __restrict__ g_row = (dtype_{_g}*)(PyArray_BYTES({_g}) + PyArray_STRIDES({_g})[0] * i); + double ip = 0.0; + + // make sure that row index is not bigger than actual number of rows + // Note: wouldn't the above operation fail if that were the case ? + // when would this ever be true anyway ? + if (i >= PyArray_DIMS({_g})[0]) + {{PyErr_SetString(PyExc_NotImplementedError, "H"); {fail};}} + + // perform dot product of dense and sparse rows + for(int k = 0; k < K; ++k) + {{ + ip += d_row[k * Sd1] * g_row[k*Sg1]; + }} + + // write resulting gradient to sparse output + ((dtype_{_zout}* __restrict__)(PyArray_BYTES({_zout}) + i_idx * PyArray_STRIDES({_zout})[0]))[0] = ip; + }} + }} + }} + + """ + + def infer_shape(self, fgraph, node, shapes): + return [shapes[0]] + + +sdg_csc = StructuredDotGradCSC() + + +class StructuredDotGradCSR(COp): + # Op that produces the grad of StructuredDot. + + # :param a_indices: Matrix indices + # :param a_indptr: Matrix indptr + # :param b: Right operand + # :param g_ab: Accumulated gradient. + + # :return: The grad of `a`.`b` for `a` accumulated + # with g_ab. + + # :note: The grad implemented is structured. + # :note: a_* are the corresponding properties of a sparse + # matrix in csr format. + __props__ = () + + def make_node(self, a_indices, a_indptr, b, g_ab): + return Apply( + self, [a_indices, a_indptr, b, g_ab], [tensor(dtype=b.dtype, shape=(None,))] + ) + + def perform(self, node, inputs, outputs): + (a_indices, a_indptr, b, g_ab) = inputs + (out,) = outputs + g_a_data = np.zeros(a_indices.shape, dtype=g_ab.dtype) + for i in range(len(a_indptr) - 1): # loop over rows + ind0 = a_indptr[i] + ind1 = a_indptr[i + 1] + # loop over values in that row (columns) + for j_idx in range(ind0, ind1): + j = a_indices[j_idx] + # grad is dot product of i-th row of gradient with j-th row of b + # Depending on the type of g_ab and b (sparse or dense), + # the following dot product can result in a scalar or + # a (1, 1) sparse matrix. + dot_val = np.dot(g_ab[i], b[j].T) + if isinstance(dot_val, scipy_sparse.spmatrix): + dot_val = dot_val[0, 0] + g_a_data[j_idx] = dot_val + out[0] = g_a_data + + def c_code_cache_version(self): + return (2,) + + def c_code(self, node, name, inputs, outputs, sub): + (_indices, _indptr, _d, _g) = inputs + (_zout,) = outputs + if node.inputs[2].type.dtype in ("complex64", "complex128"): + raise NotImplementedError("Complex types are not supported for b") + if node.inputs[3].type.dtype in ("complex64", "complex128"): + raise NotImplementedError("Complex types are not supported for g_ab") + + fail = sub["fail"] + return f""" + if (PyArray_NDIM({_d}) != 2) {{PyErr_SetString(PyExc_NotImplementedError, "rank(d) != 2"); {fail};}} + if (PyArray_NDIM({_g}) != 2) {{PyErr_SetString(PyExc_NotImplementedError, "rank(g) != 2"); {fail};}} + if (PyArray_NDIM({_indices}) != 1) {{PyErr_SetString(PyExc_NotImplementedError, "rank(indices) != 1"); {fail};}} + if (PyArray_NDIM({_indptr}) != 1) {{PyErr_SetString(PyExc_NotImplementedError, "rank(indptr) != 1"); {fail};}} + + if( PyArray_TYPE({_indices}) != NPY_INT32) {{ + PyErr_SetString(PyExc_NotImplementedError, "C"); {fail};}} + + if( PyArray_TYPE({_indptr}) != NPY_INT32) + {{PyErr_SetString(PyExc_NotImplementedError, "D"); {fail};}} + + if( PyArray_DIMS({_d})[1] != PyArray_DIMS({_g})[1]) + {{PyErr_SetString(PyExc_NotImplementedError, "d and g have different numbers of columns"); {fail};}} + + if (!{_zout} + || (PyArray_DIMS({_zout})[0] != PyArray_DIMS({_indices})[0])) + {{ + Py_XDECREF({_zout}); + {_zout} = (PyArrayObject*) PyArray_SimpleNew(1, PyArray_DIMS({_indices}), PyArray_TYPE({_g})); + }} + + {{ //makes it compile even though labels jump over variable definitions. + npy_intp nnz = PyArray_DIMS({_indices})[0]; + // extract number of rows + npy_intp N = PyArray_DIMS({_indptr})[0]-1; //TODO: error checking with this + + npy_intp Sindices = PyArray_STRIDES({_indices})[0]/PyArray_ITEMSIZE({_indices}); + npy_intp Sindptr = PyArray_STRIDES({_indptr})[0]/PyArray_ITEMSIZE({_indptr}); + + const npy_intp Sd1 = PyArray_STRIDES({_d})[1]/PyArray_ITEMSIZE({_d}); + const npy_intp Sg1 = PyArray_STRIDES({_g})[1]/PyArray_ITEMSIZE({_g}); + + const npy_intp K = PyArray_DIMS({_d})[1]; + + const npy_int32 * __restrict__ indptr = (npy_int32 *)PyArray_DATA({_indptr}); + const npy_int32 * __restrict__ indices = (npy_int32 *)PyArray_DATA({_indices}); + + // loop over columns of sparse matrix + for (npy_int32 i = 0; i < N; ++i) + {{ + // for each non-null value in the sparse row + for (npy_int32 j_idx = indptr[i * Sindptr]; j_idx < indptr[(i+1) * Sindptr]; ++j_idx) + {{ + // extract column index of non-null value + npy_int32 j = indices[j_idx * Sindices]; + + // extract j-th row of dense matrix + const dtype_{_d}* __restrict__ d_row = (dtype_{_d}*)(PyArray_BYTES({_d}) + PyArray_STRIDES({_d})[0] * j); + if(j >= PyArray_DIMS({_d})[0]) {{PyErr_SetString(PyExc_NotImplementedError, "G"); {fail};}} + + // extract corresponding row in gradient + const dtype_{_g}* __restrict__ g_row = (dtype_{_g}*)(PyArray_BYTES({_g}) + PyArray_STRIDES({_g})[0] * i); + double ip = 0.0; + + // make sure that row index is not bigger than actual number of rows + // Note: wouldn't the above operation fail if that were the case ? + // when would this ever be true anyway ? + if (i >= PyArray_DIMS({_g})[0]) + {{PyErr_SetString(PyExc_NotImplementedError, "H"); {fail};}} + + // perform dot product of dense and sparse rows + for(int k = 0; k < K; ++k) + {{ + ip += d_row[k * Sd1] * g_row[k*Sg1]; + }} + + // write resulting gradient to sparse output + ((dtype_{_zout}* __restrict__)(PyArray_BYTES({_zout}) + j_idx * PyArray_STRIDES({_zout})[0]))[0] = ip; + }} + }} + }} + + """ + + def infer_shape(self, fgraph, node, shapes): + return [shapes[0]] + + +sdg_csr = StructuredDotGradCSR() + + +def structured_dot_grad(sparse_A, dense_B, ga): + if sparse_A.type.format in ("csc", "csr"): + if sparse_A.type.format == "csc": + sdgcsx = sdg_csc + CSx = psb.CSC + else: + sdgcsx = sdg_csr + CSx = psb.CSR + + g_A_data = sdgcsx( + psb.csm_indices(sparse_A), psb.csm_indptr(sparse_A), dense_B, ga + ) + return CSx( + g_A_data, + psb.csm_indices(sparse_A), + psb.csm_indptr(sparse_A), + psb.csm_shape(sparse_A), + ) + else: + raise NotImplementedError() + + +class SamplingDot(Op): + """Compute the dot product ``dot(x, y.T) = z`` for only a subset of `z`. + + This is equivalent to ``p * (x . y.T)`` where ``*`` is the element-wise + product, ``x`` and ``y`` operands of the dot product and ``p`` is a matrix that + contains 1 when the corresponding element of ``z`` should be calculated + and ``0`` when it shouldn't. Note that `SamplingDot` has a different interface + than `dot` because it requires ``x`` to be a ``m x k`` matrix while + ``y`` is a ``n x k`` matrix instead of the usual ``k x n`` matrix. + + Notes + ----- + It will work if the pattern is not binary value, but if the + pattern doesn't have a high sparsity proportion it will be slower + then a more optimized dot followed by a normal elemwise + multiplication. + + The grad implemented is regular, i.e. not structured. + + """ + + __props__ = () + + def make_node(self, x, y, p): + """ + Parameters + ---------- + x + Tensor matrix. + y + Tensor matrix. + p + Sparse matrix in csr format. + + """ + x = ptb.as_tensor_variable(x) + y = ptb.as_tensor_variable(y) + p = psb.as_sparse_variable(p) + assert p.format in ("csr", "csc") + + if not psb._is_sparse_variable(p): + raise TypeError(p) + + # TODO: use it. + # dtype_out = ps.upcast(x.type.dtype, y.type.dtype, p.type.dtype) + + return Apply(self, [x, y, p], [p.type()]) + + def perform(self, node, inputs, outputs): + (x, y, p) = inputs + (out,) = outputs + if psb._is_sparse(x): + raise TypeError(x) + + if psb._is_sparse(y): + raise TypeError(y) + + if not psb._is_sparse(p): + raise TypeError(p) + + out[0] = p.__class__(p.multiply(np.dot(x, y.T))) + + def grad(self, inputs, gout): + (x, y, p) = inputs + (gz,) = gout + rval = [dot(p * gz, y), dot((p * gz).T, x), grad_not_implemented(self, 2, p)] + + return rval + + def infer_shape(self, fgraph, node, ins_shapes): + return [ins_shapes[2]] + + +sampling_dot = SamplingDot() + + +class Dot(Op): + __props__ = () + + def __str__(self): + return "Sparse" + self.__class__.__name__ + + def infer_shape(self, fgraph, node, shapes): + xshp, yshp = shapes + x, y = node.inputs + if x.ndim == 2 and y.ndim == 2: + return [(xshp[0], yshp[1])] + if x.ndim == 1 and y.ndim == 2: + return [(yshp[1],)] + if x.ndim == 2 and y.ndim == 1: + return [(xshp[0],)] + if x.ndim == 1 and y.ndim == 1: + return [()] + raise NotImplementedError() + + def make_node(self, x, y): + dtype_out = ps.upcast(x.dtype, y.dtype) + + # Sparse dot product should have at least one sparse variable + # as input. If the other one is not sparse, it has to be converted + # into a tensor. + if isinstance(x, scipy_sparse.spmatrix): + x = psb.as_sparse_variable(x) + if isinstance(y, scipy_sparse.spmatrix): + y = psb.as_sparse_variable(y) + + x_is_sparse_var = psb._is_sparse_variable(x) + y_is_sparse_var = psb._is_sparse_variable(y) + + if not x_is_sparse_var and not y_is_sparse_var: + raise TypeError( + "Sparse dot product should have at least one " + "sparse variable as inputs, but the inputs are " + f"{x} ({x.type}) and {y} ({y.type})." + ) + + if x_is_sparse_var: + shape_x = (None,) * x.type.ndim + else: + x = ptb.as_tensor_variable(x) + shape_x = x.type.shape + assert y.format in ("csr", "csc") + if x.ndim not in (1, 2): + raise TypeError( + "Input 0 (0-indexed) must have ndim of " + f"1 or 2, {int(x.type.ndim)} given." + ) + + if y_is_sparse_var: + shape_y = (None,) * y.type.ndim + else: + y = ptb.as_tensor_variable(y) + shape_y = y.type.shape + assert x.format in ("csr", "csc") + if y.ndim not in (1, 2): + raise TypeError( + "Input 1 (1-indexed) must have ndim of " + f"1 or 2, {int(y.type.ndim)} given." + ) + + if len(shape_y) == 2: + shape_out = shape_x[:-1] + shape_y[1:] + elif len(shape_y) == 1: + shape_out = shape_x[:-1] + + return Apply(self, [x, y], [tensor(dtype=dtype_out, shape=shape_out)]) + + def perform(self, node, inputs, out): + x, y = inputs + out = out[0] + x_is_sparse = psb._is_sparse(x) + y_is_sparse = psb._is_sparse(y) + + if not x_is_sparse and not y_is_sparse: + raise TypeError(x) + + rval = x * y + + if x_is_sparse and y_is_sparse: + rval = rval.toarray() + + out[0] = np.asarray(rval, dtype=node.outputs[0].dtype) + + def grad(self, inputs, gout): + (x, y) = inputs + (gz,) = gout + assert psb._is_sparse_variable(x) or psb._is_sparse_variable(y) + rval = [] + + if psb._is_dense_variable(y): + rval.append(ptm.dot(gz, y.T)) + else: + rval.append(dot(gz, y.T)) + if psb._is_dense_variable(x): + rval.append(ptm.dot(x.T, gz)) + else: + rval.append(dot(x.T, gz)) + + return rval + + +_dot = Dot() + + +def dot(x, y): + """Efficiently compute the dot product when one or all operands are sparse. + + Supported formats are CSC and CSR. The output of the operation is dense. + + Parameters + ---------- + x + Sparse or dense matrix variable. + y + Sparse or dense matrix variable. + + Returns + ------- + The dot product ``x @ y`` in a dense format. + + Notes + ----- + The grad implemented is regular, i.e. not structured. + + At least one of `x` or `y` must be a sparse matrix. + + When the operation has the form ``dot(csr_matrix, dense)`` + the gradient of this operation can be performed inplace + by `UsmmCscDense`. This leads to significant speed-ups. + + """ + + if hasattr(x, "getnnz"): + x = psb.as_sparse_variable(x) + if hasattr(y, "getnnz"): + y = psb.as_sparse_variable(y) + + x_is_sparse_variable = psb._is_sparse_variable(x) + y_is_sparse_variable = psb._is_sparse_variable(y) + + if not x_is_sparse_variable and not y_is_sparse_variable: + raise TypeError() + + return _dot(x, y) + + +class Usmm(Op): + """Computes the dense matrix resulting from ``alpha * x @ y + z``. + + Notes + ----- + At least one of `x` or `y` must be a sparse matrix. + + """ + + __props__ = () + + def __str__(self): + return "Usmm{no_inplace}" + + def make_node(self, alpha, x, y, z): + """ + + Parameters + ---------- + alpha + A scalar. + x + Matrix variable. + y + Matrix variable. + z + Dense matrix. + + """ + if not psb._is_sparse_variable(x) and not psb._is_sparse_variable(y): + # If x and y are tensor, we don't want to use this class + # We should use Dot22 and Gemm in that case. + raise TypeError(x) + + dtype_out = ps.upcast( + alpha.type.dtype, x.type.dtype, y.type.dtype, z.type.dtype + ) + alpha = ptb.as_tensor_variable(alpha) + z = ptb.as_tensor_variable(z) + + assert z.type.ndim == 2 + assert alpha.type.shape == (1,) * alpha.type.ndim + if not psb._is_sparse_variable(x): + x = ptb.as_tensor_variable(x) + assert y.format in ("csr", "csc") + assert x.type.ndim == 2 + if not psb._is_sparse_variable(y): + y = ptb.as_tensor_variable(y) + assert x.format in ("csr", "csc") + assert y.type.ndim == 2 + + return Apply( + self, + [alpha, x, y, z], + [tensor(dtype=dtype_out, shape=(None, None))], + ) + + def perform(self, node, inputs, outputs): + (alpha, x, y, z) = inputs + (out,) = outputs + x_is_sparse = psb._is_sparse(x) + y_is_sparse = psb._is_sparse(y) + + if not x_is_sparse and not y_is_sparse: + raise TypeError(x) + + rval = x * y + if isinstance(rval, scipy_sparse.spmatrix): + rval = rval.toarray() + if rval.dtype == alpha.dtype: + rval *= alpha # Faster because operation is inplace + else: + rval = rval * alpha + if rval.dtype == z.dtype: + rval += z # Faster because operation is inplace + else: + rval = rval + z + + out[0] = rval + + +usmm = Usmm() diff --git a/pytensor/sparse/rewriting.py b/pytensor/sparse/rewriting.py index 041688fdc6..d992635298 100644 --- a/pytensor/sparse/rewriting.py +++ b/pytensor/sparse/rewriting.py @@ -3,6 +3,8 @@ import pytensor import pytensor.scalar as ps +import pytensor.sparse.basic as sparse +import pytensor.sparse.math as spm from pytensor.configdefaults import config from pytensor.graph.basic import Apply from pytensor.graph.rewriting.basic import ( @@ -11,17 +13,8 @@ node_rewriter, ) from pytensor.link.c.op import COp, _NoPythonCOp -from pytensor.sparse import basic as sparse -from pytensor.sparse.basic import ( - CSC, - CSR, - csm_data, - csm_grad, - csm_indices, - csm_indptr, - csm_properties, - usmm, -) +from pytensor.sparse.basic import csm_properties +from pytensor.sparse.math import usmm from pytensor.tensor import blas from pytensor.tensor.basic import as_tensor_variable, cast from pytensor.tensor.math import mul, neg, sub @@ -34,25 +27,22 @@ _is_dense = sparse._is_dense -@node_rewriter([csm_properties]) +@register_specialize +@node_rewriter([sparse.csm_properties]) def local_csm_properties_csm(fgraph, node): """ If we find csm_properties(CSM(*args)), then we can replace that with the *args directly. """ - if node.op == csm_properties: + if node.op == sparse.csm_properties: (csm,) = node.inputs - if csm.owner and (csm.owner.op == CSC or csm.owner.op == CSR): + if csm.owner and (csm.owner.op == sparse.CSC or csm.owner.op == sparse.CSR): return csm.owner.inputs return False -register_specialize(local_csm_properties_csm) - - -# This is tested in tests/test_basic.py:test_remove0 @node_rewriter([sparse.Remove0]) def local_inplace_remove0(fgraph, node): """Rewrite to insert inplace versions of `Remove0`.""" @@ -120,7 +110,11 @@ def make_node(self, x, y): if self.inplace: assert out_dtype == y.dtype - indices, indptr, data = csm_indices(x), csm_indptr(x), csm_data(x) + indices, indptr, data = ( + sparse.csm_indices(x), + sparse.csm_indptr(x), + sparse.csm_data(x), + ) # We either use CSC or CSR depending on the format of input assert self.format == x.type.format # The magic number two here arises because L{scipy.sparse} @@ -189,10 +183,10 @@ def c_code_cache_version(self): return (3,) -@node_rewriter([sparse.AddSD]) +@node_rewriter([spm.AddSD]) def local_inplace_addsd_ccode(fgraph, node): """Rewrite to insert inplace versions of `AddSD`.""" - if isinstance(node.op, sparse.AddSD) and config.cxx: + if isinstance(node.op, spm.AddSD) and config.cxx: out_dtype = ps.upcast(*[inp.type.dtype for inp in node.inputs]) if out_dtype != node.inputs[1].dtype: return @@ -225,13 +219,13 @@ def local_dense_from_sparse_sparse_from_dense(fgraph, node): return inp.owner.inputs -@node_rewriter([sparse.AddSD]) +@node_rewriter([spm.AddSD]) def local_addsd_ccode(fgraph, node): """ Convert AddSD to faster AddSD_ccode. """ - if isinstance(node.op, sparse.AddSD) and config.cxx: + if isinstance(node.op, spm.AddSD) and config.cxx: new_node = AddSD_ccode(format=node.inputs[0].type.format)(*node.inputs) return [new_node] return False @@ -624,9 +618,9 @@ def c_code_cache_version(self): # register a specialization to replace StructuredDot -> StructuredDotCSx # This is tested in tests/test_basic.py:792 -@node_rewriter([sparse._structured_dot]) +@node_rewriter([spm._structured_dot]) def local_structured_dot(fgraph, node): - if node.op == sparse._structured_dot: + if node.op == spm._structured_dot: a, b = node.inputs if a.type.format == "csc": a_val, a_ind, a_ptr, a_shape = csm_properties(a) @@ -918,10 +912,10 @@ def c_code_cache_version(self): all(s == 1 for s in expr.type.shape) and config.blas__ldflags ), }, - (sparse._dot, "x", "y"), + (spm._dot, "x", "y"), ), ), - (usmm, (neg, "alpha"), "x", "y", "z"), + (spm.usmm, (neg, "alpha"), "x", "y", "z"), ) register_specialize(local_usmm, name="local_usmm") @@ -938,7 +932,7 @@ def local_usmm_csc_dense_inplace(fgraph, node): # This is tested in tests/test_basic.py:UsmmTests -@node_rewriter([usmm]) +@node_rewriter([spm.usmm]) def local_usmm_csx(fgraph, node): """ usmm -> usmm_csc_dense @@ -1094,13 +1088,13 @@ def c_code_cache_version(self): csm_grad_c = CSMGradC() -@node_rewriter([csm_grad(None)]) +@node_rewriter([sparse.csm_grad(None)]) def local_csm_grad_c(fgraph, node): """ csm_grad(None) -> csm_grad_c """ - if node.op == csm_grad(None): + if node.op == sparse.csm_grad(None): return [csm_grad_c(*node.inputs)] return False @@ -1386,9 +1380,9 @@ def __str__(self): # register a specialization to replace MulSD -> MulSDCSX -@node_rewriter([sparse.mul_s_d]) +@node_rewriter([spm.mul_s_d]) def local_mul_s_d(fgraph, node): - if node.op == sparse.mul_s_d: + if node.op == spm.mul_s_d: x, y = node.inputs x_is_sparse_variable = _is_sparse_variable(x) @@ -1572,9 +1566,9 @@ def __str__(self): # register a specialization to replace MulSV -> MulSVCSR -@node_rewriter([sparse.mul_s_v]) +@node_rewriter([spm.mul_s_v]) def local_mul_s_v(fgraph, node): - if node.op == sparse.mul_s_v: + if node.op == spm.mul_s_v: x, y = node.inputs x_is_sparse_variable = _is_sparse_variable(x) @@ -1753,9 +1747,9 @@ def __str__(self): # register a specialization to replace # structured_add_s_v -> structured_add_s_v_csr -@node_rewriter([sparse.structured_add_s_v]) +@node_rewriter([spm.structured_add_s_v]) def local_structured_add_s_v(fgraph, node): - if node.op == sparse.structured_add_s_v: + if node.op == spm.structured_add_s_v: x, y = node.inputs x_is_sparse_variable = _is_sparse_variable(x) @@ -2044,12 +2038,12 @@ def c_code(self, node, name, inputs, outputs, sub): # register a specialization to replace SamplingDot -> SamplingDotCsr -@node_rewriter([sparse.sampling_dot]) +@node_rewriter([spm.sampling_dot]) def local_sampling_dot_csr(fgraph, node): if not config.blas__ldflags: # The C implementation of SamplingDotCsr relies on BLAS routines return - if node.op == sparse.sampling_dot: + if node.op == spm.sampling_dot: x, y, p = node.inputs if p.type.format == "csr": p_data, p_ind, p_ptr, p_shape = sparse.csm_properties(p) diff --git a/pytensor/sparse/sharedvar.py b/pytensor/sparse/sharedvar.py index 60b09656be..ce9f9ca6d0 100644 --- a/pytensor/sparse/sharedvar.py +++ b/pytensor/sparse/sharedvar.py @@ -3,7 +3,7 @@ import scipy.sparse from pytensor.compile import shared_constructor -from pytensor.sparse.basic import SparseTensorType, SparseVariable +from pytensor.sparse.variable import SparseTensorType, SparseVariable from pytensor.tensor.sharedvar import TensorSharedVariable diff --git a/pytensor/sparse/variable.py b/pytensor/sparse/variable.py new file mode 100644 index 0000000000..fc462b278b --- /dev/null +++ b/pytensor/sparse/variable.py @@ -0,0 +1,476 @@ +from warnings import warn + +import numpy as np +import scipy.sparse as scipy_sparse + +from pytensor.sparse.basic import ( + cast, + csm_data, + dense_from_sparse, + get_item_2d, + get_item_2lists, + get_item_list, + get_item_scalar, + neg, + sp_ones_like, + sp_zeros_like, + transpose, +) +from pytensor.sparse.math import ( + add, + ge, + gt, + le, + lt, + mul, + sp_sum, + structured_abs, + structured_arccos, + structured_arccosh, + structured_arcsin, + structured_arcsinh, + structured_arctan, + structured_arctanh, + structured_ceil, + structured_conjugate, + structured_cos, + structured_cosh, + structured_deg2rad, + structured_dot, + structured_exp, + structured_exp2, + structured_expm1, + structured_floor, + structured_log, + structured_log1p, + structured_log2, + structured_log10, + structured_rad2deg, + structured_sin, + structured_sinh, + structured_sqrt, + structured_tan, + structured_tanh, + structured_trunc, + sub, +) +from pytensor.sparse.type import SparseTensorType +from pytensor.sparse.utils import hash_from_sparse +from pytensor.tensor import iscalar +from pytensor.tensor.shape import shape +from pytensor.tensor.variable import ( + TensorConstant, + TensorVariable, + _tensor_py_operators, +) + + +def constant(x, name=None): + if not isinstance(x, scipy_sparse.spmatrix): + raise TypeError("sparse.constant must be called on a scipy.sparse.spmatrix") + try: + return SparseConstant( + SparseTensorType(format=x.format, dtype=x.dtype), x.copy(), name=name + ) + except TypeError: + raise TypeError(f"Could not convert {x} to SparseTensorType", type(x)) + + +def override_dense(method): + dense_method = getattr(_tensor_py_operators, method.__name__) + + def to_dense(self, *args, **kwargs): + self = self.toarray() + new_args = [ + arg.toarray() + if hasattr(arg, "type") and isinstance(arg.type, SparseTensorType) + else arg + for arg in args + ] + warn( + f"Method {method} is not implemented for sparse variables. The variable will be converted to dense." + ) + return dense_method(self, *new_args, **kwargs) + + to_dense._is_dense_override = True + return to_dense + + +class _sparse_py_operators: + T = property( + lambda self: transpose(self), doc="Return aliased transpose of self (read-only)" + ) + + def astype(self, dtype): + return cast(self, dtype) + + def __neg__(self): + return neg(self) + + def __add__(left, right): + return add(left, right) + + def __radd__(right, left): + return add(left, right) + + def __sub__(left, right): + return sub(left, right) + + def __rsub__(right, left): + return sub(left, right) + + def __mul__(left, right): + return mul(left, right) + + def __rmul__(left, right): + return mul(left, right) + + # comparison operators + + def __lt__(self, other): + return lt(self, other) + + def __le__(self, other): + return le(self, other) + + def __gt__(self, other): + return gt(self, other) + + def __ge__(self, other): + return ge(self, other) + + def __dot__(left, right): + return structured_dot(left, right) + + def __rdot__(right, left): + return structured_dot(left, right) + + def sum(self, axis=None, sparse_grad=False): + return sp_sum(self, axis=axis, sparse_grad=sparse_grad) + + dot = __dot__ + + def toarray(self): + return dense_from_sparse(self) + + @property + def shape(self): + # TODO: The plan is that the ShapeFeature in ptb.opt will do shape + # propagation and remove the dense_from_sparse from the graph. This + # will *NOT* actually expand your sparse matrix just to get the shape. + return shape(dense_from_sparse(self)) + + ndim = property(lambda self: self.type.ndim) + dtype = property(lambda self: self.type.dtype) + + # Note that the `size` attribute of sparse matrices behaves differently + # from dense matrices: it is the number of elements stored in the matrix + # rather than the total number of elements that may be stored. Note also + # that stored zeros *do* count in the size. + size = property(lambda self: csm_data(self).size) + + def zeros_like(model): + return sp_zeros_like(model) + + def ones_like(self): + return sp_ones_like(self) + + def __getitem__(self, args): + if not isinstance(args, tuple): + args = (args,) + + if len(args) == 2: + scalar_arg_1 = ( + np.isscalar(args[0]) or getattr(args[0], "type", None) == iscalar + ) + scalar_arg_2 = ( + np.isscalar(args[1]) or getattr(args[1], "type", None) == iscalar + ) + if scalar_arg_1 and scalar_arg_2: + ret = get_item_scalar(self, args) + elif isinstance(args[0], list): + ret = get_item_2lists(self, args[0], args[1]) + else: + ret = get_item_2d(self, args) + elif isinstance(args[0], list): + ret = get_item_list(self, args[0]) + else: + ret = get_item_2d(self, args) + return ret + + def conj(self): + return structured_conjugate(self) + + def __abs__(self): + return structured_abs(self) + + def __ceil__(self): + return structured_ceil(self) + + def __floor__(self): + return structured_floor(self) + + def __trunc__(self): + return structured_trunc(self) + + def transpose(self): + return transpose(self) + + @override_dense + def any(self, axis=None, keepdims=False): + raise NotImplementedError + + @override_dense + def all(self, axis=None, keepdims=False): + raise NotImplementedError + + @override_dense + def nonzero(self): + raise NotImplementedError + + @override_dense + def nonzero_values(self): + raise NotImplementedError + + @override_dense + def flatten(self, ndim=1): + raise NotImplementedError + + @override_dense + def ravel(self): + raise NotImplementedError + + def arccos(self): + return structured_arccos(self) + + def arcsin(self): + return structured_arcsin(self) + + def arctan(self): + return structured_arctan(self) + + def arccosh(self): + return structured_arccosh(self) + + def arcsinh(self): + return structured_arcsinh(self) + + def arctanh(self): + return structured_arctanh(self) + + def ceil(self): + return structured_ceil(self) + + def cos(self): + return structured_cos(self) + + def cosh(self): + return structured_cosh(self) + + def deg2rad(self): + return structured_deg2rad(self) + + def exp(self): + return structured_exp(self) + + def exp2(self): + return structured_exp2(self) + + def expm1(self): + return structured_expm1(self) + + def floor(self): + return structured_floor(self) + + def log(self): + return structured_log(self) + + def log10(self): + return structured_log10(self) + + def log1p(self): + return structured_log1p(self) + + def log2(self): + return structured_log2(self) + + def rad2deg(self): + return structured_rad2deg(self) + + def sin(self): + return structured_sin(self) + + def sinh(self): + return structured_sinh(self) + + def sqrt(self): + return structured_sqrt(self) + + def tan(self): + return structured_tan(self) + + def tanh(self): + return structured_tanh(self) + + @override_dense + def copy(self, name=None): + raise NotImplementedError + + @override_dense + def prod(self, axis=None, dtype=None, keepdims=False, acc_dtype=None): + raise NotImplementedError + + @override_dense + def mean(self, axis=None, dtype=None, keepdims=False, acc_dtype=None): + raise NotImplementedError + + @override_dense + def var(self, axis=None, ddof=0, keepdims=False, corrected=False): + raise NotImplementedError + + @override_dense + def std(self, axis=None, ddof=0, keepdims=False, corrected=False): + raise NotImplementedError + + @override_dense + def min(self, axis=None, keepdims=False): + raise NotImplementedError + + @override_dense + def max(self, axis=None, keepdims=False): + raise NotImplementedError + + @override_dense + def argmin(self, axis=None, keepdims=False): + raise NotImplementedError + + @override_dense + def argmax(self, axis=None, keepdims=False): + raise NotImplementedError + + @override_dense + def argsort(self, axis=-1, kind="quicksort", order=None): + raise NotImplementedError + + @override_dense + def round(self, mode=None): + raise NotImplementedError + + @override_dense + def trace(self): + raise NotImplementedError + + @override_dense + def cumsum(self, axis=None): + raise NotImplementedError + + @override_dense + def cumprod(self, axis=None): + raise NotImplementedError + + @override_dense + def ptp(self, axis=None): + raise NotImplementedError + + @override_dense + def squeeze(self, axis=None): + raise NotImplementedError + + @override_dense + def diagonal(self, offset=0, axis1=0, axis2=1): + raise NotImplementedError + + @override_dense + def __and__(self, other): + raise NotImplementedError + + @override_dense + def __or__(self, other): + raise NotImplementedError + + @override_dense + def __xor__(self, other): + raise NotImplementedError + + @override_dense + def __pow__(self, other): + raise NotImplementedError + + @override_dense + def __mod__(self, other): + raise NotImplementedError + + @override_dense + def __divmod__(self, other): + raise NotImplementedError + + @override_dense + def __truediv__(self, other): + raise NotImplementedError + + @override_dense + def __floordiv__(self, other): + raise NotImplementedError + + @override_dense + def reshape(self, shape, *, ndim=None): + raise NotImplementedError + + @override_dense + def dimshuffle(self, *pattern): + raise NotImplementedError + + +class SparseVariable(_sparse_py_operators, TensorVariable): + format = property(lambda self: self.type.format) + + def __str__(self): + return f"{self.__class__.__name__}{{{self.format},{self.dtype}}}" + + def __repr__(self): + return str(self) + + +class SparseConstantSignature(tuple): + def __eq__(self, other): + (a, b), (x, y) = self, other + return ( + a == x + and (b.dtype == y.dtype) + and (type(b) is type(y)) + and (b.shape == y.shape) + and (abs(b - y).sum() < 1e-6 * b.nnz) + ) + + def __ne__(self, other): + return not self == other + + def __hash__(self): + (a, b) = self + return hash(type(self)) ^ hash(a) ^ hash(type(b)) + + def pytensor_hash(self): + (_, d) = self + return hash_from_sparse(d) + + +class SparseConstant(SparseVariable, TensorConstant): + format = property(lambda self: self.type.format) + + def signature(self): + assert self.data is not None + return SparseConstantSignature((self.type, self.data)) + + def __str__(self): + return f"{self.__class__.__name__}{{{self.format},{self.dtype},shape={self.data.shape},nnz={self.data.nnz}}}" + + def __repr__(self): + return str(self) + + @property + def unique_value(self): + return None + + +SparseTensorType.variable_type = SparseVariable +SparseTensorType.constant_type = SparseConstant diff --git a/tests/sparse/test_basic.py b/tests/sparse/test_basic.py index 86d703740e..6f14652471 100644 --- a/tests/sparse/test_basic.py +++ b/tests/sparse/test_basic.py @@ -1,57 +1,45 @@ -import time -from itertools import product - import numpy as np import pytest +import scipy.sparse as scipy_sparse from packaging import version +from scipy import __version__ as scipy_version import pytensor +import pytensor.sparse.math import pytensor.tensor as pt from pytensor import sparse from pytensor.compile.function import function -from pytensor.compile.io import In, Out +from pytensor.compile.io import In from pytensor.configdefaults import config from pytensor.gradient import GradientError -from pytensor.graph.basic import Apply, Constant +from pytensor.graph.basic import Apply from pytensor.graph.op import Op -from pytensor.graph.traversal import applys_between -from pytensor.sparse import ( +from pytensor.sparse.basic import ( CSC, CSM, CSR, - AddSD, - AddSS, - AddSSData, Cast, ConstructSparseFromList, CSMGrad, CSMProperties, DenseFromSparse, Diag, - Dot, EnsureSortedIndices, GetItemScalar, HStack, - MulSD, - MulSS, Neg, Remove0, - SamplingDot, SparseFromDense, SparseTensorType, SquareDiagonal, - StructuredDot, - StructuredDotGradCSC, - StructuredDotGradCSR, Transpose, - TrueDot, - Usmm, VStack, - add, - add_s_s_data, + _is_sparse, + _is_sparse_variable, + _mtypes, + all_dtypes, as_sparse_or_tensor_variable, as_sparse_variable, - block_diag, cast, clean, construct_sparse_from_list, @@ -61,74 +49,40 @@ dense_from_sparse, diag, ensure_sorted_indices, - ge, - gt, - le, - lt, - mul_s_v, - multiply, - sampling_dot, sp_ones_like, + sparse_formats, square_diagonal, - structured_add, - structured_add_s_v, - structured_dot, - structured_maximum, - structured_minimum, transpose, - true_dot, -) -from pytensor.sparse.basic import ( - SparseConstant, - _is_dense_variable, - _is_sparse, - _is_sparse_variable, - _mtypes, ) from pytensor.sparse.rewriting import ( - AddSD_ccode, CSMGradC, - StructuredDotCSC, - UsmmCscDense, ) +from pytensor.sparse.variable import SparseConstant from pytensor.tensor.basic import MakeVector -from pytensor.tensor.elemwise import DimShuffle, Elemwise from pytensor.tensor.math import sum as pt_sum from pytensor.tensor.shape import Shape_i from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, AdvancedSubtensor1, - Subtensor, ) from pytensor.tensor.type import ( TensorType, float_dtypes, - fscalar, iscalar, ivector, lvector, matrix, - scalar, tensor, vector, ) from tests import unittest_tools as utt -from tests.tensor.test_sharedvar import makeSharedTester - - -sp = pytest.importorskip("scipy", minversion="0.7.0") - - -# Probability distributions are currently tested in test_sp2.py -# from pytensor.sparse import ( -# Poisson, poisson, Binomial, Multinomial, multinomial) def as_sparse_format(data, format): if format == "csc": - return sp.sparse.csc_matrix(data) + return scipy_sparse.csc_matrix(data) elif format == "csr": - return sp.sparse.csr_matrix(data) + return scipy_sparse.csr_matrix(data) else: raise NotImplementedError() @@ -147,7 +101,7 @@ def as_ndarray(val): def random_lil(shape, dtype, nnz): - rval = sp.sparse.lil_matrix(shape, dtype=dtype) + rval = scipy_sparse.lil_matrix(shape, dtype=dtype) huge = 2**30 for k in range(nnz): # set non-zeros in random locations (row x, col y) @@ -231,7 +185,7 @@ def _rand(): getattr(pytensor.sparse, format + "_matrix")(dtype=out_dtype) for k in range(n) ] data = [ - getattr(sp.sparse, format + "_matrix")(_rand(), dtype=out_dtype) + getattr(scipy_sparse, format + "_matrix")(_rand(), dtype=out_dtype) for k in range(n) ] if unsorted_indices: @@ -239,11 +193,11 @@ def _rand(): d = data[idx] # these flip the matrix, but it's random anyway if format == "csr": - d = sp.sparse.csr_matrix( + d = scipy_sparse.csr_matrix( (d.data, d.shape[1] - 1 - d.indices, d.indptr), shape=d.shape ) if format == "csc": - d = sp.sparse.csc_matrix( + d = scipy_sparse.csc_matrix( (d.data, d.shape[0] - 1 - d.indices, d.indptr), shape=d.shape ) assert not d.has_sorted_indices @@ -254,7 +208,7 @@ def _rand(): d_idx = np.random.default_rng().integers(data[idx].nnz) data[idx].data[d_idx] = 0 - # numpy 1.5.0 with scipy 0.9.0 have sp.sparse.XXX_matrix return + # numpy 1.5.0 with scipy 0.9.0 have scipy_sparse.XXX_matrix return # typenum 10(ulonglong) instead of 8(uint64) event if they are the same! # PyTensor don't like ulonglong type_num dtype = np.dtype(out_dtype) # Convert into dtype object. @@ -379,19 +333,19 @@ def test_grad_fail(self): with pytest.raises(GradientError): verify_grad_sparse( self.FailOp(structured=False), - [sp.sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], + [scipy_sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], ) with pytest.raises(GradientError): verify_grad_sparse( self.FailOp(structured=True), - [sp.sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], + [scipy_sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], ) class TestTranspose: def test_transpose_csc(self): - spe = sp.sparse.csc_matrix(sp.sparse.eye(5, 3)) + spe = scipy_sparse.csc_matrix(scipy_sparse.eye(5, 3)) a = as_sparse_variable(spe) assert a.data is not spe assert a.data.shape == (5, 3) @@ -405,7 +359,7 @@ def test_transpose_csc(self): assert vta.shape == (3, 5) def test_transpose_csr(self): - a = as_sparse_variable(sp.sparse.csr_matrix(sp.sparse.eye(5, 3))) + a = as_sparse_variable(scipy_sparse.csr_matrix(scipy_sparse.eye(5, 3))) assert a.data.shape == (5, 3) assert a.type.dtype == "float64" assert a.type.format == "csr" @@ -427,7 +381,7 @@ def test_getitem_scalar(self): self._compile_and_check( [x], [x[2, 2]], - [sp.sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], + [scipy_sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], GetItemScalar, ) @@ -437,7 +391,7 @@ def test_csm(self): y = ivector() z = ivector() s = ivector() - call = getattr(sp.sparse, sparsetype + "_matrix") + call = getattr(scipy_sparse, sparsetype + "_matrix") spm = call(random_lil((300, 400), config.floatX, 5)) out = CSM(sparsetype)(x, y, z, s) self._compile_and_check( @@ -450,7 +404,7 @@ def test_csm_grad(self): y = ivector() z = ivector() s = ivector() - call = getattr(sp.sparse, sparsetype + "_matrix") + call = getattr(scipy_sparse, sparsetype + "_matrix") spm = call(random_lil((300, 400), config.floatX, 5)) out = pytensor.grad(dense_from_sparse(CSM(sparsetype)(x, y, z, s)).sum(), x) self._compile_and_check( @@ -465,7 +419,7 @@ def test_transpose(self): self._compile_and_check( [x], [x.T], - [sp.sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], + [scipy_sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], Transpose, ) @@ -474,152 +428,25 @@ def test_neg(self): self._compile_and_check( [x], [-x], - [sp.sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], + [scipy_sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], Neg, ) - def test_add_ss(self): - x = SparseTensorType("csr", dtype=config.floatX)() - y = SparseTensorType("csr", dtype=config.floatX)() - self._compile_and_check( - [x, y], - [x + y], - [ - sp.sparse.csr_matrix(random_lil((10, 40), config.floatX, 3)), - sp.sparse.csr_matrix(random_lil((10, 40), config.floatX, 3)), - ], - AddSS, - ) - - def test_add_sd(self): - x = SparseTensorType("csr", dtype=config.floatX)() - y = matrix() - self._compile_and_check( - [x, y], - [x + y], - [ - sp.sparse.csr_matrix(random_lil((10, 40), config.floatX, 3)), - np.random.standard_normal((10, 40)).astype(config.floatX), - ], - (AddSD, AddSD_ccode), - ) - - def test_mul_ss(self): - x = SparseTensorType("csr", dtype=config.floatX)() - y = SparseTensorType("csr", dtype=config.floatX)() - self._compile_and_check( - [x, y], - [x * y], - [ - sp.sparse.csr_matrix(random_lil((10, 40), config.floatX, 3)), - ] - * 2, - MulSS, - ) - - def test_mul_sd(self): - x = SparseTensorType("csr", dtype=config.floatX)() - y = matrix() - self._compile_and_check( - [x, y], - [x * y], - [ - sp.sparse.csr_matrix(random_lil((10, 40), config.floatX, 3)), - np.random.standard_normal((10, 40)).astype(config.floatX), - ], - MulSD, - excluding=["local_mul_s_d"], - ) - def test_remove0(self): x = SparseTensorType("csr", dtype=config.floatX)() self._compile_and_check( [x], [Remove0()(x)], - [sp.sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], + [scipy_sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], Remove0, ) - def test_dot(self): - x = SparseTensorType("csc", dtype=config.floatX)() - y = SparseTensorType("csc", dtype=config.floatX)() - self._compile_and_check( - [x, y], - [Dot()(x, y)], - [ - sp.sparse.csc_matrix(random_lil((4, 5), config.floatX, 3)), - sp.sparse.csc_matrix(random_lil((5, 3), config.floatX, 3)), - ], - Dot, - ) - - def test_dot_broadcast(self): - for x, y in [ - (SparseTensorType("csr", "float32")(), vector()[:, None]), - (SparseTensorType("csr", "float32")(), vector()[None, :]), - (SparseTensorType("csr", "float32")(), matrix()), - (vector()[:, None], SparseTensorType("csr", "float32")()), - (vector()[None, :], SparseTensorType("csr", "float32")()), - (matrix(), SparseTensorType("csr", "float32")()), - ]: - sparse_out = pt.dot(x, y) - if isinstance(x, sparse.SparseVariable): - x = matrix() - if isinstance(y, sparse.SparseVariable): - y = matrix() - dense_out = pt.dot(x, y) - assert dense_out.broadcastable == sparse_out.broadcastable - - def test_structured_dot(self): - x = SparseTensorType("csc", dtype=config.floatX)() - y = SparseTensorType("csc", dtype=config.floatX)() - self._compile_and_check( - [x, y], - [structured_dot(x, y)], - [ - sp.sparse.csc_matrix(random_lil((4, 5), config.floatX, 3)), - sp.sparse.csc_matrix(random_lil((5, 3), config.floatX, 3)), - ], - StructuredDot, - ) - - @pytest.mark.skip( - reason="infer_shape not implemented for the grad of structured_dot" - ) - def test_structured_dot_grad(self): - # We also need the grad of CSM to be implemetned. - for format, op in [ - ("csc", StructuredDotGradCSC), - ("csr", StructuredDotGradCSR), - ]: - x = SparseTensorType(format, dtype=config.floatX)() - y = SparseTensorType(format, dtype=config.floatX)() - grads = pytensor.grad(dense_from_sparse(structured_dot(x, y)).sum(), [x, y]) - self._compile_and_check( - [x, y], - [grads[0]], - [ - as_sparse_format(random_lil((4, 5), config.floatX, 3), format), - as_sparse_format(random_lil((5, 3), config.floatX, 3), format), - ], - op, - ) - self._compile_and_check( - [x, y], - [grads[1]], - [ - as_sparse_format(random_lil((4, 5), config.floatX, 3), format), - as_sparse_format(random_lil((5, 3), config.floatX, 3), format), - ], - op, - ) - def test_dense_from_sparse(self): x = SparseTensorType("csr", dtype=config.floatX)() self._compile_and_check( [x], [dense_from_sparse(x)], - [sp.sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], + [scipy_sparse.csr_matrix(random_lil((10, 40), config.floatX, 3))], dense_from_sparse.__class__, ) @@ -712,312 +539,6 @@ def test_err(self): pytensor.grad(sub.sum(), t) -class TestAddMul: - def test_AddSS(self): - self._testSS(add) - - def test_AddSD(self): - self._testSD(add) - - def test_AddDS(self): - self._testDS(add) - - def test_MulSS(self): - self._testSS( - multiply, - np.array([[1.0, 0], [3, 0], [0, 6]]), - np.array([[1.0, 2], [3, 0], [0, 6]]), - ) - - def test_MulSD(self): - self._testSD( - multiply, - np.array([[1.0, 0], [3, 0], [0, 6]]), - np.array([[1.0, 2], [3, 0], [0, 6]]), - ) - - def test_MulDS(self): - self._testDS( - multiply, - np.array([[1.0, 0], [3, 0], [0, 6]]), - np.array([[1.0, 2], [3, 0], [0, 6]]), - ) - - def _testSS( - self, - op, - array1=None, - array2=None, - ): - if array1 is None: - array1 = np.array([[1.0, 0], [3, 0], [0, 6]]) - if array2 is None: - array2 = np.asarray([[0, 2.0], [0, 4], [5, 0]]) - - for mtype1, mtype2 in product(_mtypes, _mtypes): - for dtype1, dtype2 in [ - ("float64", "int8"), - ("int8", "float64"), - ("float64", "float64"), - ]: - a = mtype1(array1).astype(dtype1) - aR = as_sparse_variable(a) - assert aR.data is not a - assert _is_sparse(a) - assert _is_sparse_variable(aR) - - b = mtype2(array2).astype(dtype2) - bR = as_sparse_variable(b) - assert bR.data is not b - assert _is_sparse(b) - assert _is_sparse_variable(bR) - - apb = op(aR, bR) - assert _is_sparse_variable(apb) - - assert apb.type.format == aR.type.format, apb.type.format - - val = eval_outputs([apb]) - assert val.shape == (3, 2) - if op is add: - assert np.all(val.todense() == array1 + array2) - if dtype1.startswith("float") and dtype2.startswith("float"): - verify_grad_sparse(op, [a, b], structured=False) - elif op is multiply: - assert np.all(val.todense() == array1 * array2) - if dtype1.startswith("float") and dtype2.startswith("float"): - verify_grad_sparse(op, [a, b], structured=False) - - def _testSD( - self, - op, - array1=None, - array2=None, - ): - if array1 is None: - array1 = np.array([[1.0, 0], [3, 0], [0, 6]]) - if array2 is None: - array2 = np.asarray([[0, 2.0], [0, 4], [5, 0]]) - - for mtype in _mtypes: - for a in [ - np.array(array1), - pt.as_tensor_variable(array1), - pytensor.shared(array1), - ]: - for dtype1, dtype2 in [ - ("float64", "int8"), - ("int8", "float64"), - # Needed to test the grad - ("float32", "float64"), - ]: - a = a.astype(dtype1) - b = mtype(array2).astype(dtype2) - bR = as_sparse_variable(b) - assert bR.data is not b # constants are copied - assert _is_sparse(b) - assert _is_sparse_variable(bR) - - apb = op(a, bR) - - val = eval_outputs([apb]) - assert val.shape == (3, 2) - if op is add: - assert _is_dense_variable(apb) - assert np.all(val == array1 + b) - ans = np.array([[1.0, 2], [3, 4], [5, 6]]) - assert np.all(val == ans) - if isinstance(a, Constant): - a = a.data - if getattr(a, "owner", None): - continue - if dtype1.startswith("float") and dtype2.startswith("float"): - verify_grad_sparse(op, [a, b], structured=True) - elif op is multiply: - assert _is_sparse_variable(apb) - assert np.all(val.todense() == b.multiply(array1)) - assert np.all( - val.todense() == np.array([[1, 0], [9, 0], [0, 36]]) - ) - if isinstance(a, Constant): - a = a.data - if getattr(a, "owner", None): - continue - if dtype1.startswith("float") and dtype2.startswith("float"): - verify_grad_sparse(op, [a, b], structured=False) - - def _testDS( - self, - op, - array1=None, - array2=None, - ): - if array1 is None: - array1 = np.array([[1.0, 0], [3, 0], [0, 6]]) - if array2 is None: - array2 = np.asarray([[0, 2.0], [0, 4], [5, 0]]) - - for mtype in _mtypes: - for b in [ - np.asarray(array2), - pt.as_tensor_variable(array2), - pytensor.shared(array2), - ]: - for dtype1, dtype2 in [ - ("float64", "int8"), - ("int8", "float64"), - ]: - a = mtype(array1).astype(dtype1) - aR = as_sparse_variable(a) - assert aR.data is not a - assert _is_sparse(a) - assert _is_sparse_variable(aR) - b = b.astype(dtype2) - - apb = op(aR, b) - - val = eval_outputs([apb]) - assert val.shape == (3, 2) - if op is add: - assert _is_dense_variable(apb) - assert np.all(val == a + array2) - ans = np.array([[1.0, 2], [3, 4], [5, 6]]) - assert np.all(val == ans) - if isinstance(b, Constant): - b = b.data - if dtype1.startswith("float") and dtype2.startswith("float"): - verify_grad_sparse(op, [a, b], structured=True) - elif op is multiply: - assert _is_sparse_variable(apb) - ans = np.array([[1, 0], [9, 0], [0, 36]]) - assert np.all(val.todense() == (a.multiply(array2))) - assert np.all(val.todense() == ans) - if isinstance(b, Constant): - b = b.data - if dtype1.startswith("float") and dtype2.startswith("float"): - verify_grad_sparse(op, [a, b], structured=False) - - -class TestComparison: - # took from tensor basic_test.py - def _rand_ranged(self, min, max, shape): - return np.asarray( - np.random.random(shape) * (max - min) + min, dtype=config.floatX - ) - - tests = [ - lambda x, y: x > y, - lambda x, y: x < y, - lambda x, y: x >= y, - lambda x, y: x <= y, - ] - - testsDic = { - gt: lambda x, y: x > y, - lt: lambda x, y: x < y, - ge: lambda x, y: x >= y, - le: lambda x, y: x <= y, - } - - @pytest.mark.skipif( - version.parse(sp.__version__) < version.parse("0.13"), - reason="Comparison operators need newer release of scipy", - ) - def __generalized_ss_test(self, pytensorp, symbolicType, testOp, scipyType): - x = symbolicType() - y = symbolicType() - - op = pytensorp(x, y) - - f = pytensor.function([x, y], op) - - m1 = scipyType(random_lil((10, 40), config.floatX, 3)) - m2 = scipyType(random_lil((10, 40), config.floatX, 3)) - - assert np.array_equal(f(m1, m2).data, testOp(m1, m2).data) - - @pytest.mark.skipif( - version.parse(sp.__version__) < version.parse("0.13"), - reason="Comparison operators need newer release of scipy", - ) - def __generalized_sd_test(self, pytensorp, symbolicType, testOp, scipyType): - x = symbolicType() - y = matrix() - - op = pytensorp(x, y) - - f = pytensor.function([x, y], op) - - m1 = scipyType(random_lil((10, 40), config.floatX, 3)) - m2 = self._rand_ranged(1000, -1000, [10, 40]) - - assert np.array_equal(f(m1, m2).data, testOp(m1, m2).data) - - @pytest.mark.skipif( - version.parse(sp.__version__) < version.parse("0.13"), - reason="Comparison operators need newer release of scipy", - ) - def __generalized_ds_test(self, pytensorp, symbolicType, testOp, scipyType): - x = symbolicType() - y = matrix() - - op = pytensorp(y, x) - - f = pytensor.function([y, x], op) - - m1 = scipyType(random_lil((10, 40), config.floatX, 3)) - m2 = self._rand_ranged(1000, -1000, [10, 40]) - - assert np.array_equal(f(m2, m1).data, testOp(m2, m1).data) - - def test_ss_csr_comparison(self): - for op in self.tests: - self.__generalized_ss_test(op, sparse.csr_matrix, op, sp.sparse.csr_matrix) - - def test_ss_csc_comparison(self): - for op in self.tests: - self.__generalized_ss_test(op, sparse.csc_matrix, op, sp.sparse.csc_matrix) - - def test_sd_csr_comparison(self): - for op in self.tests: - self.__generalized_sd_test(op, sparse.csr_matrix, op, sp.sparse.csr_matrix) - - def test_sd_csc_comparison(self): - for op in self.tests: - self.__generalized_sd_test(op, sparse.csc_matrix, op, sp.sparse.csc_matrix) - - def test_ds_csc_comparison(self): - for op in self.testsDic: - self.__generalized_ds_test( - op, sparse.csc_matrix, self.testsDic[op], sp.sparse.csc_matrix - ) - - def test_ds_csr_comparison(self): - for op in self.testsDic: - self.__generalized_ds_test( - op, sparse.csr_matrix, self.testsDic[op], sp.sparse.csr_matrix - ) - - @pytest.mark.skipif( - version.parse(sp.__version__) < version.parse("0.13"), - reason="Comparison operators need newer release of scipy", - ) - def test_equality_case(self): - # Test assuring normal behaviour when values - # in the matrices are equal - x = sparse.csc_matrix() - y = matrix() - - m1 = sp.sparse.csc_matrix((2, 2), dtype=pytensor.config.floatX) - m2 = np.asarray([[0, 0], [0, 0]], dtype=pytensor.config.floatX) - - for func in self.testsDic: - op = func(y, x) - f = pytensor.function([y, x], op) - - assert np.array_equal(f(m2, m1), self.testsDic[func](m2, m1)) - - class TestConversion: def test_basic(self): test_val = np.random.random((5,)).astype(config.floatX) @@ -1034,34 +555,22 @@ def test_basic(self): assert val.format == "csr" test_val = np.eye(3).astype(config.floatX) - a = sp.sparse.csr_matrix(test_val) + a = scipy_sparse.csr_matrix(test_val) s = as_sparse_or_tensor_variable(a) res = pt.as_tensor_variable(s) assert isinstance(res, SparseConstant) - a = sp.sparse.csr_matrix(test_val) + a = scipy_sparse.csr_matrix(test_val) s = as_sparse_or_tensor_variable(a) from pytensor.tensor.exceptions import NotScalarConstantError with pytest.raises(NotScalarConstantError): pt.get_underlying_scalar_constant_value(s, only_process_constants=True) - # TODO: - # def test_sparse_as_tensor_variable(self): - # csr = sp.sparse.csr_matrix(np.eye(3)) - # val = aet.as_tensor_variable(csr) - # assert str(val.dtype) == config.floatX - # assert val.format == "csr" - # - # csr = sp.sparse.csc_matrix(np.eye(3)) - # val = aet.as_tensor_variable(csr) - # assert str(val.dtype) == config.floatX - # assert val.format == "csc" - def test_dense_from_sparse(self): # call dense_from_sparse for t in _mtypes: - s = t(sp.sparse.identity(5)) + s = t(scipy_sparse.identity(5)) s = as_sparse_variable(s) d = dense_from_sparse(s) val = eval_outputs([d]) @@ -1071,7 +580,7 @@ def test_dense_from_sparse(self): def test_todense(self): # call sparse_var.todense() for t in _mtypes: - s = t(sp.sparse.identity(5)) + s = t(scipy_sparse.identity(5)) s = as_sparse_variable(s) d = s.toarray() val = eval_outputs([d]) @@ -1104,7 +613,7 @@ def test_format_ndim(self): class TestCsmProperties: def test_csm_properties_grad(self): - sp_types = {"csc": sp.sparse.csc_matrix, "csr": sp.sparse.csr_matrix} + sp_types = {"csc": scipy_sparse.csc_matrix, "csr": scipy_sparse.csr_matrix} for format in ("csc", "csr"): for dtype in ("float32", "float64"): @@ -1127,7 +636,7 @@ def test_csm_properties_grad(self): ) def test_csm_properties(self): - sp_types = {"csc": sp.sparse.csc_matrix, "csr": sp.sparse.csr_matrix} + sp_types = {"csc": scipy_sparse.csc_matrix, "csr": scipy_sparse.csr_matrix} for format in ("csc", "csr"): for dtype in ("float32", "float64"): @@ -1146,7 +655,7 @@ def test_csm_properties(self): class TestCsm: def test_csm_grad(self): - sp_types = {"csc": sp.sparse.csc_matrix, "csr": sp.sparse.csr_matrix} + sp_types = {"csc": scipy_sparse.csc_matrix, "csr": scipy_sparse.csr_matrix} for format in ("csc", "csr"): for dtype in ("float32", "float64"): @@ -1161,13 +670,13 @@ def test_csm_grad(self): ) @pytest.mark.skipif( - version.parse(sp.__version__) >= version.parse("1.16.0"), + version.parse(scipy_version) >= version.parse("1.16.0"), reason="Scipy 1.16 introduced some changes that make this test fail", ) def test_csm_sparser(self): # Test support for gradients sparser than the input. - sp_types = {"csc": sp.sparse.csc_matrix, "csr": sp.sparse.csr_matrix} + sp_types = {"csc": scipy_sparse.csc_matrix, "csr": scipy_sparse.csr_matrix} for format in ("csc", "csr"): for dtype in ("float32", "float64"): @@ -1197,7 +706,7 @@ def test_csm_sparser(self): assert len(spmat.data) == len(res) @pytest.mark.skipif( - version.parse(sp.__version__) >= version.parse("1.16.0"), + version.parse(scipy_version) >= version.parse("1.16.0"), reason="Scipy 1.16 introduced some changes that make this test fail", ) def test_csm_unsorted(self): @@ -1224,7 +733,7 @@ def my_op(x): verify_grad_sparse(my_op, [a.data]) def test_csm(self): - sp_types = {"csc": sp.sparse.csc_matrix, "csr": sp.sparse.csr_matrix} + sp_types = {"csc": scipy_sparse.csc_matrix, "csr": scipy_sparse.csr_matrix} for format in ("csc", "csr"): for dtype in ("float32", "float64"): @@ -1249,636 +758,11 @@ def test_csm(self): assert np.all(res.shape == spmat.shape) -class TestStructuredDot: - def test_structureddot_csc_grad(self): - # shortcut: testing csc in float32, testing csr in float64 - - # allocate a random sparse matrix - spmat = sp.sparse.csc_matrix(random_lil((4, 3), "float32", 3)) - - mat = np.asarray(np.random.standard_normal((3, 2)), "float32") - - verify_grad_sparse(structured_dot, [spmat, mat], structured=True) - - def buildgraph_T(spmat, mat): - return structured_dot(mat.T, spmat.T) - - verify_grad_sparse(buildgraph_T, [spmat, mat], structured=True) - - def test_structureddot_csr_grad(self): - # shortcut: testing csc in float32, testing csr in float64 - - # allocate a random sparse matrix - spmat = sp.sparse.csr_matrix(random_lil((4, 3), "float64", 3)) - - mat = np.asarray(np.random.standard_normal((3, 2)), "float64") - - verify_grad_sparse(structured_dot, [spmat, mat], structured=True) - - def buildgraph_T(spmat, mat): - return structured_dot(mat.T, spmat.T) - - verify_grad_sparse(buildgraph_T, [spmat, mat], structured=True) - - def test_upcast(self): - typenames = ( - "float32", - "int64", - "int8", - "int32", - "int16", - "float64", - "complex64", - "complex128", - ) - for dense_dtype in typenames: - for sparse_dtype in typenames: - correct_dtype = pytensor.scalar.upcast(sparse_dtype, dense_dtype) - a = SparseTensorType("csc", dtype=sparse_dtype)() - b = matrix(dtype=dense_dtype) - d = structured_dot(a, b) - assert d.type.dtype == correct_dtype - - # compile and run a function - - f = pytensor.function([a, b], d) - - M, N, K, nnz = (4, 3, 5, 3) - spmat = sp.sparse.csc_matrix(random_lil((M, N), sparse_dtype, nnz)) - # the following madness is necessary to workaround - # an intc vs. int32 bug. - # The lil makes an intc on my computer when sparse_dtype - # is int32. - spmat.dtype = np.dtype(sparse_dtype) - mat = np.asarray( - np.random.standard_normal((N, K)) * 9, dtype=dense_dtype - ) - # print 'DTYPES', sparse_dtype, dense_dtype - # print 'sym types', a.type, b.type - # print 'dtype strings', spmat.dtype, mat.dtype - # print 'numpy dtype num', mat.dtype.num - # print 'scipy dtype num', spmat.data.dtype.num - pytensor_result = f(spmat, mat) - scipy_result = spmat * mat - assert pytensor_result.shape == scipy_result.shape - assert pytensor_result.dtype == scipy_result.dtype - utt.assert_allclose(scipy_result, pytensor_result) - - def test_opt_unpack(self): - # - # Test that a graph involving - # structured_dot(assembled_csc_matrix) is optimized to be just - # a structured_dot_csc Op and no assembly of a csc_matrix. - # - # The optimization from structured_dot -> structured_dot_csc - # is currently disabled, So this test is not expected to pass - - return - # - kerns = TensorType(dtype="int64", shape=(None,))("kerns") - spmat = sp.sparse.lil_matrix((4, 6), dtype="int64") - for i in range(5): - # set non-zeros in random locations (row x, col y) - x = np.floor(np.random.random() * spmat.shape[0]) - y = np.floor(np.random.random() * spmat.shape[1]) - spmat[x, y] = np.random.random() * 10 - spmat = sp.sparse.csc_matrix(spmat) - - images = TensorType(dtype="float32", shape=(None, None))("images") - - cscmat = CSC(kerns, spmat.indices[: spmat.size], spmat.indptr, spmat.shape) - f = pytensor.function([kerns, images], structured_dot(cscmat, images.T)) - - sdcscpresent = False - for node in f.maker.fgraph.toposort(): - # print node.op - assert not isinstance(node.op, CSM) - assert not isinstance(node.op, CSMProperties) - if isinstance(f.maker.fgraph.toposort()[1].op, StructuredDotCSC): - sdcscpresent = True - assert sdcscpresent - - kernvals = np.array(spmat.data[: spmat.size]) - # print 'kdtype', kernvals.dtype, kernvals.shape, - # print kernvals.ndim, kernvals.dtype.num - # print 'type of kernvals = ', kernvals.dtype - bsize = 3 - imvals = 1.0 * np.array( - np.arange(bsize * spmat.shape[1]).reshape(bsize, spmat.shape[1]), - dtype="float32", - ) - f(kernvals, imvals) - # print outvals - - def test_dot_sparse_sparse(self): - # test dot for 2 input sparse matrix - sparse_dtype = "float64" - sp_mat = { - "csc": sp.sparse.csc_matrix, - "csr": sp.sparse.csr_matrix, - "bsr": sp.sparse.csr_matrix, - } - - for sparse_format_a in ["csc", "csr", "bsr"]: - for sparse_format_b in ["csc", "csr", "bsr"]: - a = SparseTensorType(sparse_format_a, dtype=sparse_dtype)() - b = SparseTensorType(sparse_format_b, dtype=sparse_dtype)() - d = pt.dot(a, b) - f = pytensor.function([a, b], Out(d, borrow=True)) - for M, N, K, nnz in [ - (4, 3, 2, 3), - (40, 30, 20, 3), - (40, 30, 20, 30), - (400, 3000, 200, 6000), - ]: - a_val = sp_mat[sparse_format_a]( - random_lil((M, N), sparse_dtype, nnz) - ) - b_val = sp_mat[sparse_format_b]( - random_lil((N, K), sparse_dtype, nnz) - ) - f(a_val, b_val) - - def test_csc_correct_output_faster_than_scipy(self): - sparse_dtype = "float64" - dense_dtype = "float64" - - a = SparseTensorType("csc", dtype=sparse_dtype)() - b = matrix(dtype=dense_dtype) - d = pt.dot(a, b) - f = pytensor.function([a, b], Out(d, borrow=True)) - - for M, N, K, nnz in [ - (4, 3, 2, 3), - (40, 30, 20, 3), - (40, 30, 20, 30), - (400, 3000, 200, 6000), - ]: - spmat = sp.sparse.csc_matrix(random_lil((M, N), sparse_dtype, nnz)) - mat = np.asarray(np.random.standard_normal((N, K)), dense_dtype) - pytensor_times = [] - scipy_times = [] - for i in range(5): - t0 = time.perf_counter() - pytensor_result = f(spmat, mat) - t1 = time.perf_counter() - scipy_result = spmat * mat - t2 = time.perf_counter() - - pytensor_times.append(t1 - t0) - scipy_times.append(t2 - t1) - - pytensor_time = np.min(pytensor_times) - scipy_time = np.min(scipy_times) - - # speedup = scipy_time / pytensor_time - # print scipy_times - # print pytensor_times - # print ('M=%(M)s N=%(N)s K=%(K)s nnz=%(nnz)s pytensor_time' - # '=%(pytensor_time)s speedup=%(speedup)s') % locals() - - # fail if PyTensor is slower than scipy by more than a certain amount - overhead_tol = 0.003 # seconds overall - overhead_rtol = 1.2 # times as long - utt.assert_allclose(scipy_result, pytensor_result) - if pytensor.config.mode == "FAST_RUN" and pytensor.config.cxx: - assert pytensor_time <= overhead_rtol * scipy_time + overhead_tol - - def test_csr_correct_output_faster_than_scipy(self): - # contrast with test_grad, we put csr in float32, csc in float64 - - sparse_dtype = "float32" - dense_dtype = "float32" - - a = SparseTensorType("csr", dtype=sparse_dtype)() - b = matrix(dtype=dense_dtype) - d = pt.dot(a, b) - f = pytensor.function([a, b], d) - - for M, N, K, nnz in [ - (4, 3, 2, 3), - (40, 30, 20, 3), - (40, 30, 20, 30), - (400, 3000, 200, 6000), - ]: - spmat = sp.sparse.csr_matrix(random_lil((M, N), sparse_dtype, nnz)) - mat = np.asarray(np.random.standard_normal((N, K)), dense_dtype) - t0 = time.perf_counter() - pytensor_result = f(spmat, mat) - t1 = time.perf_counter() - scipy_result = spmat * mat - t2 = time.perf_counter() - - pytensor_time = t1 - t0 - scipy_time = t2 - t1 - # print 'pytensor took', pytensor_time, - # print 'scipy took', scipy_time - overhead_tol = 0.002 # seconds - overhead_rtol = 1.1 # times as long - utt.assert_allclose(scipy_result, pytensor_result) - if pytensor.config.mode == "FAST_RUN" and pytensor.config.cxx: - assert pytensor_time <= overhead_rtol * scipy_time + overhead_tol, ( - pytensor_time, - overhead_rtol * scipy_time + overhead_tol, - scipy_time, - overhead_rtol, - overhead_tol, - ) - - -class TestDots(utt.InferShapeTester): - def setup_method(self): - super().setup_method() - x_size = (10, 100) - y_size = (100, 1000) - - self.x_csr = sp.sparse.csr_matrix( - np.random.binomial(1, 0.5, x_size), dtype=pytensor.config.floatX - ) - self.x_csc = sp.sparse.csc_matrix( - np.random.binomial(1, 0.5, x_size), dtype=pytensor.config.floatX - ) - self.y = np.asarray( - np.random.uniform(-1, 1, y_size), dtype=pytensor.config.floatX - ) - self.y_csr = sp.sparse.csr_matrix( - np.random.binomial(1, 0.5, y_size), dtype=pytensor.config.floatX - ) - self.y_csc = sp.sparse.csc_matrix( - np.random.binomial(1, 0.5, y_size), dtype=pytensor.config.floatX - ) - self.v_10 = np.asarray( - np.random.uniform(-1, 1, 10), dtype=pytensor.config.floatX - ) - self.v_100 = np.asarray( - np.random.uniform(-1, 1, 100), dtype=pytensor.config.floatX - ) - - def test_csr_dense(self): - x = sparse.csr_matrix("x") - y = matrix("y") - v = vector("v") - - for x, y, x_v, y_v in [ - (x, y, self.x_csr, self.y), - (x, v, self.x_csr, self.v_100), - (v, x, self.v_10, self.x_csr), - ]: - f_a = pytensor.function([x, y], sparse.dot(x, y)) - - def f_b(x, y): - return x * y - - utt.assert_allclose(f_a(x_v, y_v), f_b(x_v, y_v)) - - # Test infer_shape - self._compile_and_check( - [x, y], [sparse.dot(x, y)], [x_v, y_v], (Dot, Usmm, UsmmCscDense) - ) - - def test_csc_dense(self): - x = sparse.csc_matrix("x") - y = matrix("y") - v = vector("v") - - for x, y, x_v, y_v in [ - (x, y, self.x_csc, self.y), - (x, v, self.x_csc, self.v_100), - (v, x, self.v_10, self.x_csc), - ]: - f_a = pytensor.function([x, y], sparse.dot(x, y)) - - def f_b(x, y): - return x * y - - utt.assert_allclose(f_a(x_v, y_v), f_b(x_v, y_v)) - - # Test infer_shape - self._compile_and_check( - [x, y], [sparse.dot(x, y)], [x_v, y_v], (Dot, Usmm, UsmmCscDense) - ) - - def test_sparse_sparse(self): - for d1, d2 in [ - ("float32", "float32"), - ("float32", "float64"), - ("float64", "float32"), - ("float64", "float64"), - ("float32", "int16"), - ("float32", "complex64"), - ]: - for x_f, y_f in [ - ("csc", "csc"), - ("csc", "csr"), - ("csr", "csc"), - ("csr", "csr"), - ]: - x = sparse.SparseTensorType(format=x_f, dtype=d1)("x") - y = sparse.SparseTensorType(format=x_f, dtype=d2)("x") - - def f_a(x, y): - return x * y - - f_b = pytensor.function([x, y], sparse.dot(x, y)) - - vx = getattr(self, "x_" + x_f).astype(d1) - vy = getattr(self, "y_" + y_f).astype(d2) - utt.assert_allclose(f_a(vx, vy).toarray(), f_b(vx, vy)) - - # Test infer_shape - f_a = pytensor.function([x, y], sparse.dot(x, y).shape) - - def f_b(x, y): - return (x * y).shape - - assert np.all(f_a(vx, vy) == f_b(vx, vy)) - topo = f_a.maker.fgraph.toposort() - assert not any( - isinstance(node.op, Dot | Usmm | UsmmCscDense) for node in topo - ) - - def test_int32_dtype(self): - # Reported on the theano-user mailing-list: - # https://groups.google.com/d/msg/theano-users/MT9ui8LtTsY/rwatwEF9zWAJ - size = 9 - intX = "int32" - - C = matrix("C", dtype=intX) - I = matrix("I", dtype=intX) - - fI = I.flatten() - data = pt.ones_like(fI) - indptr = pt.arange(data.shape[0] + 1, dtype="int32") - - m1 = sparse.CSR(data, fI, indptr, (8, size)) - m2 = sparse.dot(m1, C) - y = m2.reshape(shape=(2, 4, 9), ndim=3) - - f = pytensor.function(inputs=[I, C], outputs=y) - i = np.asarray([[4, 3, 7, 7], [2, 8, 4, 5]], dtype=intX) - a = np.asarray( - np.random.default_rng().integers(0, 100, (size, size)), dtype=intX - ) - f(i, a) - - def test_tensor_dot_types(self): - x = sparse.csc_matrix("x") - x_d = pt.matrix("x_d") - y = sparse.csc_matrix("y") - - res = pt.dot(x, y) - op_types = {type(n.op) for n in applys_between([x, y], [res])} - assert sparse.basic.StructuredDot in op_types - assert pt.math.Dot not in op_types - - res = pt.dot(x_d, y) - op_types = {type(n.op) for n in applys_between([x, y], [res])} - assert sparse.basic.StructuredDot in op_types - assert pt.math.Dot not in op_types - - res = pt.dot(x, x_d) - op_types = {type(n.op) for n in applys_between([x, y], [res])} - assert sparse.basic.StructuredDot in op_types - assert pt.math.Dot not in op_types - - res = pt.dot(pt.second(1, x), y) - op_types = {type(n.op) for n in applys_between([x, y], [res])} - assert sparse.basic.StructuredDot in op_types - assert pt.math.Dot not in op_types - - def test_csr_dense_grad(self): - # shortcut: testing csc in float32, testing csr in float64 - - # allocate a random sparse matrix - spmat = sp.sparse.csr_matrix(random_lil((4, 3), "float64", 3)) - - mat = np.asarray(np.random.standard_normal((2, 4)), "float64") - - def buildgraph_T(mat): - return Dot()(mat, spmat) - - utt.verify_grad(buildgraph_T, [mat]) - - -class TestUsmm: - """ - Test the Usmm and UsmmCscDense class and related optimization - """ - - def setup_method(self): - x_size = (10, 100) - y_size = (100, 200) - z_size = (x_size[0], y_size[1]) - - self.rng = np.random.default_rng(seed=utt.fetch_seed()) - self.x = np.asarray( - self.rng.binomial(1, 0.5, x_size), dtype=pytensor.config.floatX - ) - self.y = np.asarray( - self.rng.uniform(-1, 1, y_size), dtype=pytensor.config.floatX - ) - self.z = np.asarray( - self.rng.uniform(-1, 1, z_size), dtype=pytensor.config.floatX - ) - - @pytest.mark.slow - def test_basic(self): - # this is slow, but it's the only test for the op. - def mat(format, name, dtype): - if format == "dense": - return matrix(name, dtype=dtype) - else: - return sparse.matrix(format, name, dtype=dtype) - - params = product( - *( - [["float32", "float64", "int16", "complex64"]] * 4 - + [["dense", "csc", "csr"]] * 2 - ) - ) - - # All test are too slow, so we randomly take 100 of them. - # The buildbot change the seed, so we will finish by running them all. - # As of this writing they where all passing. - # params = self.rng.permutation(list(params))[:500] - - for dtype1, dtype2, dtype3, dtype4, format1, format2 in params: - if format1 == "dense" and format2 == "dense": - # Usmm won't be used! - continue - x = mat(format1, "x", dtype1) - y = mat(format2, "y", dtype2) - a = scalar("a", dtype=dtype3) - z = pytensor.shared(np.asarray(self.z, dtype=dtype4).copy()) - - def f_b(z, a, x, y): - return z - a * (x * y) - - x_data = np.asarray(self.x, dtype=dtype1) - if format1 != "dense": - x_data = as_sparse_format(x_data, format1) - y_data = np.asarray(self.y, dtype=dtype2) - if format2 != "dense": - y_data = as_sparse_format(y_data, format2) - a_data = np.asarray(1.5, dtype=dtype3) - z_data = np.asarray(self.z, dtype=dtype4) - - f_b_out = f_b(z_data, a_data, x_data, y_data) - - # Can it work inplace? - inplace = dtype4 == pytensor.scalar.upcast(dtype1, dtype2, dtype3) - - # To make it easier to check the toposort - mode = pytensor.compile.mode.get_default_mode().excluding("fusion") - - if inplace: - updates = [(z, z - a * sparse.dot(x, y))] - f_a = pytensor.function([a, x, y], [], updates=updates, mode=mode) - f_a(a_data, x_data, y_data) - f_a_out = z.get_value(borrow=True) - else: - f_a = pytensor.function([a, x, y], z - a * sparse.dot(x, y), mode=mode) - # In DebugMode there is a strange difference with complex - # So we raise a little the threshold a little. - try: - orig_atol = pytensor.tensor.math.float64_atol - orig_rtol = pytensor.tensor.math.float64_rtol - pytensor.tensor.math.float64_atol = 1e-7 - pytensor.tensor.math.float64_rtol = 1e-6 - f_a_out = f_a(a_data, x_data, y_data) - finally: - pytensor.tensor.math.float64_atol = orig_atol - pytensor.tensor.math.float64_rtol = orig_rtol - - # As we do a dot product of 2 vector of 100 element, - # This mean we can have 2*100*eps abs error. - if f_a_out.dtype in ["float64", "complex128"]: - atol = 3e-8 - rtol = 1e-5 - else: - atol = None - rtol = None - utt.assert_allclose(f_a_out, f_b_out, rtol=rtol, atol=atol) - topo = f_a.maker.fgraph.toposort() - up = pytensor.scalar.upcast(dtype1, dtype2, dtype3, dtype4) - - fast_compile = pytensor.config.mode == "FAST_COMPILE" - - if not pytensor.config.blas__ldflags: - # Usmm should not be inserted, because it relies on BLAS - assert len(topo) == 4, topo - assert isinstance(topo[0].op, sparse.Dot) - assert isinstance(topo[1].op, DimShuffle) - assert isinstance(topo[2].op, Elemwise) and isinstance( - topo[2].op.scalar_op, pytensor.scalar.Mul - ) - assert isinstance(topo[3].op, Elemwise) and isinstance( - topo[3].op.scalar_op, pytensor.scalar.Sub - ) - elif ( - y.type.dtype == up - and format1 == "csc" - and format2 == "dense" - and not fast_compile - and pytensor.config.cxx - and up in ("float32", "float64") - ): - # The op UsmmCscDense should be inserted - assert ( - sum( - isinstance(node.op, Elemwise) - and isinstance(node.op.scalar_op, pytensor.scalar.basic.Cast) - for node in topo - ) - == len(topo) - 5 - ) - new_topo = [ - node - for node in topo - if not ( - isinstance(node.op, Elemwise) - and isinstance(node.op.scalar_op, pytensor.scalar.basic.Cast) - ) - ] - topo = new_topo - assert len(topo) == 5, topo - # Usmm is tested at the same time in debugmode - # Check if the optimization local_usmm and local_usmm_csx is - # applied - - def check_once(x): - assert sum(isinstance(n.op, x) for n in topo) == 1 - - check_once(sparse.basic.CSMProperties) - check_once(DimShuffle) - check_once(Subtensor) - check_once(UsmmCscDense) - check_once(Elemwise) - if inplace: - assert topo[4].op.inplace - elif not fast_compile: - # The op Usmm should be inserted - assert len(topo) == 3, topo - assert isinstance(topo[0].op, DimShuffle) - assert topo[1].op == pytensor.tensor.neg - assert isinstance(topo[2].op, sparse.Usmm) - - def test_infer_shape(self): - def mat(format, name, dtype): - if format == "dense": - return matrix(name, dtype=dtype) - else: - return sparse.matrix(format, name, dtype=dtype) - - params = [ - ("float32", "float64", "int16", "complex64", "csc", "dense"), - ("float32", "float64", "int16", "complex64", "csr", "dense"), - ] - for dtype1, dtype2, dtype3, dtype4, format1, format2 in params: - if format1 == "dense" and format2 == "dense": - # Usmm won't be used! - continue - x = mat(format1, "x", dtype1) - y = mat(format2, "y", dtype2) - a = scalar("a", dtype=dtype3) - z = pytensor.shared(np.asarray(self.z, dtype=dtype4).copy()) - - def f_b(z, a, x, y): - return z - a * (x * y) - - x_data = np.asarray(self.x, dtype=dtype1) - if format1 != "dense": - x_data = as_sparse_format(x_data, format1) - y_data = np.asarray(self.y, dtype=dtype2) - if format2 != "dense": - y_data = as_sparse_format(y_data, format2) - a_data = np.asarray(1.5, dtype=dtype3) - z_data = np.asarray(self.z, dtype=dtype4) - - f_b_out = f_b(z_data, a_data, x_data, y_data) - - # Can it work inplace? - # inplace = dtype4 == pytensor.scalar.upcast(dtype1, dtype2, dtype3) - - # To make it easier to check the toposort - mode = pytensor.compile.mode.get_default_mode().excluding("fusion") - - # test infer_shape of Dot got applied - f_shape = pytensor.function( - [a, x, y], (z - a * sparse.dot(x, y)).shape, mode=mode - ) - assert all(f_shape(a_data, x_data, y_data) == f_b_out.shape) - topo = f_shape.maker.fgraph.toposort() - assert not any( - isinstance(node.op, Dot | Usmm | UsmmCscDense) for node in topo - ) - - class TestZerosLike: def test(self): x = sparse.csr_matrix() f = pytensor.function([x], sparse.sp_zeros_like(x)) - vx = sp.sparse.csr_matrix( + vx = scipy_sparse.csr_matrix( np.asarray( np.random.binomial(1, 0.5, (100, 100)), dtype=pytensor.config.floatX ) @@ -1895,7 +779,7 @@ def test_shape_i(): a = SparseTensorType("csr", dtype=sparse_dtype)() f = pytensor.function([a], a.shape[1]) - assert f(sp.sparse.csr_matrix(random_lil((100, 10), sparse_dtype, 3))) == 10 + assert f(scipy_sparse.csr_matrix(random_lil((100, 10), sparse_dtype, 3))) == 10 def test_shape(): @@ -1906,7 +790,7 @@ def test_shape(): a = SparseTensorType("csr", dtype=sparse_dtype)() f = pytensor.function([a], a.shape) assert np.all( - f(sp.sparse.csr_matrix(random_lil((100, 10), sparse_dtype, 3))) == (100, 10) + f(scipy_sparse.csr_matrix(random_lil((100, 10), sparse_dtype, 3))) == (100, 10) ) if pytensor.config.mode != "FAST_COMPILE": topo = f.maker.fgraph.toposort() @@ -1917,8 +801,8 @@ def test_shape(): def test_may_share_memory(): - a = sp.sparse.csc_matrix(sp.sparse.eye(5, 3)) - b = sp.sparse.csc_matrix(sp.sparse.eye(4, 3)) + a = scipy_sparse.csc_matrix(scipy_sparse.eye(5, 3)) + b = scipy_sparse.csc_matrix(scipy_sparse.eye(4, 3)) def as_ar(a): return np.asarray(a, dtype="int32") @@ -1964,7 +848,7 @@ def test_sparse_shared_memory(): x = SparseTensorType("csr", dtype="float32")() y = SparseTensorType("csr", dtype="float32")() - sdot = sparse.structured_dot + sdot = sparse.math.structured_dot z = sdot(x * 3, m1) + sdot(y * 2, m2) f = pytensor.function( @@ -1985,7 +869,7 @@ def test_size(): for sparse_type in ("csc_matrix", "csr_matrix"): x = getattr(pytensor.sparse, sparse_type)() - y = getattr(sp.sparse, sparse_type)((5, 7)).astype(config.floatX) + y = getattr(scipy_sparse, sparse_type)((5, 7)).astype(config.floatX) get_size = pytensor.function([x], x.size) def check(): @@ -2074,55 +958,6 @@ def test_grad(self): verify_grad_sparse(self.op, data, structured=True) -class TestSpSum(utt.InferShapeTester): - possible_axis = [None, 0, 1] - - def setup_method(self): - super().setup_method() - self.op_class = sparse.SpSum - self.op = sparse.sp_sum - - @pytest.mark.parametrize("op_type", ["func", "method"]) - def test_op(self, op_type): - for format in sparse.sparse_formats: - for axis in self.possible_axis: - variable, data = sparse_random_inputs(format, shape=(10, 10)) - - if op_type == "func": - z = sparse.sp_sum(variable[0], axis=axis) - if op_type == "method": - z = variable[0].sum(axis=axis) - - if axis is None: - assert z.type.broadcastable == () - else: - assert z.type.broadcastable == (False,) - - f = pytensor.function(variable, self.op(variable[0], axis=axis)) - tested = f(*data) - expected = data[0].todense().sum(axis).ravel() - utt.assert_allclose(expected, tested) - - def test_infer_shape(self): - for format in sparse.sparse_formats: - for axis in self.possible_axis: - variable, data = sparse_random_inputs(format, shape=(9, 10)) - self._compile_and_check( - variable, [self.op(variable[0], axis=axis)], data, self.op_class - ) - - def test_grad(self): - for format in sparse.sparse_formats: - for axis in self.possible_axis: - for struct in [True, False]: - _variable, data = sparse_random_inputs(format, shape=(9, 10)) - verify_grad_sparse( - self.op_class(axis=axis, sparse_grad=struct), - data, - structured=struct, - ) - - class TestDiag(utt.InferShapeTester): def setup_method(self): super().setup_method() @@ -2264,8 +1099,8 @@ def setup_method(self): def test_remove0(self): configs = [ # structure type, numpy matching class - ("csc", sp.sparse.csc_matrix), - ("csr", sp.sparse.csr_matrix), + ("csc", scipy_sparse.csc_matrix), + ("csr", scipy_sparse.csr_matrix), ] for format, matrix_class in configs: @@ -2323,21 +1158,21 @@ def test_infer_shape(self): mat[0, 1] = mat[1, 0] = mat[2, 2] = 0 x_csc = sparse.csc_matrix(dtype=pytensor.config.floatX) - mat_csc = sp.sparse.csc_matrix(mat, dtype=pytensor.config.floatX) + mat_csc = scipy_sparse.csc_matrix(mat, dtype=pytensor.config.floatX) self._compile_and_check([x_csc], [Remove0()(x_csc)], [mat_csc], self.op_class) x_csr = sparse.csr_matrix(dtype=pytensor.config.floatX) - mat_csr = sp.sparse.csr_matrix(mat, dtype=pytensor.config.floatX) + mat_csr = scipy_sparse.csr_matrix(mat, dtype=pytensor.config.floatX) self._compile_and_check([x_csr], [Remove0()(x_csr)], [mat_csr], self.op_class) def test_grad(self): mat = (np.arange(9) + 1).reshape((3, 3)) mat[0, 1] = mat[1, 0] = mat[2, 2] = 0 - mat_csc = sp.sparse.csc_matrix(mat, dtype=pytensor.config.floatX) + mat_csc = scipy_sparse.csc_matrix(mat, dtype=pytensor.config.floatX) verify_grad_sparse(Remove0(), [mat_csc]) - mat_csr = sp.sparse.csr_matrix(mat, dtype=pytensor.config.floatX) + mat_csr = scipy_sparse.csr_matrix(mat, dtype=pytensor.config.floatX) verify_grad_sparse(Remove0(), [mat_csr]) @@ -2357,8 +1192,8 @@ def test_GetItemList(self): t_geta = fa(A[0]).todense() t_getb = fb(B[0]).todense() - s_geta = sp.sparse.csr_matrix(A[0])[[0, 1, 2, 3, 1]].todense() - s_getb = sp.sparse.csc_matrix(B[0])[[0, 1, 2, 3, 1]].todense() + s_geta = scipy_sparse.csr_matrix(A[0])[[0, 1, 2, 3, 1]].todense() + s_getb = scipy_sparse.csc_matrix(B[0])[[0, 1, 2, 3, 1]].todense() utt.assert_allclose(t_geta, s_geta) utt.assert_allclose(t_getb, s_getb) @@ -2396,8 +1231,8 @@ def test_GetItem2Lists(self): t_geta = fa(A[0]) t_getb = fb(B[0]) - s_geta = np.asarray(sp.sparse.csr_matrix(A[0])[[0, 0, 1, 3], [0, 1, 2, 4]]) - s_getb = np.asarray(sp.sparse.csc_matrix(B[0])[[0, 0, 1, 3], [0, 1, 2, 4]]) + s_geta = np.asarray(scipy_sparse.csr_matrix(A[0])[[0, 0, 1, 3], [0, 1, 2, 4]]) + s_getb = np.asarray(scipy_sparse.csc_matrix(B[0])[[0, 0, 1, 3], [0, 1, 2, 4]]) utt.assert_allclose(t_geta, s_geta) utt.assert_allclose(t_getb, s_getb) @@ -2426,8 +1261,6 @@ def op_with_fixed_index(x): verify_grad_sparse(op_with_fixed_index, x_val) def test_GetItem2D(self): - is_supported_version = version.parse(sp.__version__) >= version.parse("0.14") - sparse_formats = ("csc", "csr") for format in sparse_formats: x = sparse.matrix(format, name="x") @@ -2443,12 +1276,8 @@ def test_GetItem2D(self): n = 5 p = 10 q = 15 - if is_supported_version: - j = 2 - k = 3 - else: - j = None - k = None + j = 2 + k = 3 vx = as_sparse_format(self.rng.binomial(1, 0.5, (100, 97)), format).astype( pytensor.config.floatX @@ -2457,14 +1286,10 @@ def test_GetItem2D(self): # mode_no_debug = pytensor.compile.mode.get_default_mode() # if isinstance(mode_no_debug, pytensor.compile.debugmode.DebugMode): # mode_no_debug = 'FAST_RUN' - if is_supported_version: - f1 = pytensor.function([x, a, b, c, d, e, f], x[a:b:e, c:d:f]) - r1 = f1(vx, m, n, p, q, j, k) - t1 = vx[m:n:j, p:q:k] - else: - f1 = pytensor.function([x, a, b, c, d], x[a:b, c:d]) - r1 = f1(vx, m, n, p, q) - t1 = vx[m:n, p:q] + f1 = pytensor.function([x, a, b, c, d, e, f], x[a:b:e, c:d:f]) + r1 = f1(vx, m, n, p, q, j, k) + t1 = vx[m:n:j, p:q:k] + assert r1.shape == t1.shape assert np.all(t1.toarray() == r1.toarray()) @@ -2499,14 +1324,10 @@ def test_GetItem2D(self): assert r7.shape == t7.shape assert np.all(r7.toarray() == t7.toarray()) """ - if is_supported_version: - f4 = pytensor.function([x, a, b, e], x[a:b:e]) - r4 = f4(vx, m, n, j) - t4 = vx[m:n:j] - else: - f4 = pytensor.function([x, a, b], x[a:b]) - r4 = f4(vx, m, n) - t4 = vx[m:n] + f4 = pytensor.function([x, a, b, e], x[a:b:e]) + r4 = f4(vx, m, n, j) + t4 = vx[m:n:j] + assert r4.shape == t4.shape assert np.all(t4.toarray() == r4.toarray()) @@ -2520,14 +1341,10 @@ def test_GetItem2D(self): # ---------------------------------------------------------- # test cases with indexing both with pytensor variable and int - if is_supported_version: - f8 = pytensor.function([x, a, b, e], x[a:b:e, 10:20:1]) - r8 = f8(vx, m, n, j) - t8 = vx[m:n:j, 10:20:1] - else: - f8 = pytensor.function([x, a, b], x[a:b, 10:20]) - r8 = f8(vx, m, n) - t8 = vx[m:n, 10:20] + f8 = pytensor.function([x, a, b, e], x[a:b:e, 10:20:1]) + r8 = f8(vx, m, n, j) + t8 = vx[m:n:j, 10:20:1] + assert r8.shape == t8.shape assert np.all(r8.toarray() == t8.toarray()) @@ -2567,24 +1384,6 @@ def test_GetItem2D(self): with pytest.raises(NotImplementedError): x.__getitem__((slice(a, b), c)) - # x[a:b:step, c:d] is not accepted because scipy silently drops - # the step (!) - if not is_supported_version: - with pytest.raises(ValueError): - x.__getitem__((slice(a, b, -1), slice(c, d))) - with pytest.raises(ValueError): - x.__getitem__((slice(a, b), slice(c, d, 2))) - - # Advanced indexing is not supported - with pytest.raises(ValueError): - x.__getitem__((ivector("l"), slice(a, b))) - - # Indexing with random things is not supported either - with pytest.raises(ValueError): - x.__getitem__(slice(fscalar("f"), None)) - with pytest.raises(ValueError): - x.__getitem__((slice(None), slice([1, 3, 4], None))) - def test_GetItemScalar(self): sparse_formats = ("csc", "csr") for format in sparse_formats: @@ -2701,7 +1500,7 @@ def _format_info(nb): for format in sparse.sparse_formats: variable = getattr(pytensor.sparse, format + "_matrix") - spa = getattr(sp.sparse, format + "_matrix") + spa = getattr(scipy_sparse, format + "_matrix") x[format] = [variable() for t in range(nb)] mat[format] = [ @@ -2719,9 +1518,9 @@ class _TestHVStack(utt.InferShapeTester): x, mat = _format_info(nb) def test_op(self): - for format in sparse.sparse_formats: - for out_f in sparse.sparse_formats: - for dtype in sparse.all_dtypes: + for format in sparse_formats: + for out_f in sparse_formats: + for dtype in all_dtypes: blocks = self.mat[format] f = pytensor.function( @@ -2779,222 +1578,8 @@ def expected_f(self, a, format=None, dtype=None): return TestXStack -TestHStack = _hv_switch(HStack, sp.sparse.hstack) -TestVStack = _hv_switch(VStack, sp.sparse.vstack) - - -class TestAddSSData(utt.InferShapeTester): - x = {} - a = {} - - def setup_method(self): - super().setup_method() - self.op_class = AddSSData - - for format in sparse.sparse_formats: - variable = getattr(pytensor.sparse, format + "_matrix") - - a_val = np.array( - np.random.default_rng(utt.fetch_seed()).integers(1, 4, size=(3, 4)) - 1, - dtype=pytensor.config.floatX, - ) - constant = as_sparse_format(a_val, format) - - self.x[format] = [variable() for t in range(2)] - self.a[format] = [constant for t in range(2)] - - def test_op(self): - for format in sparse.sparse_formats: - f = pytensor.function(self.x[format], add_s_s_data(*self.x[format])) - - tested = f(*self.a[format]) - expected = 2 * self.a[format][0] - - utt.assert_allclose(expected.toarray(), tested.toarray()) - assert tested.format == expected.format - assert tested.dtype == expected.dtype - - def test_infer_shape(self): - for format in sparse.sparse_formats: - self._compile_and_check( - self.x[format], - [add_s_s_data(*self.x[format])], - self.a[format], - self.op_class, - ) - - def test_grad(self): - for format in sparse.sparse_formats: - verify_grad_sparse(self.op_class(), self.a[format], structured=True) - - -def elemwise_checker( - op, expected_f, gap=None, test_dtypes=None, grad_test=True, name=None, gap_grad=None -): - """ - Return the appropriate test class for the elemwise on sparse. - - :param op: Op to test. - :expected_f: Function use to compare. This function must act - on dense matrix. If the op is structured - see the `structure_function` decorator to make - this function structured. - :param gap: Tuple for the range of the random sample. When - length is 1, it is assumed to be the exclusive - max, when `gap` = (`a`, `b`) it provide a sample - from [a, b[. If `None` is used, it provide [0, 1] - for float dtypes and [0, 50[ for integer dtypes. - :param test_dtypes: Particular dtypes for testing the op. - If `None`, this is set to the most common - dtypes. - :param grad_test: True for testing the grad. False will - skip this test. - :param gap_grad: If None, we reuse gap. Otherwise it is the same as gap - but for testing the gradiant of the op. - - :return: The class that perform the tests, not an instance - of the class. - """ - - if test_dtypes is None: - test_dtypes = sparse.all_dtypes - - class TestElemwise: - def setup_method(self): - super().setup_method() - self.op = op - self.expected_f = expected_f - self.gap = gap - if gap_grad is not None: - self.gap_grad = gap_grad - else: - self.gap_grad = gap - # Ensure the test's name is correct. - assert eval(self.__class__.__name__) is self.__class__ - - def test_op(self): - for format in sparse.sparse_formats: - for dtype in test_dtypes: - if dtype == "int8" or dtype == "uint8": - continue - - # When testing with unsigned integers, - # we must check if the gap contains - # negative numbers. - if dtype.startswith("uint"): - if self.gap and len(self.gap) == 2 and self.gap[0] < 0: - if self.gap[1] >= 1: - self.gap = (0, self.gap[1]) - else: - raise TypeError( - "Gap not suitable for", dtype, self.__name__ - ) - - variable, data = sparse_random_inputs( - format, shape=(4, 7), out_dtype=dtype, gap=self.gap - ) - - f = pytensor.function(variable, self.op(*variable)) - - tested = f(*data) - data = [m.toarray() for m in data] - expected = self.expected_f(*data) - - assert tested.format == format - tested = tested.toarray() - - try: - utt.assert_allclose(expected, tested) - except AssertionError: - raise AssertionError(self.__name__) - - # Test with int8 as dtype - # These tests are not in the loop for two reasons. - # First, in recent version of numpy, when a numpy - # function have int8 as input dtype, it returns a - # float16 as output dtype. Since this does not provide - # enough precision, we upcast the data before we apply the - # function. - # Second, the tolerance for the checkup in DebugMode - # is too high. - for dtype in ["int8", "uint8"]: - if dtype in test_dtypes: - if self.gap: - domain = self.gap - # When testing with unsigned integers, - # we must check if the gap contains - # negative numbers. - if dtype == "uint8": - if len(domain) == 2 and domain[0] < 0: - if domain[1] >= 1: - domain = (0, domain[1]) - else: - raise TypeError( - "Gap not suitable for", dtype, self.__name__ - ) - - else: - domain = (0, 5) - - variable, data = sparse_random_inputs( - format, shape=(4, 7), out_dtype=dtype, gap=domain - ) - - f = pytensor.function(variable, self.op(*variable)) - - old_value = ( - tensor.math.float32_atol, - tensor.math.float32_rtol, - tensor.math.float64_atol, - tensor.math.float64_rtol, - ) - tensor.math.float32_atol = 1e-4 - tensor.math.float32_rtol = 1e-3 - tensor.math.float64_atol = 1e-3 - tensor.math.float64_rtol = 1e-4 - try: - tested = f(*data) - finally: - ( - tensor.math.float32_atol, - tensor.math.float32_rtol, - tensor.math.float64_atol, - tensor.math.float64_rtol, - ) = old_value - - data = [m.toarray().astype("float32") for m in data] - expected = self.expected_f(*data) - - assert tested.format == format - tested = tested.toarray() - - try: - utt.assert_allclose(tested, expected, rtol=1e-2) - except AssertionError: - raise AssertionError(self.__name__) - - if grad_test: - - def test_grad(self): - for format in sparse.sparse_formats: - for dtype in sparse.float_dtypes: - _variable, data = sparse_random_inputs( - format, shape=(4, 7), out_dtype=dtype, gap=self.gap_grad - ) - - verify_grad_sparse(self.op, data, structured=True) - - # Set proper class name to uniquely identify tests. - # Note that it is important to run this code *outside* of the `Tester` - # class itself, otherwise it will not work properly for some reason. - if name is None: - name = op.__name__.capitalize() + "Tester" - TestElemwise.__name__ = name - if hasattr(TestElemwise, "__qualname__"): - TestElemwise.__qualname__ = name - assert "Roundhalftoeven" not in TestElemwise.__name__ - - return TestElemwise +TestHStack = _hv_switch(HStack, scipy_sparse.hstack) +TestVStack = _hv_switch(VStack, scipy_sparse.vstack) def test_hstack_vstack(): @@ -3021,399 +1606,3 @@ def get_expected_dtype(blocks, to_dtype): stacked_blocks = stack_function(blocks, dtype=to_dtype) expected_dtype = get_expected_dtype(blocks, to_dtype) assert stacked_blocks.dtype == expected_dtype - - -def structure_function(f, index=0): - """ - Decorator to structure a function which - apply on dense matrix. - - Here, the inputs of the function must be - dense matrix. The sparse pattern is - determined by finding the zeros. - - :param index: The index of the parameter - from which the function must - be structured. - - :return: The structured function for its - `index` parameter. - """ - - def structured_function(*args): - pattern = args[index] - evaluated = f(*args) - evaluated[pattern == 0] = 0 - return evaluated - - return structured_function - - -StructuredSigmoidTester = elemwise_checker( - sparse.structured_sigmoid, - structure_function(lambda x: 1.0 / (1.0 + np.exp(-x))), - test_dtypes=[ - m - for m in sparse.all_dtypes - if (m not in sparse.complex_dtypes and not m.startswith("uint")) - ], - gap=(-5, 5), - name="StructuredSigmoidTester", -) - -StructuredExpTester = elemwise_checker( - sparse.structured_exp, structure_function(np.exp), name="StructuredExpTester" -) - -StructuredLogTester = elemwise_checker( - sparse.structured_log, - structure_function(np.log), - gap=(0.5, 10), - name="StructuredLogTester", -) - -StructuredPowTester = elemwise_checker( - lambda x: sparse.structured_pow(x, 2), - structure_function(lambda x: np.power(x, 2)), - name="StructuredPowTester", -) - -StructuredMinimumTester = elemwise_checker( - lambda x: structured_minimum(x, 2), - structure_function(lambda x: np.minimum(x, 2)), - name="StructuredMinimumTester", -) - -StructuredMaximumTester = elemwise_checker( - lambda x: structured_maximum(x, 2), - structure_function(lambda x: np.maximum(x, 2)), - name="StructuredMaximumTester", -) - -StructuredAddTester = elemwise_checker( - lambda x: structured_add(x, 2), - structure_function(lambda x: np.add(x, 2)), - name="StructuredAddTester", -) - -SinTester = elemwise_checker(sparse.sin, np.sin) - -TanTester = elemwise_checker(sparse.tan, np.tan, gap=(-1, 1)) - -ArcsinTester = elemwise_checker( - sparse.arcsin, np.arcsin, gap=(-1, 1), gap_grad=(-0.99, 0.99) -) - -ArctanTester = elemwise_checker(sparse.arctan, np.arctan) - -SinhTester = elemwise_checker(sparse.sinh, np.sinh) - -ArcsinhTester = elemwise_checker(sparse.arcsinh, np.arcsinh, gap=(-1, 1)) - -TanhTester = elemwise_checker(sparse.tanh, np.tanh, gap=(-1, 1)) - -ArctanhTester = elemwise_checker( - sparse.arctanh, np.arctanh, gap=(-0.9, 1), gap_grad=(-0.9, 0.95) -) - -RintTester = elemwise_checker( - sparse.rint, np.rint, grad_test=False, test_dtypes=sparse.float_dtypes -) - -SgnTester = elemwise_checker( - sparse.sign, - np.sign, - grad_test=False, - test_dtypes=[ - m - for m in sparse.all_dtypes - if (m not in sparse.complex_dtypes and not m.startswith("uint")) - ], -) - -CeilTester = elemwise_checker( - sparse.ceil, - np.ceil, - grad_test=False, - test_dtypes=[m for m in sparse.all_dtypes if m not in sparse.complex_dtypes], -) - -FloorTester = elemwise_checker( - sparse.floor, - np.floor, - grad_test=False, - test_dtypes=[m for m in sparse.all_dtypes if m not in sparse.complex_dtypes], -) - -Log1pTester = elemwise_checker(sparse.log1p, np.log1p, gap=(0.5, 10)) - -Expm1Tester = elemwise_checker(sparse.expm1, np.expm1) - -Deg2radTester = elemwise_checker( - sparse.deg2rad, - np.deg2rad, - test_dtypes=[m for m in sparse.all_dtypes if m not in sparse.complex_dtypes], -) - -Rad2degTester = elemwise_checker( - sparse.rad2deg, - np.rad2deg, - test_dtypes=[m for m in sparse.all_dtypes if m not in sparse.complex_dtypes], -) - - -TruncTester = elemwise_checker( - sparse.trunc, - np.trunc, - test_dtypes=[m for m in sparse.all_dtypes if m not in sparse.complex_dtypes], - grad_test=False, -) - - -SqrTester = elemwise_checker(sparse.sqr, lambda x: x * x) - -SqrtTester = elemwise_checker(sparse.sqrt, np.sqrt, gap=(0, 10)) - -ConjTester = elemwise_checker(sparse.conj, np.conj, grad_test=False) - - -def test_useless_conj(): - x = sparse.SparseTensorType("csr", dtype="complex128")() - assert x.conj() is not x - - # No conjugate when the data type isn't complex - x = sparse.SparseTensorType("csr", dtype="float64")() - assert x.conj() is x - - -class TestMulSV: - def test_mul_s_v_grad(self): - sp_types = {"csc": sp.sparse.csc_matrix, "csr": sp.sparse.csr_matrix} - - for format in ("csr", "csc"): - for dtype in ("float32", "float64"): - spmat = sp_types[format](random_lil((4, 3), dtype, 3)) - mat = np.asarray(np.random.random(3), dtype=dtype) - - verify_grad_sparse(mul_s_v, [spmat, mat], structured=True) - - def test_mul_s_v(self): - sp_types = {"csc": sp.sparse.csc_matrix, "csr": sp.sparse.csr_matrix} - - for format in ("csr", "csc"): - for dtype in ("float32", "float64"): - x = sparse.SparseTensorType(format, dtype=dtype)() - y = vector(dtype=dtype) - f = pytensor.function([x, y], mul_s_v(x, y)) - - spmat = sp_types[format](random_lil((4, 3), dtype, 3)) - mat = np.asarray(np.random.random(3), dtype=dtype) - - out = f(spmat, mat) - - utt.assert_allclose(spmat.toarray() * mat, out.toarray()) - - -class TestStructuredAddSV: - def test_structured_add_s_v_grad(self): - sp_types = {"csc": sp.sparse.csc_matrix, "csr": sp.sparse.csr_matrix} - - for format in ("csr", "csc"): - for dtype in ("float32", "float64"): - spmat = sp_types[format](random_lil((4, 3), dtype, 3)) - mat = np.asarray(np.random.random(3), dtype=dtype) - - verify_grad_sparse(structured_add_s_v, [spmat, mat], structured=True) - - def test_structured_add_s_v(self): - sp_types = {"csc": sp.sparse.csc_matrix, "csr": sp.sparse.csr_matrix} - - for format in ("csr", "csc"): - for dtype in ("float32", "float64"): - x = sparse.SparseTensorType(format, dtype=dtype)() - y = vector(dtype=dtype) - f = pytensor.function([x, y], structured_add_s_v(x, y)) - - spmat = sp_types[format](random_lil((4, 3), dtype, 3)) - spones = spmat.copy() - spones.data = np.ones_like(spones.data) - mat = np.asarray(np.random.random(3), dtype=dtype) - - out = f(spmat, mat) - - utt.assert_allclose( - as_ndarray(spones.multiply(spmat + mat)), out.toarray() - ) - - -class TestTrueDot(utt.InferShapeTester): - def setup_method(self): - super().setup_method() - self.op = true_dot - self.op_class = TrueDot - - def test_op_ss(self): - for format in sparse.sparse_formats: - for dtype in sparse.all_dtypes: - variable, data = sparse_random_inputs( - format, shape=(10, 10), out_dtype=dtype, n=2, p=0.1 - ) - - f = pytensor.function(variable, self.op(*variable)) - - tested = f(*data) - - x, y = (m.toarray() for m in data) - expected = np.dot(x, y) - - assert tested.format == format - assert tested.dtype == expected.dtype - tested = tested.toarray() - utt.assert_allclose(tested, expected) - - def test_op_sd(self): - for format in sparse.sparse_formats: - for dtype in sparse.all_dtypes: - variable, data = sparse_random_inputs( - format, shape=(10, 10), out_dtype=dtype, n=2, p=0.1 - ) - variable[1] = TensorType(dtype=dtype, shape=(None, None))() - data[1] = data[1].toarray() - - f = pytensor.function(variable, self.op(*variable)) - - tested = f(*data) - expected = np.dot(data[0].toarray(), data[1]) - - assert tested.format == format - assert tested.dtype == expected.dtype - tested = tested.toarray() - utt.assert_allclose(tested, expected) - - def test_infer_shape(self): - for format in sparse.sparse_formats: - for dtype in sparse.all_dtypes: - (x,), (x_value,) = sparse_random_inputs( - format, shape=(9, 10), out_dtype=dtype, p=0.1 - ) - (y,), (y_value,) = sparse_random_inputs( - format, shape=(10, 24), out_dtype=dtype, p=0.1 - ) - variable = [x, y] - data = [x_value, y_value] - self._compile_and_check( - variable, [self.op(*variable)], data, self.op_class - ) - - def test_grad(self): - for format in sparse.sparse_formats: - for dtype in sparse.float_dtypes: - (_x,), (x_value,) = sparse_random_inputs( - format, shape=(9, 10), out_dtype=dtype, p=0.1 - ) - (_y,), (y_value,) = sparse_random_inputs( - format, shape=(10, 24), out_dtype=dtype, p=0.1 - ) - data = [x_value, y_value] - verify_grad_sparse(self.op, data, structured=False) - - -class TestSamplingDot(utt.InferShapeTester): - x = [matrix() for t in range(2)] - x.append(sparse.csr_matrix()) - # unsquare shape - a = [ - np.array( - np.random.default_rng().integers(1, 6, size=(4, 3)) - 1, - dtype=pytensor.config.floatX, - ), - np.array( - np.random.default_rng().integers(1, 6, size=(5, 3)) - 1, - dtype=pytensor.config.floatX, - ), - np.array( - np.random.default_rng().integers(1, 3, size=(4, 5)) - 1, - dtype=pytensor.config.floatX, - ), - ] - a[2] = sp.sparse.csr_matrix(a[2]) - - def setup_method(self): - super().setup_method() - self.op_class = SamplingDot - - def test_op(self): - f = pytensor.function(self.x, sampling_dot(*self.x)) - - tested = f(*self.a) - x, y, p = self.a - expected = p.multiply(np.dot(x, y.T)) - - utt.assert_allclose(as_ndarray(expected), tested.toarray()) - assert tested.format == "csr" - assert tested.dtype == expected.dtype - - def test_negative_stride(self): - f = pytensor.function(self.x, sampling_dot(*self.x)) - - a2 = [self.a[0][::-1, :], self.a[1][:, ::-1], self.a[2]] - tested = f(*a2) - x, y, p = a2 - expected = p.multiply(np.dot(x, y.T)) - - utt.assert_allclose(as_ndarray(expected), tested.toarray()) - assert tested.format == "csr" - assert tested.dtype == expected.dtype - - def test_infer_shape(self): - self._compile_and_check( - self.x, - [sampling_dot(*self.x)], - self.a, - self.op_class, - excluding=["local_sampling_dot_csr"], - ) - - def test_grad(self): - def _helper(x, y): - return sampling_dot(x, y, self.a[2]) - - verify_grad_sparse(_helper, self.a[:2]) - - -@makeSharedTester( - shared_constructor_=sparse.shared, - dtype_="float64", - get_value_borrow_true_alias_=True, - shared_borrow_true_alias_=True, - set_value_borrow_true_alias_=True, - set_value_inplace_=False, - set_cast_value_inplace_=False, - shared_constructor_accept_ndarray_=False, - internal_type_=sp.sparse.csc_matrix, - check_internal_type_=sp.sparse.issparse, - pytensor_fct_=lambda a: dense_from_sparse(a * 2.0), - ref_fct_=lambda a: np.asarray((a * 2).todense()), - cast_value_=sp.sparse.csr_matrix, - expect_fail_fast_shape_inplace=False, -) -class TestSharedOptions: - pass - - -@pytest.mark.parametrize("format", ["csc", "csr"], ids=["csc", "csr"]) -@pytest.mark.parametrize("sparse_input", [True, False], ids=["sparse", "dense"]) -def test_block_diagonal(format, sparse_input): - from scipy import sparse as sp_sparse - - f_array = sp_sparse.csr_matrix if sparse_input else np.array - A = f_array([[1, 2], [3, 4]]).astype(config.floatX) - B = f_array([[5, 6], [7, 8]]).astype(config.floatX) - - result = block_diag(A, B, format=format) - assert result.owner.op._props_dict() == {"n_inputs": 2, "format": format} - - sp_result = sp_sparse.block_diag([A, B], format=format) - - assert isinstance(result.eval(), type(sp_result)) - np.testing.assert_allclose(result.eval().toarray(), sp_result.toarray()) diff --git a/tests/sparse/test_linalg.py b/tests/sparse/test_linalg.py new file mode 100644 index 0000000000..7f1360287a --- /dev/null +++ b/tests/sparse/test_linalg.py @@ -0,0 +1,26 @@ +import numpy as np +import pytest + +from pytensor.configdefaults import config +from pytensor.sparse.linalg import block_diag + + +sp = pytest.importorskip("scipy", minversion="0.7.0") + + +@pytest.mark.parametrize("format", ["csc", "csr"], ids=["csc", "csr"]) +@pytest.mark.parametrize("sparse_input", [True, False], ids=["sparse", "dense"]) +def test_block_diagonal(format, sparse_input): + from scipy import sparse as sp_sparse + + f_array = sp_sparse.csr_matrix if sparse_input else np.array + A = f_array([[1, 2], [3, 4]]).astype(config.floatX) + B = f_array([[5, 6], [7, 8]]).astype(config.floatX) + + result = block_diag(A, B, format=format) + assert result.owner.op._props_dict() == {"n_inputs": 2, "format": format} + + sp_result = sp_sparse.block_diag([A, B], format=format) + + assert isinstance(result.eval(), type(sp_result)) + np.testing.assert_allclose(result.eval().toarray(), sp_result.toarray()) diff --git a/tests/sparse/test_math.py b/tests/sparse/test_math.py new file mode 100644 index 0000000000..811d80e808 --- /dev/null +++ b/tests/sparse/test_math.py @@ -0,0 +1,1462 @@ +import time +from itertools import product + +import numpy as np +import pytest +import scipy.sparse as scipy_sparse + +import pytensor +import pytensor.sparse.math as psm +import pytensor.tensor as pt +from pytensor.configdefaults import config +from pytensor.sparse.basic import ( + CSR, + CSMProperties, + SparseTensorType, + _is_dense_variable, + _is_sparse, + _is_sparse_variable, + _mtypes, + all_dtypes, + as_sparse_variable, + complex_dtypes, + csc_matrix, + csr_matrix, + float_dtypes, + sparse_formats, +) +from pytensor.sparse.math import ( + Dot, + SamplingDot, + StructuredDot, + TrueDot, + Usmm, + add, + ge, + gt, + le, + lt, + mul_s_v, + multiply, + sampling_dot, + structured_add_s_v, + structured_dot, + true_dot, +) +from pytensor.sparse.rewriting import UsmmCscDense +from pytensor.tensor.elemwise import DimShuffle, Elemwise +from pytensor.tensor.subtensor import Subtensor +from pytensor.tensor.type import ( + TensorType, + matrix, + scalar, + vector, +) +from tests import unittest_tools as utt +from tests.sparse.test_basic import ( + as_ndarray, + as_sparse_format, + random_lil, + sparse_random_inputs, + verify_grad_sparse, +) + + +def eval_outputs(outputs): + return pytensor.function([], outputs)()[0] + + +class TestAddMul: + def test_AddSS(self): + self._testSS(add) + + def test_AddSD(self): + self._testSD(add) + + def test_AddDS(self): + self._testDS(add) + + def test_MulSS(self): + self._testSS( + multiply, + np.array([[1.0, 0], [3, 0], [0, 6]]), + np.array([[1.0, 2], [3, 0], [0, 6]]), + ) + + def test_MulSD(self): + self._testSD( + multiply, + np.array([[1.0, 0], [3, 0], [0, 6]]), + np.array([[1.0, 2], [3, 0], [0, 6]]), + ) + + def test_MulDS(self): + self._testDS( + multiply, + np.array([[1.0, 0], [3, 0], [0, 6]]), + np.array([[1.0, 2], [3, 0], [0, 6]]), + ) + + def _testSS(self, op, array1=None, array2=None): + if array1 is None: + array1 = np.array([[1.0, 0], [3, 0], [0, 6]]) + if array2 is None: + array2 = np.asarray([[0, 2.0], [0, 4], [5, 0]]) + + for mtype1, mtype2 in product(_mtypes, _mtypes): + for dtype1, dtype2 in [ + ("float64", "int8"), + ("int8", "float64"), + ("float64", "float64"), + ]: + a = mtype1(array1).astype(dtype1) + aR = as_sparse_variable(a) + assert aR.data is not a + assert _is_sparse(a) + assert _is_sparse_variable(aR) + + b = mtype2(array2).astype(dtype2) + bR = as_sparse_variable(b) + assert bR.data is not b + assert _is_sparse(b) + assert _is_sparse_variable(bR) + + apb = op(aR, bR) + assert _is_sparse_variable(apb) + + assert apb.type.format == aR.type.format, apb.type.format + + val = eval_outputs([apb]) + assert val.shape == (3, 2) + if op is add: + assert np.all(val.todense() == array1 + array2) + if dtype1.startswith("float") and dtype2.startswith("float"): + verify_grad_sparse(op, [a, b], structured=False) + elif op is multiply: + assert np.all(val.todense() == array1 * array2) + if dtype1.startswith("float") and dtype2.startswith("float"): + verify_grad_sparse(op, [a, b], structured=False) + + def _testSD(self, op, array1=None, array2=None): + if array1 is None: + array1 = np.array([[1.0, 0], [3, 0], [0, 6]]) + if array2 is None: + array2 = np.asarray([[0, 2.0], [0, 4], [5, 0]]) + + for mtype in _mtypes: + for a in [ + np.array(array1), + pt.as_tensor_variable(array1), + pytensor.shared(array1), + ]: + for dtype1, dtype2 in [ + ("float64", "int8"), + ("int8", "float64"), + # Needed to test the grad + ("float32", "float64"), + ]: + a = a.astype(dtype1) + b = mtype(array2).astype(dtype2) + bR = as_sparse_variable(b) + assert bR.data is not b # constants are copied + assert _is_sparse(b) + assert _is_sparse_variable(bR) + + apb = op(a, bR) + + val = eval_outputs([apb]) + assert val.shape == (3, 2) + if op is add: + assert _is_dense_variable(apb) + assert np.all(val == array1 + b) + ans = np.array([[1.0, 2], [3, 4], [5, 6]]) + assert np.all(val == ans) + if hasattr(a, "owner") and a.owner is not None: + continue + if dtype1.startswith("float") and dtype2.startswith("float"): + verify_grad_sparse(op, [a, b], structured=True) + elif op is multiply: + assert _is_sparse_variable(apb) + assert np.all(val.todense() == b.multiply(array1)) + assert np.all( + val.todense() == np.array([[1, 0], [9, 0], [0, 36]]) + ) + if hasattr(a, "owner") and a.owner is not None: + continue + if dtype1.startswith("float") and dtype2.startswith("float"): + verify_grad_sparse(op, [a, b], structured=False) + + def _testDS(self, op, array1=None, array2=None): + if array1 is None: + array1 = np.array([[1.0, 0], [3, 0], [0, 6]]) + if array2 is None: + array2 = np.asarray([[0, 2.0], [0, 4], [5, 0]]) + + for mtype in _mtypes: + for b in [ + np.asarray(array2), + pt.as_tensor_variable(array2), + pytensor.shared(array2), + ]: + for dtype1, dtype2 in [ + ("float64", "int8"), + ("int8", "float64"), + ]: + a = mtype(array1).astype(dtype1) + aR = as_sparse_variable(a) + assert aR.data is not a + assert _is_sparse(a) + assert _is_sparse_variable(aR) + b = b.astype(dtype2) + + apb = op(aR, b) + + val = eval_outputs([apb]) + assert val.shape == (3, 2) + if op is add: + assert _is_dense_variable(apb) + assert np.all(val == a + array2) + ans = np.array([[1.0, 2], [3, 4], [5, 6]]) + assert np.all(val == ans) + if dtype1.startswith("float") and dtype2.startswith("float"): + verify_grad_sparse(op, [a, b], structured=True) + elif op is multiply: + assert _is_sparse_variable(apb) + ans = np.array([[1, 0], [9, 0], [0, 36]]) + assert np.all(val.todense() == (a.multiply(array2))) + assert np.all(val.todense() == ans) + if dtype1.startswith("float") and dtype2.startswith("float"): + verify_grad_sparse(op, [a, b], structured=False) + + +class TestComparison: + # took from tensor basic_test.py + def _rand_ranged(self, min, max, shape): + return np.asarray( + np.random.random(shape) * (max - min) + min, dtype=config.floatX + ) + + tests = [ + lambda x, y: x > y, + lambda x, y: x < y, + lambda x, y: x >= y, + lambda x, y: x <= y, + ] + + testsDic = { + gt: lambda x, y: x > y, + lt: lambda x, y: x < y, + ge: lambda x, y: x >= y, + le: lambda x, y: x <= y, + } + + def __generalized_ss_test(self, pytensorp, symbolicType, testOp, scipyType): + x = symbolicType() + y = symbolicType() + + op = pytensorp(x, y) + f = pytensor.function([x, y], op) + + m1 = scipyType(random_lil((10, 40), config.floatX, 3)) + m2 = scipyType(random_lil((10, 40), config.floatX, 3)) + + assert np.array_equal(f(m1, m2).data, testOp(m1, m2).data) + + def __generalized_sd_test(self, pytensorp, symbolicType, testOp, scipyType): + x = symbolicType() + y = matrix() + + op = pytensorp(x, y) + f = pytensor.function([x, y], op) + + m1 = scipyType(random_lil((10, 40), config.floatX, 3)) + m2 = self._rand_ranged(1000, -1000, [10, 40]) + + assert np.array_equal(f(m1, m2).data, testOp(m1, m2).data) + + def __generalized_ds_test(self, pytensorp, symbolicType, testOp, scipyType): + x = symbolicType() + y = matrix() + + op = pytensorp(y, x) + f = pytensor.function([y, x], op) + + m1 = scipyType(random_lil((10, 40), config.floatX, 3)) + m2 = self._rand_ranged(1000, -1000, [10, 40]) + + assert np.array_equal(f(m2, m1).data, testOp(m2, m1).data) + + def test_ss_csr_comparison(self): + for op in self.tests: + self.__generalized_ss_test(op, csr_matrix, op, scipy_sparse.csr_matrix) + + def test_ss_csc_comparison(self): + for op in self.tests: + self.__generalized_ss_test(op, csc_matrix, op, scipy_sparse.csc_matrix) + + def test_sd_csr_comparison(self): + for op in self.tests: + self.__generalized_sd_test(op, csr_matrix, op, scipy_sparse.csr_matrix) + + def test_sd_csc_comparison(self): + for op in self.tests: + self.__generalized_sd_test(op, csc_matrix, op, scipy_sparse.csc_matrix) + + def test_ds_csc_comparison(self): + for op in self.testsDic: + self.__generalized_ds_test( + op, + csc_matrix, + self.testsDic[op], + scipy_sparse.csc_matrix, + ) + + def test_ds_csr_comparison(self): + for op in self.testsDic: + self.__generalized_ds_test( + op, + csr_matrix, + self.testsDic[op], + scipy_sparse.csr_matrix, + ) + + def test_equality_case(self): + # Test assuring normal behaviour when values in the matrices are equal + x = csc_matrix() + y = matrix() + + m1 = scipy_sparse.csc_matrix((2, 2), dtype=pytensor.config.floatX) + m2 = np.asarray([[0, 0], [0, 0]], dtype=pytensor.config.floatX) + + for func in self.testsDic: + op = func(y, x) + f = pytensor.function([y, x], op) + assert np.array_equal(f(m2, m1), self.testsDic[func](m2, m1)) + + +class TestStructuredDot: + def test_structureddot_csc_grad(self): + spmat = scipy_sparse.csc_matrix(random_lil((4, 3), "float32", 3)) + mat = np.asarray(np.random.standard_normal((3, 2)), "float32") + verify_grad_sparse(structured_dot, [spmat, mat], structured=True) + + def buildgraph_T(spmat, mat): + return structured_dot(mat.T, spmat.T) + + verify_grad_sparse(buildgraph_T, [spmat, mat], structured=True) + + def test_structureddot_csr_grad(self): + spmat = scipy_sparse.csr_matrix(random_lil((4, 3), "float64", 3)) + mat = np.asarray(np.random.standard_normal((3, 2)), "float64") + verify_grad_sparse(structured_dot, [spmat, mat], structured=True) + + def buildgraph_T(spmat, mat): + return structured_dot(mat.T, spmat.T) + + verify_grad_sparse(buildgraph_T, [spmat, mat], structured=True) + + def test_upcast(self): + typenames = ( + "float32", + "int64", + "int8", + "int32", + "int16", + "float64", + "complex64", + "complex128", + ) + for dense_dtype in typenames: + for sparse_dtype in typenames: + correct_dtype = pytensor.scalar.upcast(sparse_dtype, dense_dtype) + a = SparseTensorType("csc", dtype=sparse_dtype)() + b = matrix(dtype=dense_dtype) + d = structured_dot(a, b) + assert d.type.dtype == correct_dtype + + f = pytensor.function([a, b], d) + + M, N, K, nnz = (4, 3, 5, 3) + spmat = scipy_sparse.csc_matrix(random_lil((M, N), sparse_dtype, nnz)) + spmat.dtype = np.dtype(sparse_dtype) + mat = np.asarray( + np.random.standard_normal((N, K)) * 9, dtype=dense_dtype + ) + pytensor_result = f(spmat, mat) + scipy_result = spmat * mat + assert pytensor_result.shape == scipy_result.shape + assert pytensor_result.dtype == scipy_result.dtype + utt.assert_allclose(scipy_result, pytensor_result) + + @pytest.mark.xfail( + reason="The optimization from structured_dot -> structured_dot_csc is currently disabled", + ) + def test_opt_unpack(self): + """ + Test that a graph involving structured_dot(assembled_csc_matrix) is optimized to be just a structured_dot_csc + Op and no assembly of a csc_matrix. + """ + + kerns = TensorType(dtype="int64", shape=(None,))("kerns") + spmat = scipy_sparse.random_array(shape=(4, 6), density=0.2, format="csc") + + images = TensorType(dtype="float32", shape=(None, None))("images") + + cscmat = pytensor.sparse.CSC( + kerns, spmat.indices[: spmat.size], spmat.indptr, spmat.shape + ) + f = pytensor.function([kerns, images], structured_dot(cscmat, images.T)) + + sdcscpresent = False + for node in f.maker.fgraph.toposort(): + # print node.op + assert not isinstance(node.op, pytensor.sparse.CSM) + assert not isinstance(node.op, CSMProperties) + if isinstance( + f.maker.fgraph.toposort()[1].op, + pytensor.sparse.rewriting.StructuredDotCSC, + ): + sdcscpresent = True + assert sdcscpresent + + kernvals = np.array(spmat.data[: spmat.size]) + bsize = 3 + imvals = 1.0 * np.array( + np.arange(bsize * spmat.shape[1]).reshape(bsize, spmat.shape[1]), + dtype="float32", + ) + f(kernvals, imvals) + + def test_dot_sparse_sparse(self): + sparse_dtype = "float64" + sp_mat = { + "csc": scipy_sparse.csc_matrix, + "csr": scipy_sparse.csr_matrix, + "bsr": scipy_sparse.csr_matrix, + } + + for sparse_format_a in ["csc", "csr", "bsr"]: + for sparse_format_b in ["csc", "csr", "bsr"]: + a = SparseTensorType(sparse_format_a, dtype=sparse_dtype)() + b = SparseTensorType(sparse_format_b, dtype=sparse_dtype)() + d = pt.dot(a, b) + f = pytensor.function([a, b], d) + for M, N, K, nnz in [ + (4, 3, 2, 3), + (40, 30, 20, 3), + (40, 30, 20, 30), + (400, 3000, 200, 6000), + ]: + a_val = sp_mat[sparse_format_a]( + random_lil((M, N), sparse_dtype, nnz) + ) + b_val = sp_mat[sparse_format_b]( + random_lil((N, K), sparse_dtype, nnz) + ) + f(a_val, b_val) + + def test_tensor_dot_types(self): + x = csc_matrix("x") + x_d = pt.matrix("x_d") + y = csc_matrix("y") + + res = pt.dot(x, y) + op_types = { + type(n.op) for n in pytensor.graph.traversal.applys_between([x, y], [res]) + } + assert StructuredDot in op_types + assert pt.math.Dot not in op_types + + res = pt.dot(x_d, y) + op_types = { + type(n.op) for n in pytensor.graph.traversal.applys_between([x, y], [res]) + } + assert StructuredDot in op_types + assert pt.math.Dot not in op_types + + res = pt.dot(x, x_d) + op_types = { + type(n.op) for n in pytensor.graph.traversal.applys_between([x, y], [res]) + } + assert StructuredDot in op_types + assert pt.math.Dot not in op_types + + res = pt.dot(pt.second(1, x), y) + op_types = { + type(n.op) for n in pytensor.graph.traversal.applys_between([x, y], [res]) + } + assert StructuredDot in op_types + assert pt.math.Dot not in op_types + + def test_csc_correct_output_faster_than_scipy(self): + sparse_dtype = "float64" + dense_dtype = "float64" + + a = SparseTensorType("csc", dtype=sparse_dtype)() + b = matrix(dtype=dense_dtype) + d = pt.dot(a, b) + f = pytensor.function([a, b], d) + + for M, N, K, nnz in [ + (4, 3, 2, 3), + (40, 30, 20, 3), + (40, 30, 20, 30), + (400, 3000, 200, 6000), + ]: + spmat = scipy_sparse.csc_matrix(random_lil((M, N), sparse_dtype, nnz)) + mat = np.asarray(np.random.standard_normal((N, K)), dense_dtype) + pytensor_times = [] + scipy_times = [] + for i in range(5): + t0 = time.perf_counter() + pytensor_result = f(spmat, mat) + t1 = time.perf_counter() + scipy_result = spmat * mat + t2 = time.perf_counter() + + pytensor_times.append(t1 - t0) + scipy_times.append(t2 - t1) + + pytensor_time = np.min(pytensor_times) + scipy_time = np.min(scipy_times) + + # speedup = scipy_time / pytensor_time + # print scipy_times + # print pytensor_times + # print ('M=%(M)s N=%(N)s K=%(K)s nnz=%(nnz)s pytensor_time' + # '=%(pytensor_time)s speedup=%(speedup)s') % locals() + + # fail if PyTensor is slower than scipy by more than a certain amount + overhead_tol = 0.003 # seconds overall + overhead_rtol = 1.2 # times as long + utt.assert_allclose(scipy_result, pytensor_result) + if pytensor.config.mode == "FAST_RUN" and pytensor.config.cxx: + assert pytensor_time <= overhead_rtol * scipy_time + overhead_tol + + def test_csr_correct_output_faster_than_scipy(self): + sparse_dtype = "float32" + dense_dtype = "float32" + + a = SparseTensorType("csr", dtype=sparse_dtype)() + b = matrix(dtype=dense_dtype) + d = pt.dot(a, b) + f = pytensor.function([a, b], d) + + for M, N, K, nnz in [ + (4, 3, 2, 3), + (40, 30, 20, 3), + (40, 30, 20, 30), + (400, 3000, 200, 6000), + ]: + spmat = scipy_sparse.csr_matrix(random_lil((M, N), sparse_dtype, nnz)) + mat = np.asarray(np.random.standard_normal((N, K)), dense_dtype) + t0 = time.perf_counter() + pytensor_result = f(spmat, mat) + t1 = time.perf_counter() + scipy_result = spmat * mat + t2 = time.perf_counter() + + pytensor_time = t1 - t0 + scipy_time = t2 - t1 + # print 'pytensor took', pytensor_time, + # print 'scipy took', scipy_time + overhead_tol = 0.002 # seconds + overhead_rtol = 1.1 # times as long + utt.assert_allclose(scipy_result, pytensor_result) + if pytensor.config.mode == "FAST_RUN" and pytensor.config.cxx: + assert pytensor_time <= overhead_rtol * scipy_time + overhead_tol, ( + pytensor_time, + overhead_rtol * scipy_time + overhead_tol, + scipy_time, + overhead_rtol, + overhead_tol, + ) + + +class TestDots(utt.InferShapeTester): + def setup_method(self): + super().setup_method() + x_size = (10, 100) + y_size = (100, 1000) + + self.x_csr = scipy_sparse.csr_matrix( + np.random.binomial(1, 0.5, x_size), dtype=pytensor.config.floatX + ) + self.x_csc = scipy_sparse.csc_matrix( + np.random.binomial(1, 0.5, x_size), dtype=pytensor.config.floatX + ) + self.y = np.asarray( + np.random.uniform(-1, 1, y_size), dtype=pytensor.config.floatX + ) + self.y_csr = scipy_sparse.csr_matrix( + np.random.binomial(1, 0.5, y_size), dtype=pytensor.config.floatX + ) + self.y_csc = scipy_sparse.csc_matrix( + np.random.binomial(1, 0.5, y_size), dtype=pytensor.config.floatX + ) + self.v_10 = np.asarray( + np.random.uniform(-1, 1, 10), dtype=pytensor.config.floatX + ) + self.v_100 = np.asarray( + np.random.uniform(-1, 1, 100), dtype=pytensor.config.floatX + ) + + def test_csr_dense(self): + x = csr_matrix("x") + y = matrix("y") + v = vector("v") + + for x, y, x_v, y_v in [ + (x, y, self.x_csr, self.y), + (x, v, self.x_csr, self.v_100), + (v, x, self.v_10, self.x_csr), + ]: + f_a = pytensor.function([x, y], psm.dot(x, y)) + + def f_b(x, y): + return x * y + + utt.assert_allclose(f_a(x_v, y_v), f_b(x_v, y_v)) + + # Test infer_shape + self._compile_and_check( + [x, y], [psm.dot(x, y)], [x_v, y_v], (Dot, Usmm, UsmmCscDense) + ) + + def test_csc_dense(self): + x = csc_matrix("x") + y = matrix("y") + v = vector("v") + + for x, y, x_v, y_v in [ + (x, y, self.x_csc, self.y), + (x, v, self.x_csc, self.v_100), + (v, x, self.v_10, self.x_csc), + ]: + f_a = pytensor.function([x, y], psm.dot(x, y)) + + def f_b(x, y): + return x * y + + utt.assert_allclose(f_a(x_v, y_v), f_b(x_v, y_v)) + + self._compile_and_check( + [x, y], [psm.dot(x, y)], [x_v, y_v], (Dot, Usmm, UsmmCscDense) + ) + + def test_sparse_sparse(self): + for d1, d2 in [ + ("float32", "float32"), + ("float32", "float64"), + ("float64", "float32"), + ("float64", "float64"), + ("float32", "int16"), + ("float32", "complex64"), + ]: + for x_f, y_f in [ + ("csc", "csc"), + ("csc", "csr"), + ("csr", "csc"), + ("csr", "csr"), + ]: + x = SparseTensorType(format=x_f, dtype=d1)("x") + y = SparseTensorType(format=x_f, dtype=d2)("x") + + def f_a(x, y): + return x * y + + f_b = pytensor.function([x, y], psm.dot(x, y)) + + vx = getattr(self, "x_" + x_f).astype(d1) + vy = getattr(self, "y_" + y_f).astype(d2) + utt.assert_allclose(f_a(vx, vy).toarray(), f_b(vx, vy)) + + f_a = pytensor.function([x, y], psm.dot(x, y).shape) + + def f_b2(x, y): + return (x * y).shape + + assert np.all(f_a(vx, vy) == f_b2(vx, vy)) + topo = f_a.maker.fgraph.toposort() + assert not any( + isinstance(node.op, Dot | Usmm | UsmmCscDense) for node in topo + ) + + def test_int32_dtype(self): + size = 9 + intX = "int32" + + C = matrix("C", dtype=intX) + I = matrix("I", dtype=intX) + + fI = I.flatten() + data = pt.ones_like(fI) + indptr = pt.arange(data.shape[0] + 1, dtype="int32") + + m1 = CSR(data, fI, indptr, (8, size)) + m2 = psm.dot(m1, C) + y = m2.reshape(shape=(2, 4, 9), ndim=3) + + f = pytensor.function(inputs=[I, C], outputs=y) + i = np.asarray([[4, 3, 7, 7], [2, 8, 4, 5]], dtype=intX) + a = np.asarray( + np.random.default_rng().integers(0, 100, (size, size)), dtype=intX + ) + f(i, a) + + def test_csr_dense_grad(self): + # shortcut: testing csc in float32, testing csr in float64 + + # allocate a random sparse matrix + spmat = scipy_sparse.csr_matrix(random_lil((4, 3), "float64", 3)) + mat = np.asarray(np.random.standard_normal((2, 4)), "float64") + + def buildgraph_T(mat): + return Dot()(mat, spmat) + + utt.verify_grad(buildgraph_T, [mat]) + + +class TestUsmm: + def setup_method(self): + x_size = (10, 100) + y_size = (100, 200) + z_size = (x_size[0], y_size[1]) + + self.rng = np.random.default_rng(seed=utt.fetch_seed()) + self.x = np.asarray( + self.rng.binomial(1, 0.5, x_size), dtype=pytensor.config.floatX + ) + self.y = np.asarray( + self.rng.uniform(-1, 1, y_size), dtype=pytensor.config.floatX + ) + self.z = np.asarray( + self.rng.uniform(-1, 1, z_size), dtype=pytensor.config.floatX + ) + + @pytest.mark.slow + @pytest.mark.parametrize("dtype1", ["float32", "float64", "int16", "complex64"]) + @pytest.mark.parametrize("dtype2", ["float32", "float64", "int16", "complex64"]) + @pytest.mark.parametrize("dtype3", ["float32", "float64", "int16", "complex64"]) + @pytest.mark.parametrize("dtype4", ["float32", "float64", "int16", "complex64"]) + @pytest.mark.parametrize("format1", ["dense", "csc", "csr"]) + @pytest.mark.parametrize("format2", ["dense", "csc", "csr"]) + def test_basic(self, dtype1, dtype2, dtype3, dtype4, format1, format2): + def mat(format, name, dtype): + if format == "dense": + return matrix(name, dtype=dtype) + else: + return pytensor.sparse.matrix(format, name, dtype=dtype) + + if format1 == "dense" and format2 == "dense": + pytest.skip("Skipping dense-dense case") + + x = mat(format1, "x", dtype1) + y = mat(format2, "y", dtype2) + a = scalar("a", dtype=dtype3) + z = pytensor.shared(np.asarray(self.z, dtype=dtype4).copy()) + + def f_b(z, a, x, y): + return z - a * (x * y) + + x_data = np.asarray(self.x, dtype=dtype1) + if format1 != "dense": + x_data = as_sparse_format(x_data, format1) + y_data = np.asarray(self.y, dtype=dtype2) + if format2 != "dense": + y_data = as_sparse_format(y_data, format2) + a_data = np.asarray(1.5, dtype=dtype3) + z_data = np.asarray(self.z, dtype=dtype4) + + f_b_out = f_b(z_data, a_data, x_data, y_data) + + # Can it work inplace? + inplace = dtype4 == pytensor.scalar.upcast(dtype1, dtype2, dtype3) + + # To make it easier to check the toposort + mode = pytensor.compile.mode.get_default_mode().excluding("fusion") + + if inplace: + updates = [(z, z - a * psm.dot(x, y))] + f_a = pytensor.function([a, x, y], [], updates=updates, mode=mode) + f_a(a_data, x_data, y_data) + f_a_out = z.get_value(borrow=True) + else: + f_a = pytensor.function([a, x, y], z - a * psm.dot(x, y), mode=mode) + # In DebugMode there is a strange difference with complex + # So we raise a little the threshold a little. + try: + orig_atol = pytensor.tensor.math.float64_atol + orig_rtol = pytensor.tensor.math.float64_rtol + pytensor.tensor.math.float64_atol = 1e-7 + pytensor.tensor.math.float64_rtol = 1e-6 + f_a_out = f_a(a_data, x_data, y_data) + finally: + pytensor.tensor.math.float64_atol = orig_atol + pytensor.tensor.math.float64_rtol = orig_rtol + + # As we do a dot product of 2 vector of 100 element, + # This mean we can have 2*100*eps abs error. + if f_a_out.dtype in ["float64", "complex128"]: + atol = 3e-8 + rtol = 1e-5 + else: + atol = None + rtol = None + utt.assert_allclose(f_a_out, f_b_out, rtol=rtol, atol=atol) + topo = f_a.maker.fgraph.toposort() + up = pytensor.scalar.upcast(dtype1, dtype2, dtype3, dtype4) + + fast_compile = pytensor.config.mode == "FAST_COMPILE" + + if not pytensor.config.blas__ldflags: + # Usmm should not be inserted, because it relies on BLAS + assert len(topo) == 4, topo + assert isinstance(topo[0].op, psm.Dot) + assert isinstance(topo[1].op, DimShuffle) + assert isinstance(topo[2].op, Elemwise) and isinstance( + topo[2].op.scalar_op, pytensor.scalar.Mul + ) + assert isinstance(topo[3].op, Elemwise) and isinstance( + topo[3].op.scalar_op, pytensor.scalar.Sub + ) + elif ( + y.type.dtype == up + and format1 == "csc" + and format2 == "dense" + and not fast_compile + and pytensor.config.cxx + and up in ("float32", "float64") + ): + # The op UsmmCscDense should be inserted + assert ( + sum( + isinstance(node.op, Elemwise) + and isinstance(node.op.scalar_op, pytensor.scalar.basic.Cast) + for node in topo + ) + == len(topo) - 5 + ) + new_topo = [ + node + for node in topo + if not ( + isinstance(node.op, Elemwise) + and isinstance(node.op.scalar_op, pytensor.scalar.basic.Cast) + ) + ] + topo = new_topo + assert len(topo) == 5, topo + + # Usmm is tested at the same time in debugmode + # Check if the optimization local_usmm and local_usmm_csx is + # applied + + def check_once(x): + assert sum(isinstance(n.op, x) for n in topo) == 1 + + check_once(CSMProperties) + check_once(DimShuffle) + check_once(Subtensor) + check_once(UsmmCscDense) + check_once(Elemwise) + if inplace: + assert topo[4].op.inplace + elif not fast_compile: + # The op Usmm should be inserted + assert len(topo) == 3, topo + assert isinstance(topo[0].op, DimShuffle) + assert topo[1].op == pytensor.tensor.neg + assert isinstance(topo[2].op, psm.Usmm) + + @pytest.mark.parametrize( + "params", + [ + ("float32", "float64", "int16", "complex64", "csc", "dense"), + ("float32", "float64", "int16", "complex64", "csr", "dense"), + ], + ) + def test_infer_shape(self, params): + def mat(format, name, dtype): + if format == "dense": + return matrix(name, dtype=dtype) + else: + return pytensor.sparse.matrix(format, name, dtype=dtype) + + dtype1, dtype2, dtype3, dtype4, format1, format2 = params + + if format1 == "dense" and format2 == "dense": + pytest.skip("Skipping dense-dense case; Usmm won't be used.") + + x = mat(format1, "x", dtype1) + y = mat(format2, "y", dtype2) + a = scalar("a", dtype=dtype3) + z = pytensor.shared(np.asarray(self.z, dtype=dtype4).copy()) + + def f_b(z, a, x, y): + return z - a * (x * y) + + x_data = np.asarray(self.x, dtype=dtype1) + if format1 != "dense": + x_data = as_sparse_format(x_data, format1) + y_data = np.asarray(self.y, dtype=dtype2) + if format2 != "dense": + y_data = as_sparse_format(y_data, format2) + a_data = np.asarray(1.5, dtype=dtype3) + z_data = np.asarray(self.z, dtype=dtype4) + + f_b_out = f_b(z_data, a_data, x_data, y_data) + + # Can it work inplace? + # inplace = dtype4 == pytensor.scalar.upcast(dtype1, dtype2, dtype3) + + # To make it easier to check the toposort + mode = pytensor.compile.mode.get_default_mode().excluding("fusion") + + # test infer_shape of Dot got applied + f_shape = pytensor.function([a, x, y], (z - a * psm.dot(x, y)).shape, mode=mode) + assert all(f_shape(a_data, x_data, y_data) == f_b_out.shape) + topo = f_shape.maker.fgraph.toposort() + assert not any(isinstance(node.op, Dot | Usmm | UsmmCscDense) for node in topo) + + +class TestTrueDot(utt.InferShapeTester): + def setup_method(self): + super().setup_method() + self.op = true_dot + self.op_class = TrueDot + + def test_op_ss(self): + for format in ["csc", "csr"]: + for dtype in all_dtypes: + variable, data = sparse_random_inputs( + format, shape=(10, 10), out_dtype=dtype, n=2, p=0.1 + ) + + f = pytensor.function(variable, self.op(*variable)) + + tested = f(*data) + + x, y = (m.toarray() for m in data) + expected = np.dot(x, y) + + assert tested.format == format + assert tested.dtype == expected.dtype + tested = tested.toarray() + utt.assert_allclose(tested, expected) + + def test_op_sd(self): + for format in ["csc", "csr"]: + for dtype in all_dtypes: + variable, data = sparse_random_inputs( + format, shape=(10, 10), out_dtype=dtype, n=2, p=0.1 + ) + variable[1] = TensorType(dtype=dtype, shape=(None, None))() + data[1] = data[1].toarray() + + f = pytensor.function(variable, self.op(*variable)) + + tested = f(*data) + expected = np.dot(data[0].toarray(), data[1]) + + assert tested.format == format + assert tested.dtype == expected.dtype + tested = tested.toarray() + utt.assert_allclose(tested, expected) + + def test_infer_shape(self): + for format in ["csc", "csr"]: + for dtype in all_dtypes: + (x,), (x_value,) = sparse_random_inputs( + format, shape=(9, 10), out_dtype=dtype, p=0.1 + ) + (y,), (y_value,) = sparse_random_inputs( + format, shape=(10, 24), out_dtype=dtype, p=0.1 + ) + variable = [x, y] + data = [x_value, y_value] + self._compile_and_check( + variable, [self.op(*variable)], data, self.op_class + ) + + def test_grad(self): + for format in ["csc", "csr"]: + for dtype in float_dtypes: + (_x,), (x_value,) = sparse_random_inputs( + format, shape=(9, 10), out_dtype=dtype, p=0.1 + ) + (_y,), (y_value,) = sparse_random_inputs( + format, shape=(10, 24), out_dtype=dtype, p=0.1 + ) + data = [x_value, y_value] + verify_grad_sparse(self.op, data, structured=False) + + +class TestSamplingDot(utt.InferShapeTester): + x = [matrix() for t in range(2)] + x.append(csr_matrix()) + a = [ + np.array( + np.random.default_rng().integers(1, 6, size=(4, 3)) - 1, + dtype=pytensor.config.floatX, + ), + np.array( + np.random.default_rng().integers(1, 6, size=(5, 3)) - 1, + dtype=pytensor.config.floatX, + ), + np.array( + np.random.default_rng().integers(1, 3, size=(4, 5)) - 1, + dtype=pytensor.config.floatX, + ), + ] + a[2] = scipy_sparse.csr_matrix(a[2]) + + def setup_method(self): + super().setup_method() + self.op_class = SamplingDot + + def test_op(self): + f = pytensor.function(self.x, sampling_dot(*self.x)) + tested = f(*self.a) + x, y, p = self.a + expected = p.multiply(np.dot(x, y.T)) + utt.assert_allclose(as_ndarray(expected), tested.toarray()) + assert tested.format == "csr" + assert tested.dtype == expected.dtype + + def test_negative_stride(self): + f = pytensor.function(self.x, sampling_dot(*self.x)) + a2 = [self.a[0][::-1, :], self.a[1][:, ::-1], self.a[2]] + tested = f(*a2) + x, y, p = a2 + expected = p.multiply(np.dot(x, y.T)) + utt.assert_allclose(as_ndarray(expected), tested.toarray()) + assert tested.format == "csr" + assert tested.dtype == expected.dtype + + def test_infer_shape(self): + self._compile_and_check( + self.x, + [sampling_dot(*self.x)], + self.a, + self.op_class, + excluding=["local_sampling_dot_csr"], + ) + + def test_grad(self): + def _helper(x, y): + return sampling_dot(x, y, self.a[2]) + + verify_grad_sparse(_helper, self.a[:2]) + + +class TestMulSV: + def test_mul_s_v_grad(self): + sp_types = {"csc": scipy_sparse.csc_matrix, "csr": scipy_sparse.csr_matrix} + for format in ("csr", "csc"): + for dtype in ("float32", "float64"): + spmat = sp_types[format](random_lil((4, 3), dtype, 3)) + mat = np.asarray(np.random.random(3), dtype=dtype) + verify_grad_sparse(mul_s_v, [spmat, mat], structured=True) + + def test_mul_s_v(self): + sp_types = {"csc": scipy_sparse.csc_matrix, "csr": scipy_sparse.csr_matrix} + for format in ("csr", "csc"): + for dtype in ("float32", "float64"): + x = SparseTensorType(format, dtype=dtype)() + y = vector(dtype=dtype) + f = pytensor.function([x, y], mul_s_v(x, y)) + + spmat = sp_types[format](random_lil((4, 3), dtype, 3)) + mat = np.asarray(np.random.random(3), dtype=dtype) + + out = f(spmat, mat) + utt.assert_allclose(spmat.toarray() * mat, out.toarray()) + + +class TestStructuredAddSV: + def test_structured_add_s_v_grad(self): + sp_types = {"csc": scipy_sparse.csc_matrix, "csr": scipy_sparse.csr_matrix} + for format in ("csr", "csc"): + for dtype in ("float32", "float64"): + spmat = sp_types[format](random_lil((4, 3), dtype, 3)) + mat = np.asarray(np.random.random(3), dtype=dtype) + verify_grad_sparse(structured_add_s_v, [spmat, mat], structured=True) + + def test_structured_add_s_v(self): + sp_types = {"csc": scipy_sparse.csc_matrix, "csr": scipy_sparse.csr_matrix} + for format in ("csr", "csc"): + for dtype in ("float32", "float64"): + x = SparseTensorType(format, dtype=dtype)() + y = vector(dtype=dtype) + f = pytensor.function([x, y], structured_add_s_v(x, y)) + + spmat = sp_types[format](random_lil((4, 3), dtype, 3)) + spones = spmat.copy() + spones.data = np.ones_like(spones.data) + mat = np.asarray(np.random.random(3), dtype=dtype) + + out = f(spmat, mat) + utt.assert_allclose( + as_ndarray(spones.multiply(spmat + mat)), out.toarray() + ) + + +class TestAddSSData(utt.InferShapeTester): + x = {} + a = {} + + def setup_method(self): + super().setup_method() + self.op_class = psm.AddSSData + + for format in sparse_formats: + variable = getattr(pytensor.sparse, format + "_matrix") + + a_val = np.array( + np.random.default_rng(utt.fetch_seed()).integers(1, 4, size=(3, 4)) - 1, + dtype=pytensor.config.floatX, + ) + constant = as_sparse_format(a_val, format) + + self.x[format] = [variable() for t in range(2)] + self.a[format] = [constant for t in range(2)] + + def test_op(self): + for format in sparse_formats: + f = pytensor.function(self.x[format], psm.add_s_s_data(*self.x[format])) + + tested = f(*self.a[format]) + expected = 2 * self.a[format][0] + + utt.assert_allclose(expected.toarray(), tested.toarray()) + assert tested.format == expected.format + assert tested.dtype == expected.dtype + + def test_infer_shape(self): + for format in sparse_formats: + self._compile_and_check( + self.x[format], + [psm.add_s_s_data(*self.x[format])], + self.a[format], + self.op_class, + ) + + def test_grad(self): + for format in sparse_formats: + verify_grad_sparse(self.op_class(), self.a[format], structured=True) + + +def test_useless_conj(): + x = SparseTensorType("csr", dtype="complex128")() + assert x.conj() is not x + x = SparseTensorType("csr", dtype="float64")() + assert x.conj() is x + + +class TestSpSum(utt.InferShapeTester): + possible_axis = [None, 0, 1] + + def setup_method(self): + super().setup_method() + self.op_class = psm.SpSum + self.op = psm.sp_sum + + @pytest.mark.parametrize("op_type", ["func", "method"]) + def test_op(self, op_type): + for format in sparse_formats: + for axis in self.possible_axis: + variable, data = sparse_random_inputs(format, shape=(10, 10)) + + if op_type == "func": + z = psm.sp_sum(variable[0], axis=axis) + if op_type == "method": + z = variable[0].sum(axis=axis) + + if axis is None: + assert z.type.broadcastable == () + else: + assert z.type.broadcastable == (False,) + + f = pytensor.function(variable, self.op(variable[0], axis=axis)) + tested = f(*data) + expected = data[0].todense().sum(axis).ravel() + utt.assert_allclose(expected, tested) + + def test_infer_shape(self): + for format in sparse_formats: + for axis in self.possible_axis: + variable, data = sparse_random_inputs(format, shape=(9, 10)) + self._compile_and_check( + variable, [self.op(variable[0], axis=axis)], data, self.op_class + ) + + def test_grad(self): + for format in sparse_formats: + for axis in self.possible_axis: + for struct in [True, False]: + _variable, data = sparse_random_inputs(format, shape=(9, 10)) + verify_grad_sparse( + self.op_class(axis=axis, sparse_grad=struct), + data, + structured=struct, + ) + + +def structure_function(f, index=0): + def structured_function(*args): + pattern = args[index] + evaluated = f(*args) + evaluated[pattern == 0] = 0 + return evaluated + + return structured_function + + +def elemwise_checker( + op, expected_f, gap=None, test_dtypes=None, grad_test=True, name=None, gap_grad=None +): + if test_dtypes is None: + test_dtypes = all_dtypes + + class TestElemwise: + def setup_method(self): + super().setup_method() + self.op = op + self.expected_f = expected_f + self.gap = gap + if gap_grad is not None: + self.gap_grad = gap_grad + else: + self.gap_grad = gap + assert eval(self.__class__.__name__) is self.__class__ + + def test_op(self): + for format in sparse_formats: + for dtype in test_dtypes: + if dtype == "int8" or dtype == "uint8": + continue + + if dtype.startswith("uint"): + if self.gap and len(self.gap) == 2 and self.gap[0] < 0: + if self.gap[1] >= 1: + self.gap = (0, self.gap[1]) + else: + raise TypeError( + "Gap not suitable for", dtype, self.__name__ + ) + + variable, data = sparse_random_inputs( + format, shape=(4, 7), out_dtype=dtype, gap=self.gap + ) + + f = pytensor.function(variable, self.op(*variable)) + + tested = f(*data) + data = [m.toarray() for m in data] + expected = self.expected_f(*data) + + assert tested.format == format + tested = tested.toarray() + + try: + utt.assert_allclose(expected, tested) + except AssertionError: + raise AssertionError(self.__name__) + + for dtype in ["int8", "uint8"]: + if dtype in test_dtypes: + if self.gap: + domain = self.gap + if dtype == "uint8": + if len(domain) == 2 and domain[0] < 0: + if domain[1] >= 1: + domain = (0, domain[1]) + else: + raise TypeError( + "Gap not suitable for", dtype, self.__name__ + ) + + else: + domain = (0, 5) + + variable, data = sparse_random_inputs( + format, shape=(4, 7), out_dtype=dtype, gap=domain + ) + + f = pytensor.function(variable, self.op(*variable)) + + old_value = ( + pt.math.float32_atol, + pt.math.float32_rtol, + pt.math.float64_atol, + pt.math.float64_rtol, + ) + pt.math.float32_atol = 1e-4 + pt.math.float32_rtol = 1e-3 + pt.math.float64_atol = 1e-3 + pt.math.float64_rtol = 1e-4 + try: + tested = f(*data) + finally: + ( + pt.math.float32_atol, + pt.math.float32_rtol, + pt.math.float64_atol, + pt.math.float64_rtol, + ) = old_value + + data = [m.toarray().astype("float32") for m in data] + expected = self.expected_f(*data) + + assert tested.format == format + tested = tested.toarray() + + try: + utt.assert_allclose(tested, expected, rtol=1e-2) + except AssertionError: + raise AssertionError(self.__name__) + + if grad_test: + + def test_grad(self): + for format in sparse_formats: + for dtype in float_dtypes: + _variable, data = sparse_random_inputs( + format, shape=(4, 7), out_dtype=dtype, gap=self.gap_grad + ) + + verify_grad_sparse(self.op, data, structured=True) + + if name is None: + name = op.__name__.capitalize() + "Tester" + TestElemwise.__name__ = name + if hasattr(TestElemwise, "__qualname__"): + TestElemwise.__qualname__ = name + assert "Roundhalftoeven" not in TestElemwise.__name__ + + return TestElemwise + + +StructuredSigmoidTester = elemwise_checker( + psm.structured_sigmoid, + structure_function(lambda x: 1.0 / (1.0 + np.exp(-x))), + test_dtypes=[ + m for m in all_dtypes if (m not in complex_dtypes and not m.startswith("uint")) + ], + gap=(-5, 5), + name="StructuredSigmoidTester", +) + +StructuredExpTester = elemwise_checker( + psm.structured_exp, structure_function(np.exp), name="StructuredExpTester" +) + +StructuredLogTester = elemwise_checker( + psm.structured_log, + structure_function(np.log), + gap=(0.5, 10), + name="StructuredLogTester", +) + +StructuredPowTester = elemwise_checker( + lambda x: psm.structured_pow(x, 2), + structure_function(lambda x: np.power(x, 2)), + name="StructuredPowTester", +) + +StructuredMinimumTester = elemwise_checker( + lambda x: psm.structured_minimum(x, 2), + structure_function(lambda x: np.minimum(x, 2)), + name="StructuredMinimumTester", +) + +StructuredMaximumTester = elemwise_checker( + lambda x: psm.structured_maximum(x, 2), + structure_function(lambda x: np.maximum(x, 2)), + name="StructuredMaximumTester", +) + +StructuredAddTester = elemwise_checker( + lambda x: psm.structured_add(x, 2), + structure_function(lambda x: np.add(x, 2)), + name="StructuredAddTester", +) + +SinTester = elemwise_checker(psm.structured_sin, np.sin) + +TanTester = elemwise_checker(psm.structured_tan, np.tan, gap=(-1, 1)) + +ArcsinTester = elemwise_checker( + psm.structured_arcsin, np.arcsin, gap=(-1, 1), gap_grad=(-0.99, 0.99) +) + +ArctanTester = elemwise_checker(psm.structured_arctan, np.arctan) + +SinhTester = elemwise_checker(psm.structured_sinh, np.sinh) + +ArcsinhTester = elemwise_checker(psm.structured_arcsinh, np.arcsinh, gap=(-1, 1)) + +TanhTester = elemwise_checker(psm.structured_tanh, np.tanh, gap=(-1, 1)) + +ArctanhTester = elemwise_checker( + psm.structured_arctanh, np.arctanh, gap=(-0.9, 1), gap_grad=(-0.9, 0.95) +) + +RintTester = elemwise_checker( + psm.structured_rint, np.rint, grad_test=False, test_dtypes=float_dtypes +) + +SgnTester = elemwise_checker( + psm.structured_sign, + np.sign, + grad_test=False, + test_dtypes=[ + m for m in all_dtypes if (m not in complex_dtypes and not m.startswith("uint")) + ], +) + +CeilTester = elemwise_checker( + psm.structured_ceil, + np.ceil, + grad_test=False, + test_dtypes=[m for m in all_dtypes if m not in complex_dtypes], +) + +FloorTester = elemwise_checker( + psm.structured_floor, + np.floor, + grad_test=False, + test_dtypes=[m for m in all_dtypes if m not in complex_dtypes], +) + +Log1pTester = elemwise_checker(psm.structured_log1p, np.log1p, gap=(0.5, 10)) + +Expm1Tester = elemwise_checker(psm.structured_expm1, np.expm1) + +Deg2radTester = elemwise_checker( + psm.structured_deg2rad, + np.deg2rad, + test_dtypes=[m for m in all_dtypes if m not in complex_dtypes], +) + +Rad2degTester = elemwise_checker( + psm.structured_rad2deg, + np.rad2deg, + test_dtypes=[m for m in all_dtypes if m not in complex_dtypes], +) + + +TruncTester = elemwise_checker( + psm.structured_trunc, + np.trunc, + test_dtypes=[m for m in all_dtypes if m not in complex_dtypes], + grad_test=False, +) + + +SqrTester = elemwise_checker(psm.structured_sqr, lambda x: x * x) + +SqrtTester = elemwise_checker(psm.structured_sqrt, np.sqrt, gap=(0, 10)) + +ConjTester = elemwise_checker(psm.structured_conjugate, np.conj, grad_test=False) diff --git a/tests/sparse/test_rewriting.py b/tests/sparse/test_rewriting.py index 280d9dbf70..759ee17f34 100644 --- a/tests/sparse/test_rewriting.py +++ b/tests/sparse/test_rewriting.py @@ -3,6 +3,7 @@ import scipy as sp import pytensor +import pytensor.sparse.math as smath from pytensor import sparse from pytensor.compile.mode import Mode, get_default_mode from pytensor.configdefaults import config @@ -74,10 +75,10 @@ def test_local_mul_s_d(): for sp_format in sparse.sparse_formats: inputs = [getattr(pytensor.sparse, sp_format + "_matrix")(), matrix()] - f = pytensor.function(inputs, sparse.mul_s_d(*inputs), mode=mode) + f = pytensor.function(inputs, smath.mul_s_d(*inputs), mode=mode) assert not any( - isinstance(node.op, sparse.MulSD) for node in f.maker.fgraph.toposort() + isinstance(node.op, smath.MulSD) for node in f.maker.fgraph.toposort() ) @@ -91,10 +92,10 @@ def test_local_mul_s_v(): for sp_format in ["csr"]: # Not implemented for other format inputs = [getattr(pytensor.sparse, sp_format + "_matrix")(), vector()] - f = pytensor.function(inputs, sparse.mul_s_v(*inputs), mode=mode) + f = pytensor.function(inputs, smath.mul_s_v(*inputs), mode=mode) assert not any( - isinstance(node.op, sparse.MulSV) for node in f.maker.fgraph.toposort() + isinstance(node.op, smath.MulSV) for node in f.maker.fgraph.toposort() ) @@ -108,10 +109,10 @@ def test_local_structured_add_s_v(): for sp_format in ["csr"]: # Not implemented for other format inputs = [getattr(pytensor.sparse, sp_format + "_matrix")(), vector()] - f = pytensor.function(inputs, sparse.structured_add_s_v(*inputs), mode=mode) + f = pytensor.function(inputs, smath.structured_add_s_v(*inputs), mode=mode) assert not any( - isinstance(node.op, sparse.StructuredAddSV) + isinstance(node.op, smath.StructuredAddSV) for node in f.maker.fgraph.toposort() ) @@ -130,11 +131,11 @@ def test_local_sampling_dot_csr(): getattr(pytensor.sparse, sp_format + "_matrix")(), ] - f = pytensor.function(inputs, sparse.sampling_dot(*inputs), mode=mode) + f = pytensor.function(inputs, smath.sampling_dot(*inputs), mode=mode) if pytensor.config.blas__ldflags: assert not any( - isinstance(node.op, sparse.SamplingDot) + isinstance(node.op, smath.SamplingDot) for node in f.maker.fgraph.toposort() ) else: diff --git a/tests/sparse/test_var.py b/tests/sparse/test_var.py deleted file mode 100644 index 4c897184ee..0000000000 --- a/tests/sparse/test_var.py +++ /dev/null @@ -1,238 +0,0 @@ -from contextlib import ExitStack - -import numpy as np -import pytest -from scipy.sparse.csr import csr_matrix - -import pytensor -import pytensor.sparse as sparse -import pytensor.tensor as pt -from pytensor.sparse.type import SparseTensorType -from pytensor.tensor.type import DenseTensorType - - -class TestSparseVariable: - @pytest.mark.parametrize( - "method, exp_type, cm, x", - [ - ("__abs__", DenseTensorType, None, None), - ("__neg__", SparseTensorType, ExitStack(), None), - ("__ceil__", DenseTensorType, None, None), - ("__floor__", DenseTensorType, None, None), - ("__trunc__", DenseTensorType, None, None), - ("transpose", DenseTensorType, None, None), - ("any", DenseTensorType, None, None), - ("all", DenseTensorType, None, None), - ("flatten", DenseTensorType, None, None), - ("ravel", DenseTensorType, None, None), - ("arccos", DenseTensorType, None, None), - ("arcsin", DenseTensorType, None, None), - ("arctan", DenseTensorType, None, None), - ("arccosh", DenseTensorType, None, None), - ("arcsinh", DenseTensorType, None, None), - ("arctanh", DenseTensorType, None, None), - ("ceil", DenseTensorType, None, None), - ("cos", DenseTensorType, None, None), - ("cosh", DenseTensorType, None, None), - ("deg2rad", DenseTensorType, None, None), - ("exp", DenseTensorType, None, None), - ("exp2", DenseTensorType, None, None), - ("expm1", DenseTensorType, None, None), - ("floor", DenseTensorType, None, None), - ("log", DenseTensorType, None, None), - ("log10", DenseTensorType, None, None), - ("log1p", DenseTensorType, None, None), - ("log2", DenseTensorType, None, None), - ("rad2deg", DenseTensorType, None, None), - ("sin", DenseTensorType, None, None), - ("sinh", DenseTensorType, None, None), - ("sqrt", DenseTensorType, None, None), - ("tan", DenseTensorType, None, None), - ("tanh", DenseTensorType, None, None), - ("copy", DenseTensorType, None, None), - ("sum", DenseTensorType, ExitStack(), None), - ("prod", DenseTensorType, None, None), - ("mean", DenseTensorType, None, None), - ("var", DenseTensorType, None, None), - ("std", DenseTensorType, None, None), - ("min", DenseTensorType, None, None), - ("max", DenseTensorType, None, None), - ("argmin", DenseTensorType, None, None), - ("argmax", DenseTensorType, None, None), - ("nonzero", DenseTensorType, ExitStack(), None), - ("nonzero_values", DenseTensorType, None, None), - ("argsort", DenseTensorType, ExitStack(), None), - ("conj", SparseTensorType, ExitStack(), pt.cmatrix("x")), - ("round", DenseTensorType, None, None), - ("trace", DenseTensorType, None, None), - ("zeros_like", SparseTensorType, ExitStack(), None), - ("ones_like", DenseTensorType, ExitStack(), None), - ("cumsum", DenseTensorType, None, None), - ("cumprod", DenseTensorType, None, None), - ("ptp", DenseTensorType, None, None), - ("squeeze", DenseTensorType, None, None), - ("diagonal", DenseTensorType, None, None), - ], - ) - def test_unary(self, method, exp_type, cm, x): - if x is None: - x = pt.dmatrix("x") - - x = sparse.csr_from_dense(x) - - method_to_call = getattr(x, method) - - if cm is None: - cm = pytest.warns(UserWarning, match=".*converted to dense.*") - - if exp_type == SparseTensorType: - exp_res_type = csr_matrix - else: - exp_res_type = np.ndarray - - with cm: - z = method_to_call() - - if not isinstance(z, tuple): - z_outs = (z,) - else: - z_outs = z - - assert all(isinstance(out.type, exp_type) for out in z_outs) - - f = pytensor.function( - [x], z, on_unused_input="ignore", allow_input_downcast=True - ) - - res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) - - if not isinstance(res, list): - res_outs = [res] - else: - res_outs = res - - assert all(isinstance(out, exp_res_type) for out in res_outs) - - @pytest.mark.parametrize( - "method, exp_type", - [ - ("__lt__", SparseTensorType), - ("__le__", SparseTensorType), - ("__gt__", SparseTensorType), - ("__ge__", SparseTensorType), - ("__and__", DenseTensorType), - ("__or__", DenseTensorType), - ("__xor__", DenseTensorType), - ("__add__", SparseTensorType), - ("__sub__", SparseTensorType), - ("__mul__", SparseTensorType), - ("__pow__", DenseTensorType), - ("__mod__", DenseTensorType), - ("__divmod__", DenseTensorType), - ("__truediv__", DenseTensorType), - ("__floordiv__", DenseTensorType), - ], - ) - def test_binary(self, method, exp_type): - x = pt.lmatrix("x") - y = pt.lmatrix("y") - x = sparse.csr_from_dense(x) - y = sparse.csr_from_dense(y) - - method_to_call = getattr(x, method) - - if exp_type == SparseTensorType: - exp_res_type = csr_matrix - cm = ExitStack() - else: - exp_res_type = np.ndarray - cm = pytest.warns(UserWarning, match=".*converted to dense.*") - - with cm: - z = method_to_call(y) - - if not isinstance(z, tuple): - z_outs = (z,) - else: - z_outs = z - - assert all(isinstance(out.type, exp_type) for out in z_outs) - - f = pytensor.function([x, y], z) - res = f( - [[1, 0, 2], [-1, 0, 0]], - [[1, 1, 2], [1, 4, 1]], - ) - - if not isinstance(res, list): - res_outs = [res] - else: - res_outs = res - - assert all(isinstance(out, exp_res_type) for out in res_outs) - - def test_reshape(self): - x = pt.dmatrix("x") - x = sparse.csr_from_dense(x) - - with pytest.warns(UserWarning, match=".*converted to dense.*"): - z = x.reshape((3, 2)) - - assert isinstance(z.type, DenseTensorType) - - f = pytensor.function([x], z) - exp_res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) - assert isinstance(exp_res, np.ndarray) - - def test_dimshuffle(self): - x = pt.dmatrix("x") - x = sparse.csr_from_dense(x) - - with pytest.warns(UserWarning, match=".*converted to dense.*"): - z = x.dimshuffle((1, 0)) - - assert isinstance(z.type, DenseTensorType) - - f = pytensor.function([x], z) - exp_res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) - assert isinstance(exp_res, np.ndarray) - - def test_getitem(self): - x = pt.dmatrix("x") - x = sparse.csr_from_dense(x) - - z = x[:, :2] - assert isinstance(z.type, SparseTensorType) - - f = pytensor.function([x], z) - exp_res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) - assert isinstance(exp_res, csr_matrix) - - def test_dot(self): - x = pt.lmatrix("x") - y = pt.lmatrix("y") - x = sparse.csr_from_dense(x) - y = sparse.csr_from_dense(y) - - z = x.__dot__(y) - assert isinstance(z.type, SparseTensorType) - - f = pytensor.function([x, y], z) - exp_res = f( - [[1, 0, 2], [-1, 0, 0]], - [[-1], [2], [1]], - ) - assert isinstance(exp_res, csr_matrix) - - def test_repeat(self): - x = pt.dmatrix("x") - x = sparse.csr_from_dense(x) - - with pytest.warns(UserWarning, match=".*converted to dense.*"): - z = x.repeat(2, axis=1) - - assert isinstance(z.type, DenseTensorType) - - f = pytensor.function([x], z) - exp_res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) - assert isinstance(exp_res, np.ndarray) diff --git a/tests/sparse/test_variable.py b/tests/sparse/test_variable.py new file mode 100644 index 0000000000..36c46160c9 --- /dev/null +++ b/tests/sparse/test_variable.py @@ -0,0 +1,235 @@ +from contextlib import ExitStack + +import numpy as np +import pytest +from scipy.sparse.csr import csr_matrix + +import pytensor +import pytensor.sparse as sparse +import pytensor.tensor as pt +from pytensor.sparse.type import SparseTensorType +from pytensor.tensor.type import DenseTensorType + + +class TestSparseVariable: + @pytest.mark.parametrize( + "method", + [ + "__abs__", + "__neg__", + "__ceil__", + "__floor__", + "__trunc__", + "any", + "all", + "flatten", + "ravel", + "arccos", + "arcsin", + "arctan", + "arccosh", + "arcsinh", + "arctanh", + "ceil", + "cos", + "cosh", + "deg2rad", + "exp", + "exp2", + "expm1", + "floor", + "log", + "log10", + "log1p", + "log2", + "rad2deg", + "sin", + "sinh", + "sqrt", + "tan", + "tanh", + "copy", + "sum", + "prod", + "mean", + "var", + "std", + "min", + "max", + "argmin", + "argmax", + "nonzero", + "nonzero_values", + "argsort", + "conj", + "round", + "trace", + "zeros_like", + "ones_like", + "cumsum", + "cumprod", + "ptp", + "squeeze", + "diagonal", + ], + ) + def test_unary(self, method): + x = pt.dmatrix("x") if method != "conj" else pt.cmatrix("x") + x = sparse.csr_from_dense(x) + + method_to_call = getattr(x, method) + z = method_to_call() + + if not isinstance(z, tuple): + z_outs = (z,) + else: + z_outs = z + + f = pytensor.function( + [x], z, on_unused_input="ignore", allow_input_downcast=True + ) + + res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) + + if not isinstance(res, list): + res_outs = [res] + else: + res_outs = res + + # TODO: Make a separate test for methods that always reduce to dense (only sum for now) + if getattr(method_to_call, "_is_dense_override", False) or method == "sum": + assert all(isinstance(out.type, DenseTensorType) for out in z_outs) + assert all(isinstance(out, np.ndarray) for out in res_outs) + else: + assert all(isinstance(out.type, SparseTensorType) for out in z_outs) + assert all(isinstance(out, csr_matrix) for out in res_outs) + + @pytest.mark.parametrize( + "method", + [ + "__lt__", + "__le__", + "__gt__", + "__ge__", + "__and__", + "__or__", + "__xor__", + "__add__", + "__sub__", + "__mul__", + "__pow__", + "__mod__", + "__divmod__", + "__truediv__", + "__floordiv__", + ], + ) + def test_binary(self, method): + x = pt.lmatrix("x") + y = pt.lmatrix("y") + x = sparse.csr_from_dense(x) + y = sparse.csr_from_dense(y) + + method_to_call = getattr(x, method) + + exp_type = ( + DenseTensorType + if getattr(method_to_call, "_is_dense_override", False) + else SparseTensorType + ) + + if exp_type == SparseTensorType: + exp_res_type = csr_matrix + cm = ExitStack() + else: + exp_res_type = np.ndarray + cm = pytest.warns(UserWarning, match=".*converted to dense.*") + + with cm: + z = method_to_call(y) + + if not isinstance(z, tuple): + z_outs = (z,) + else: + z_outs = z + + assert all(isinstance(out.type, exp_type) for out in z_outs) + + f = pytensor.function([x, y], z) + res = f( + [[1, 0, 2], [-1, 0, 0]], + [[1, 1, 2], [1, 4, 1]], + ) + + if not isinstance(res, list): + res_outs = [res] + else: + res_outs = res + + assert all(isinstance(out, exp_res_type) for out in res_outs) + + def test_reshape(self): + x = pt.dmatrix("x") + x = sparse.csr_from_dense(x) + + with pytest.warns(UserWarning, match=".*converted to dense.*"): + z = x.reshape((3, 2)) + + assert isinstance(z.type, DenseTensorType) + + f = pytensor.function([x], z) + exp_res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) + assert isinstance(exp_res, np.ndarray) + + def test_dimshuffle(self): + x = pt.dmatrix("x") + x = sparse.csr_from_dense(x) + + with pytest.warns(UserWarning, match=".*converted to dense.*"): + z = x.dimshuffle((1, 0)) + + assert isinstance(z.type, DenseTensorType) + + f = pytensor.function([x], z) + exp_res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) + assert isinstance(exp_res, np.ndarray) + + def test_getitem(self): + x = pt.dmatrix("x") + x = sparse.csr_from_dense(x) + + z = x[:, :2] + assert isinstance(z.type, SparseTensorType) + + f = pytensor.function([x], z) + exp_res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) + assert isinstance(exp_res, csr_matrix) + + def test_dot(self): + x = pt.lmatrix("x") + y = pt.lmatrix("y") + x = sparse.csr_from_dense(x) + y = sparse.csr_from_dense(y) + + z = x.__dot__(y) + assert isinstance(z.type, SparseTensorType) + + f = pytensor.function([x, y], z) + exp_res = f( + [[1, 0, 2], [-1, 0, 0]], + [[-1], [2], [1]], + ) + assert isinstance(exp_res, csr_matrix) + + def test_repeat(self): + x = pt.dmatrix("x") + x = sparse.csr_from_dense(x) + + with pytest.warns(UserWarning, match=".*converted to dense.*"): + z = x.repeat(2, axis=1) + + assert isinstance(z.type, DenseTensorType) + + f = pytensor.function([x], z) + exp_res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) + assert isinstance(exp_res, np.ndarray) From 7d9687d1dcab50d4e0a7a30165db8c6dfd5bceae Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sat, 18 Oct 2025 13:15:54 -0500 Subject: [PATCH 2/3] Appease mypy --- pytensor/sparse/sharedvar.py | 2 +- pytensor/sparse/variable.py | 4 ++-- pytensor/tensor/slinalg.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pytensor/sparse/sharedvar.py b/pytensor/sparse/sharedvar.py index ce9f9ca6d0..12385b6f07 100644 --- a/pytensor/sparse/sharedvar.py +++ b/pytensor/sparse/sharedvar.py @@ -7,7 +7,7 @@ from pytensor.tensor.sharedvar import TensorSharedVariable -class SparseTensorSharedVariable(TensorSharedVariable, SparseVariable): +class SparseTensorSharedVariable(TensorSharedVariable, SparseVariable): # type: ignore[misc] @property def format(self): return self.type.format diff --git a/pytensor/sparse/variable.py b/pytensor/sparse/variable.py index fc462b278b..5387442282 100644 --- a/pytensor/sparse/variable.py +++ b/pytensor/sparse/variable.py @@ -421,7 +421,7 @@ def dimshuffle(self, *pattern): raise NotImplementedError -class SparseVariable(_sparse_py_operators, TensorVariable): +class SparseVariable(_sparse_py_operators, TensorVariable): # type: ignore[misc] format = property(lambda self: self.type.format) def __str__(self): @@ -454,7 +454,7 @@ def pytensor_hash(self): return hash_from_sparse(d) -class SparseConstant(SparseVariable, TensorConstant): +class SparseConstant(SparseVariable, TensorConstant): # type: ignore[misc] format = property(lambda self: self.type.format) def signature(self): diff --git a/pytensor/tensor/slinalg.py b/pytensor/tensor/slinalg.py index cf1358813e..2616ccb794 100644 --- a/pytensor/tensor/slinalg.py +++ b/pytensor/tensor/slinalg.py @@ -1654,7 +1654,7 @@ def _largest_common_dtype(tensors: Sequence[TensorVariable]) -> np.dtype: class BaseBlockDiagonal(Op): - __props__ = ("n_inputs",) + __props__: tuple[str, ...] = ("n_inputs",) def __init__(self, n_inputs): input_sig = ",".join(f"(m{i},n{i})" for i in range(n_inputs)) From c904d24ca408c39388fffbfb7603493f344057d0 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sat, 18 Oct 2025 13:21:41 -0500 Subject: [PATCH 3/3] Remove `/pytensor/sparse/basic.py` from `mypy-failing.txt` --- scripts/mypy-failing.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/scripts/mypy-failing.txt b/scripts/mypy-failing.txt index 5b39180db9..11da9ed34c 100644 --- a/scripts/mypy-failing.txt +++ b/scripts/mypy-failing.txt @@ -9,7 +9,6 @@ pytensor/link/numba/dispatch/elemwise.py pytensor/link/numba/dispatch/scan.py pytensor/printing.py pytensor/raise_op.py -pytensor/sparse/basic.py pytensor/tensor/basic.py pytensor/tensor/blas_c.py pytensor/tensor/blas_headers.py