Skip to content

P
pytensor
Unverified 提交 f2c4977b authored 作者: Jesse Grabowski's avatar Jesse Grabowski 提交者： GitHub
浏览文件
操作

Reorganize the `sparse` module (#1674)

上级 c77d1eff
隐藏空白字符变更
内嵌 并排
正在显示 包含 4526 行增加 和 4624 行删除
pytensor/link/jax/dispatch/sparse.py
浏览文件 @ f2c4977b
......@@ -3,7 +3,7 @@ from scipy.sparse import spmatrix
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
......
pytensor/sparse/__init__.py
浏览文件 @ f2c4977b
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
......
pytensor/sparse/basic.py
浏览文件 @ f2c4977b
......@@ -8,7 +8,6 @@ TODO: Automatic methods for determining best sparse format?
"""
from typing import Literal
from warnings import warn
import numpy as np
......@@ -19,55 +18,18 @@ import pytensor
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 @@ as_sparse = as_sparse_variable
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 @@ class ColScaleCSC(Op):
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 @@ class RowScaleCSC(Op):
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 @@ class ConstructSparseFromList(Op):
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)
pytensor/sparse/linalg.py 0 → 100644
浏览文件 @ f2c4977b
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)
pytensor/sparse/math.py 0 → 100644
浏览文件 @ f2c4977b
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.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, y):
"""
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, y):
"""
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()
pytensor/sparse/rewriting.py
浏览文件 @ f2c4977b
......@@ -3,6 +3,8 @@ import scipy
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 @@ from pytensor.graph.rewriting.basic import (
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_sparse_variable = sparse._is_sparse_variable
_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 @@ class AddSD_ccode(_NoPythonCOp):
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 @@ class AddSD_ccode(_NoPythonCOp):
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 @@ sd_csr = StructuredDotCSR()
# 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 @@ local_usmm = PatternNodeRewriter(
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 @@ register_specialize(local_usmm_csc_dense_inplace, "cxx_only", "inplace")
# 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 @@ class CSMGradC(_NoPythonCOp):
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 @@ mul_s_d_csr = MulSDCSR()
# 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 @@ mul_s_v_csr = MulSVCSR()
# 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 @@ structured_add_s_v_csr = StructuredAddSVCSR()
# 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 @@ sampling_dot_csr = SamplingDotCSR()
# 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)
......
pytensor/sparse/sharedvar.py
浏览文件 @ f2c4977b
......@@ -3,11 +3,11 @@ import copy
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
class SparseTensorSharedVariable(TensorSharedVariable, SparseVariable):
class SparseTensorSharedVariable(TensorSharedVariable, SparseVariable): # type: ignore[misc]
@property
def format(self):
return self.type.format
......
pytensor/sparse/variable.py 0 → 100644
浏览文件 @ f2c4977b
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_conjugate,
structured_dot,
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)
@override_dense
def __abs__(self):
raise NotImplementedError
@override_dense
def __ceil__(self):
raise NotImplementedError
@override_dense
def __floor__(self):
raise NotImplementedError
@override_dense
def __trunc__(self):
raise NotImplementedError
@override_dense
def transpose(self):
raise NotImplementedError
@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
@override_dense
def arccos(self):
raise NotImplementedError
@override_dense
def arcsin(self):
raise NotImplementedError
@override_dense
def arctan(self):
raise NotImplementedError
@override_dense
def arccosh(self):
raise NotImplementedError
@override_dense
def arcsinh(self):
raise NotImplementedError
@override_dense
def arctanh(self):
raise NotImplementedError
@override_dense
def ceil(self):
raise NotImplementedError
@override_dense
def cos(self):
raise NotImplementedError
@override_dense
def cosh(self):
raise NotImplementedError
@override_dense
def deg2rad(self):
raise NotImplementedError
@override_dense
def exp(self):
raise NotImplementedError
@override_dense
def exp2(self):
raise NotImplementedError
@override_dense
def expm1(self):
raise NotImplementedError
@override_dense
def floor(self):
raise NotImplementedError
@override_dense
def log(self):
raise NotImplementedError
@override_dense
def log10(self):
raise NotImplementedError
@override_dense
def log1p(self):
raise NotImplementedError
@override_dense
def log2(self):
raise NotImplementedError
@override_dense
def rad2deg(self):
raise NotImplementedError
@override_dense
def sin(self):
raise NotImplementedError
@override_dense
def sinh(self):
raise NotImplementedError
@override_dense
def sqrt(self):
raise NotImplementedError
@override_dense
def tan(self):
raise NotImplementedError
@override_dense
def tanh(self):
raise NotImplementedError
@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): # type: ignore[misc]
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): # type: ignore[misc]
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
pytensor/tensor/slinalg.py
浏览文件 @ f2c4977b
......@@ -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))
......
scripts/mypy-failing.txt
浏览文件 @ f2c4977b
......@@ -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
......
tests/sparse/test_basic.py
浏览文件 @ f2c4977b
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 @@ from pytensor.sparse import (
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 sparse_random_inputs(
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 sparse_random_inputs(
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 sparse_random_inputs(
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 @@ class TestVerifyGradSparse:
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 @@ class TestTranspose:
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 @@ class TestSparseInferShape(utt.InferShapeTester):
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 @@ class TestSparseInferShape(utt.InferShapeTester):
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 @@ class TestSparseInferShape(utt.InferShapeTester):
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 @@ class TestSparseInferShape(utt.InferShapeTester):
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 @@ class TestSparseInferShape(utt.InferShapeTester):
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 @@ class TestConstructSparseFromList:
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 @@ class TestConversion:
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 @@ class TestConversion:
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 @@ class TestConversion:
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 @@ class TestCsmProperties:
)
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 @@ class TestCsmProperties:
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 @@ class TestCsm:
)
@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 @@ class TestCsm:
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 @@ class TestCsm:
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 @@ class TestCsm:
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 @@ class TestRowScaleCSC(utt.InferShapeTester):
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 @@ class TestRemove0(utt.InferShapeTester):
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 @@ class TestRemove0(utt.InferShapeTester):
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 @@ class TestGetItem:
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 @@ class TestGetItem:
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 @@ class TestGetItem:
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 @@ class TestGetItem:
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 @@ class TestGetItem:
# 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 @@ class TestGetItem:
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 @@ class TestGetItem:
# ----------------------------------------------------------
# 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 @@ class TestGetItem:
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 _hv_switch(op, expected_function):
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 test_hstack_vstack():
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())
tests/sparse/test_linalg.py 0 → 100644
浏览文件 @ f2c4977b
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())
tests/sparse/test_math.py 0 → 100644
浏览文件 @ f2c4977b
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)
tests/sparse/test_rewriting.py
浏览文件 @ f2c4977b
......@@ -3,6 +3,7 @@ import pytest
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:
......
tests/sparse/test_var.py → tests/sparse/test_variable.py
浏览文件 @ f2c4977b
......@@ -13,93 +13,78 @@ from pytensor.tensor.type import DenseTensorType
class TestSparseVariable:
@pytest.mark.parametrize(
"method, exp_type, cm, x",
"method",
[
("__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),
"__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, exp_type, cm, x):
if x is None:
x = pt.dmatrix("x")
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)
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()
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
)
......@@ -111,29 +96,35 @@ class TestSparseVariable:
else:
res_outs = res
assert all(isinstance(out, exp_res_type) for out in res_outs)
# 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, exp_type",
"method",
[
("__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),
"__lt__",
"__le__",
"__gt__",
"__ge__",
"__and__",
"__or__",
"__xor__",
"__add__",
"__sub__",
"__mul__",
"__pow__",
"__mod__",
"__divmod__",
"__truediv__",
"__floordiv__",
],
)
def test_binary(self, method, exp_type):
def test_binary(self, method):
x = pt.lmatrix("x")
y = pt.lmatrix("y")
x = sparse.csr_from_dense(x)
......@@ -141,6 +132,12 @@ class TestSparseVariable:
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()
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论