From 0d3e0cba67a6cc815e528c9260de35ebf76f79c5 Mon Sep 17 00:00:00 2001 From: "Brandon T. Willard" Date: Sun, 14 Feb 2021 00:44:48 -0600 Subject: [PATCH 1/3] Make sparse tensor types extend the existing tensor types --- aesara/sparse/basic.py | 26 +++++++---- aesara/sparse/opt.py | 16 ++++++- aesara/sparse/type.py | 89 ++++++++++++++++++++++---------------- aesara/tensor/basic.py | 9 +++- aesara/tensor/basic_opt.py | 10 ++++- aesara/tensor/blas.py | 40 +++++++++++++++-- aesara/tensor/math.py | 6 +++ aesara/tensor/shape.py | 2 +- aesara/tensor/type.py | 22 ++++++++-- aesara/tensor/var.py | 31 +++++++++++++ tests/sparse/test_basic.py | 66 ++++++++++++++++++++++++---- tests/sparse/test_type.py | 22 ++++++++++ tests/tensor/test_var.py | 19 +++++++- tests/test_ifelse.py | 8 ++-- 14 files changed, 294 insertions(+), 72 deletions(-) diff --git a/aesara/sparse/basic.py b/aesara/sparse/basic.py index a89deed0d2..f0df5c1cb0 100644 --- a/aesara/sparse/basic.py +++ b/aesara/sparse/basic.py @@ -49,6 +49,7 @@ from aesara.tensor.type import continuous_dtypes as tensor_continuous_dtypes from aesara.tensor.type import discrete_dtypes as tensor_discrete_dtypes from aesara.tensor.type import iscalar, ivector, scalar, tensor, vector +from aesara.tensor.var import TensorConstant, TensorVariable, _tensor_py_operators sparse_formats = ["csc", "csr"] @@ -126,8 +127,7 @@ def _is_dense(x): return isinstance(x, np.ndarray) -# Wrapper type -def as_sparse_variable(x, name=None): +def as_sparse_variable(x, name=None, ndim=None, **kwargs): """ Wrapper around SparseVariable constructor to construct a Variable with a sparse matrix with the same dtype and @@ -250,7 +250,7 @@ def sp_zeros_like(x): ) -class _sparse_py_operators: +class _sparse_py_operators(_tensor_py_operators): T = property( lambda self: transpose(self), doc="Return aliased transpose of self (read-only)" ) @@ -361,8 +361,7 @@ def __getitem__(self, args): return ret -class SparseVariable(_sparse_py_operators, Variable): - dtype = property(lambda self: self.type.dtype) +class SparseVariable(_sparse_py_operators, TensorVariable): format = property(lambda self: self.type.format) def __str__(self): @@ -395,8 +394,7 @@ def aesara_hash(self): return hash_from_sparse(d) -class SparseConstant(Constant, _sparse_py_operators): - dtype = property(lambda self: self.type.dtype) +class SparseConstant(TensorConstant, _sparse_py_operators): format = property(lambda self: self.type.format) def signature(self): @@ -448,7 +446,7 @@ def bsr_matrix(name=None, dtype=None): csr_fmatrix = SparseType(format="csr", dtype="float32") bsr_fmatrix = SparseType(format="bsr", dtype="float32") -all_dtypes = SparseType.dtype_set +all_dtypes = list(SparseType.dtype_specs_map.keys()) complex_dtypes = [t for t in all_dtypes if t[:7] == "complex"] float_dtypes = [t for t in all_dtypes if t[:5] == "float"] int_dtypes = [t for t in all_dtypes if t[:3] == "int"] @@ -926,6 +924,12 @@ def __init__(self, structured=True): def __str__(self): return f"{self.__class__.__name__}{{structured_grad={self.sparse_grad}}}" + def __call__(self, x): + if not isinstance(x.type, SparseType): + return x + + return super().__call__(x) + def make_node(self, x): x = as_sparse_variable(x) return Apply( @@ -1003,6 +1007,12 @@ def __init__(self, format): def __str__(self): return f"{self.__class__.__name__}{{{self.format}}}" + def __call__(self, x): + if isinstance(x.type, SparseType): + return x + + return super().__call__(x) + def make_node(self, x): x = at.as_tensor_variable(x) if x.ndim > 2: diff --git a/aesara/sparse/opt.py b/aesara/sparse/opt.py index 5d0134427e..b646f43248 100644 --- a/aesara/sparse/opt.py +++ b/aesara/sparse/opt.py @@ -23,6 +23,7 @@ from aesara.tensor.basic import as_tensor_variable, cast, patternbroadcast from aesara.tensor.basic_opt import register_canonicalize, register_specialize from aesara.tensor.math import mul, neg, sub +from aesara.tensor.shape import shape, specify_shape from aesara.tensor.type import TensorType, tensor @@ -2070,8 +2071,19 @@ def local_sampling_dot_csr(fgraph, node): z_data, z_ind, z_ptr = sampling_dot_csr( x, y, p_data, p_ind, p_ptr, p_shape[1] ) - - return [sparse.CSR(z_data, z_ind, z_ptr, p_shape)] + # This is a hack that works around some missing `Type`-related + # static shape narrowing. More specifically, + # `TensorType.convert_variable` currently won't combine the static + # shape information from `old_out.type` and `new_out.type`, only + # the broadcast patterns, and, since `CSR.make_node` doesn't do + # that either, we use `specify_shape` to produce an output `Type` + # with the same level of static shape information as the original + # `old_out`. + old_out = node.outputs[0] + new_out = specify_shape( + sparse.CSR(z_data, z_ind, z_ptr, p_shape), shape(old_out) + ) + return [new_out] return False diff --git a/aesara/sparse/type.py b/aesara/sparse/type.py index 3765919b4e..6a86bb2cfd 100644 --- a/aesara/sparse/type.py +++ b/aesara/sparse/type.py @@ -2,7 +2,8 @@ import scipy.sparse import aesara -from aesara.graph.type import HasDataType, Type +from aesara.graph.type import HasDataType +from aesara.tensor.type import TensorType def _is_sparse(x): @@ -24,7 +25,7 @@ def _is_sparse(x): return isinstance(x, scipy.sparse.spmatrix) -class SparseType(Type, HasDataType): +class SparseType(TensorType, HasDataType): """ Fundamental way to create a sparse node. @@ -52,19 +53,19 @@ class SparseType(Type, HasDataType): "csc": scipy.sparse.csc_matrix, "bsr": scipy.sparse.bsr_matrix, } - dtype_set = { - "int8", - "int16", - "int32", - "int64", - "float32", - "uint8", - "uint16", - "uint32", - "uint64", - "float64", - "complex64", - "complex128", + dtype_specs_map = { + "float32": (float, "npy_float32", "NPY_FLOAT32"), + "float64": (float, "npy_float64", "NPY_FLOAT64"), + "uint8": (int, "npy_uint8", "NPY_UINT8"), + "int8": (int, "npy_int8", "NPY_INT8"), + "uint16": (int, "npy_uint16", "NPY_UINT16"), + "int16": (int, "npy_int16", "NPY_INT16"), + "uint32": (int, "npy_uint32", "NPY_UINT32"), + "int32": (int, "npy_int32", "NPY_INT32"), + "uint64": (int, "npy_uint64", "NPY_UINT64"), + "int64": (int, "npy_int64", "NPY_INT64"), + "complex128": (complex, "aesara_complex128", "NPY_COMPLEX128"), + "complex64": (complex, "aesara_complex64", "NPY_COMPLEX64"), } ndim = 2 @@ -72,28 +73,25 @@ class SparseType(Type, HasDataType): variable_type = None Constant = None - def __init__(self, format, dtype, shape=None): - dtype = str(dtype) - if dtype in self.dtype_set: - self.dtype = dtype - else: - raise NotImplementedError( - f'unsupported dtype "{dtype}" not in list', list(self.dtype_set) - ) - + def __init__(self, format, dtype, shape=None, broadcastable=None, name=None): if shape is None: shape = (None, None) self.shape = shape - assert isinstance(format, str) + if not isinstance(format, str): + raise TypeError("The sparse format parameter must be a string") + if format in self.format_cls: self.format = format else: raise NotImplementedError( f'unsupported format "{format}" not in list', - list(self.format_cls.keys()), ) + if broadcastable is None: + broadcastable = [False, False] + + super().__init__(dtype, shape, name=name) def clone(self, format=None, dtype=None, shape=None, **kwargs): if format is None: @@ -153,21 +151,11 @@ def may_share_memory(a, b): def make_variable(self, name=None): return self.variable_type(self, name=name) - def __eq__(self, other): - return ( - type(self) == type(other) - and other.dtype == self.dtype - and other.format == self.format - ) - def __hash__(self): - return hash(self.dtype) ^ hash(self.format) - - def __str__(self): - return f"Sparse[{self.dtype}, {self.format}]" + return super().__hash__() ^ hash(self.format) def __repr__(self): - return f"Sparse[{self.dtype}, {self.format}]" + return f"Sparse({self.dtype}, {self.shape}, {self.format})" def values_eq_approx(self, a, b, eps=1e-6): # WARNING: equality comparison of sparse matrices is not fast or easy @@ -210,6 +198,31 @@ def get_size(self, shape_info): + (shape_info[2] + shape_info[3]) * np.dtype("int32").itemsize ) + def value_zeros(self, shape): + matrix_constructor = self.format_cls.get(self.format) + + if matrix_constructor is None: + raise ValueError(f"Sparse matrix type {self.format} not found in SciPy") + + return matrix_constructor(shape, dtype=self.dtype) + + def __eq__(self, other): + res = super().__eq__(other) + + if isinstance(res, bool): + return res and other.format == self.format + + return res + + def is_super(self, otype): + if not super().is_super(otype): + return False + + if self.format == otype.format: + return True + + return False + # Register SparseType's C code for ViewOp. aesara.compile.register_view_op_c_code( diff --git a/aesara/tensor/basic.py b/aesara/tensor/basic.py index 0398d54b69..a227a02446 100644 --- a/aesara/tensor/basic.py +++ b/aesara/tensor/basic.py @@ -313,8 +313,13 @@ def get_scalar_constant_value( return np.array(data.item(), dtype=v.dtype) except ValueError: raise NotScalarConstantError() - else: - return data + + from aesara.sparse.type import SparseType + + if isinstance(v.type, SparseType): + raise NotScalarConstantError() + + return data if not only_process_constants and getattr(v, "owner", None) and max_recur > 0: max_recur -= 1 diff --git a/aesara/tensor/basic_opt.py b/aesara/tensor/basic_opt.py index 073432ff70..eb9ded0e7c 100644 --- a/aesara/tensor/basic_opt.py +++ b/aesara/tensor/basic_opt.py @@ -78,7 +78,12 @@ from aesara.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape, shape_padleft from aesara.tensor.sort import TopKOp from aesara.tensor.subtensor import Subtensor, get_idx_list -from aesara.tensor.type import TensorType, discrete_dtypes, integer_dtypes +from aesara.tensor.type import ( + DenseTensorType, + TensorType, + discrete_dtypes, + integer_dtypes, +) from aesara.tensor.var import TensorConstant from aesara.utils import NoDuplicateOptWarningFilter @@ -2954,7 +2959,8 @@ def constant_folding(fgraph, node): # TODO: `Type` itself should provide an interface for constructing # instances appropriate for a given constant. - if isinstance(output.type, TensorType): + # TODO: Add handling for sparse types. + if isinstance(output.type, DenseTensorType): output_type = TensorType( output.type.dtype, tuple(s == 1 for s in data.shape), diff --git a/aesara/tensor/blas.py b/aesara/tensor/blas.py index ea27f7dcba..1d14f90b2f 100644 --- a/aesara/tensor/blas.py +++ b/aesara/tensor/blas.py @@ -167,7 +167,12 @@ from aesara.tensor.elemwise import DimShuffle, Elemwise from aesara.tensor.exceptions import NotScalarConstantError from aesara.tensor.math import Dot, add, mul, neg, sub -from aesara.tensor.type import integer_dtypes, tensor, values_eq_approx_remove_inf_nan +from aesara.tensor.type import ( + DenseTensorType, + integer_dtypes, + tensor, + values_eq_approx_remove_inf_nan, +) from aesara.utils import memoize @@ -264,7 +269,13 @@ def make_node(self, y, alpha, A, x, beta): raise TypeError("gemv requires vector for x", x.type) if y.ndim != 1: raise TypeError("gemv requires vector for y", y.type) - return Apply(self, [y, alpha, A, x, beta], [y.type()]) + + inputs = [y, alpha, A, x, beta] + + if any(not isinstance(i.type, DenseTensorType) for i in inputs): + raise NotImplementedError("Only dense tensor types are supported") + + return Apply(self, inputs, [y.type()]) def perform(self, node, inputs, out_storage, params=None): y, alpha, A, x, beta = inputs @@ -361,7 +372,12 @@ def make_node(self, A, alpha, x, y): if x.dtype not in ("float32", "float64", "complex64", "complex128"): raise TypeError("only float and complex types supported", x.dtype) - return Apply(self, [A, alpha, x, y], [A.type()]) + + inputs = [A, alpha, x, y] + if any(not isinstance(i.type, DenseTensorType) for i in inputs): + raise NotImplementedError("Only dense tensor types are supported") + + return Apply(self, inputs, [A.type()]) def perform(self, node, inp, out, params=None): cA, calpha, cx, cy = inp @@ -899,6 +915,10 @@ def __getstate__(self): def make_node(self, *inputs): inputs = list(map(at.as_tensor_variable, inputs)) + + if any(not isinstance(i.type, DenseTensorType) for i in inputs): + raise NotImplementedError("Only dense tensor types are supported") + if len(inputs) != 5: raise TypeError( f"Wrong number of inputs for {self} (expected 5, got {len(inputs)})" @@ -1580,6 +1600,10 @@ class Dot22(GemmRelated): def make_node(self, x, y): x = at.as_tensor_variable(x) y = at.as_tensor_variable(y) + + if any(not isinstance(i.type, DenseTensorType) for i in (x, y)): + raise NotImplementedError("Only dense tensor types are supported") + dtypes = ("float16", "float32", "float64", "complex64", "complex128") if x.type.ndim != 2 or x.type.dtype not in dtypes: raise TypeError(x) @@ -1665,6 +1689,9 @@ def local_dot_to_dot22(fgraph, node): if not isinstance(node.op, Dot): return + if any(not isinstance(i.type, DenseTensorType) for i in node.inputs): + return False + x, y = node.inputs if y.type.dtype != x.type.dtype: # TODO: upcast one so the types match @@ -1869,6 +1896,10 @@ class Dot22Scalar(GemmRelated): check_input = False def make_node(self, x, y, a): + + if any(not isinstance(i.type, DenseTensorType) for i in (x, y, a)): + raise NotImplementedError("Only dense tensor types are supported") + if a.ndim != 0: raise TypeError(Gemm.E_scalar, a) if x.ndim != 2: @@ -2089,6 +2120,9 @@ class BatchedDot(COp): def make_node(self, *inputs): inputs = list(map(at.as_tensor_variable, inputs)) + if any(not isinstance(i.type, DenseTensorType) for i in inputs): + raise NotImplementedError("Only dense tensor types are supported") + if len(inputs) != 2: raise TypeError(f"Two arguments required, but {len(inputs)} given.") if inputs[0].ndim not in (2, 3): diff --git a/aesara/tensor/math.py b/aesara/tensor/math.py index 8eafc95add..1e08442d59 100644 --- a/aesara/tensor/math.py +++ b/aesara/tensor/math.py @@ -34,6 +34,7 @@ ) from aesara.tensor.shape import shape from aesara.tensor.type import ( + DenseTensorType, complex_dtypes, continuous_dtypes, discrete_dtypes, @@ -2076,6 +2077,11 @@ def dense_dot(a, b): """ a, b = as_tensor_variable(a), as_tensor_variable(b) + if not isinstance(a.type, DenseTensorType) or not isinstance( + b.type, DenseTensorType + ): + raise TypeError("The dense dot product is only supported for dense types") + if a.ndim == 0 or b.ndim == 0: return a * b elif a.ndim > 2 or b.ndim > 2: diff --git a/aesara/tensor/shape.py b/aesara/tensor/shape.py index 1218b3510b..05d43af505 100644 --- a/aesara/tensor/shape.py +++ b/aesara/tensor/shape.py @@ -431,7 +431,7 @@ def make_node(self, x, shape): ) if isinstance(x.type, TensorType) and all(isinstance(s, Number) for s in shape): - out_var = TensorType(x.type.dtype, shape)() + out_var = x.type.clone(shape=shape)() else: out_var = x.type() diff --git a/aesara/tensor/type.py b/aesara/tensor/type.py index 9145758c0b..ae8917484b 100644 --- a/aesara/tensor/type.py +++ b/aesara/tensor/type.py @@ -1,6 +1,6 @@ import logging import warnings -from typing import Iterable, Optional, Union +from typing import Iterable, Optional, Tuple, Union import numpy as np @@ -9,6 +9,7 @@ from aesara.configdefaults import config from aesara.graph.basic import Variable from aesara.graph.type import HasDataType +from aesara.graph.utils import MetaType from aesara.link.c.type import CType from aesara.misc.safe_asarray import _asarray from aesara.utils import apply_across_args @@ -50,8 +51,9 @@ class TensorType(CType, HasDataType): r"""Symbolic `Type` representing `numpy.ndarray`\s.""" - __props__ = ("dtype", "shape") + __props__: Tuple[str, ...] = ("dtype", "shape") + dtype_specs_map = dtype_specs_map context_name = "cpu" filter_checks_isfinite = False """ @@ -271,7 +273,7 @@ def dtype_specs(self): """ try: - return dtype_specs_map[self.dtype] + return self.dtype_specs_map[self.dtype] except KeyError: raise TypeError( f"Unsupported dtype for {self.__class__.__name__}: {self.dtype}" @@ -613,6 +615,20 @@ def c_code_cache_version(self): return () +class DenseTypeMeta(MetaType): + def __instancecheck__(self, o): + if type(o) == TensorType or isinstance(o, DenseTypeMeta): + return True + return False + + +class DenseTensorType(TensorType, metaclass=DenseTypeMeta): + r"""A `Type` for dense tensors. + + Instances of this class and `TensorType`\s are considered dense `Type`\s. + """ + + def values_eq_approx( a, b, allow_remove_inf=False, allow_remove_nan=False, rtol=None, atol=None ): diff --git a/aesara/tensor/var.py b/aesara/tensor/var.py index 911d4090a4..f3c25435fd 100644 --- a/aesara/tensor/var.py +++ b/aesara/tensor/var.py @@ -10,6 +10,7 @@ from aesara import tensor as at from aesara.configdefaults import config from aesara.graph.basic import Constant, Variable +from aesara.graph.utils import MetaType from aesara.scalar import ComplexError, IntegerDivisionError from aesara.tensor import _get_vector_length, as_tensor_variable from aesara.tensor.exceptions import AdvancedIndexingError @@ -1040,3 +1041,33 @@ def __deepcopy__(self, memo): TensorType.constant_type = TensorConstant + + +class DenseVariableMeta(MetaType): + def __instancecheck__(self, o): + if type(o) == TensorVariable or isinstance(o, DenseVariableMeta): + return True + return False + + +class DenseTensorVariable(TensorType, metaclass=DenseVariableMeta): + r"""A `Variable` for dense tensors. + + Instances of this class and `TensorVariable`\s are considered dense + `Variable`\s. + """ + + +class DenseConstantMeta(MetaType): + def __instancecheck__(self, o): + if type(o) == TensorConstant or isinstance(o, DenseConstantMeta): + return True + return False + + +class DenseTensorConstant(TensorType, metaclass=DenseConstantMeta): + r"""A `Constant` for dense tensors. + + Instances of this class and `TensorConstant`\s are considered dense + `Constant`\s. + """ diff --git a/tests/sparse/test_basic.py b/tests/sparse/test_basic.py index 2d8281efd0..2a1d6dbe5c 100644 --- a/tests/sparse/test_basic.py +++ b/tests/sparse/test_basic.py @@ -12,7 +12,7 @@ from aesara.compile.io import In, Out from aesara.configdefaults import config from aesara.gradient import GradientError -from aesara.graph.basic import Apply, Constant +from aesara.graph.basic import Apply, Constant, applys_between from aesara.graph.op import Op from aesara.misc.safe_asarray import _asarray from aesara.sparse import ( @@ -78,6 +78,7 @@ true_dot, ) from aesara.sparse.basic import ( + SparseConstant, _is_dense_variable, _is_sparse, _is_sparse_variable, @@ -1017,22 +1018,45 @@ def test_equality_case(self): class TestConversion: - @pytest.mark.skip def test_basic(self): - a = at.as_tensor_variable(np.random.random((5))) + test_val = np.random.rand(5).astype(config.floatX) + a = at.as_tensor_variable(test_val) s = csc_from_dense(a) val = eval_outputs([s]) - assert str(val.dtype) == "float64" + assert str(val.dtype) == config.floatX assert val.format == "csc" - @pytest.mark.skip - def test_basic_1(self): - a = at.as_tensor_variable(np.random.random((5))) + a = at.as_tensor_variable(test_val) s = csr_from_dense(a) val = eval_outputs([s]) - assert str(val.dtype) == "float64" + assert str(val.dtype) == config.floatX assert val.format == "csr" + test_val = np.eye(3).astype(config.floatX) + a = sp.sparse.csr_matrix(test_val) + s = as_sparse_or_tensor_variable(a) + res = at.as_tensor_variable(s) + assert isinstance(res, SparseConstant) + + a = sp.sparse.csr_matrix(test_val) + s = as_sparse_or_tensor_variable(a) + from aesara.tensor.exceptions import NotScalarConstantError + + with pytest.raises(NotScalarConstantError): + at.get_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: @@ -1591,6 +1615,32 @@ def test_int32_dtype(self): ) f(i, a) + def test_tensor_dot_types(self): + + x = sparse.csc_matrix("x") + x_d = at.matrix("x_d") + y = sparse.csc_matrix("y") + + res = at.dot(x, y) + op_types = set(type(n.op) for n in applys_between([x, y], [res])) + assert sparse.basic.StructuredDot in op_types + assert at.math.Dot not in op_types + + res = at.dot(x_d, y) + op_types = set(type(n.op) for n in applys_between([x, y], [res])) + assert sparse.basic.StructuredDot in op_types + assert at.math.Dot not in op_types + + res = at.dot(x, x_d) + op_types = set(type(n.op) for n in applys_between([x, y], [res])) + assert sparse.basic.StructuredDot in op_types + assert at.math.Dot not in op_types + + res = at.dot(at.second(1, x), y) + op_types = set(type(n.op) for n in applys_between([x, y], [res])) + assert sparse.basic.StructuredDot in op_types + assert at.math.Dot not in op_types + def test_csr_dense_grad(self): # shortcut: testing csc in float32, testing csr in float64 diff --git a/tests/sparse/test_type.py b/tests/sparse/test_type.py index 749c6ddadf..28ec9a758d 100644 --- a/tests/sparse/test_type.py +++ b/tests/sparse/test_type.py @@ -1,6 +1,28 @@ +import pytest + +from aesara.sparse import matrix as sp_matrix from aesara.sparse.type import SparseType +from aesara.tensor import dmatrix def test_clone(): st = SparseType("csr", "float64") assert st == st.clone() + + +def test_Sparse_convert_variable(): + x = dmatrix(name="x") + y = sp_matrix("csc", dtype="float64", name="y") + z = sp_matrix("csr", dtype="float64", name="z") + + assert y.type.convert_variable(z) is None + + # TODO FIXME: This is a questionable result, because `x.type` is associated + # with a dense `Type`, but, since `TensorType` is a base class of `Sparse`, + # we would need to added sparse/dense logic to `TensorType`, and we don't + # want to do that. + assert x.type.convert_variable(y) is y + + # TODO FIXME: We should be able to do this. + with pytest.raises(NotImplementedError): + y.type.convert_variable(x) diff --git a/tests/tensor/test_var.py b/tests/tensor/test_var.py index 4e57fd1c14..85429a399f 100644 --- a/tests/tensor/test_var.py +++ b/tests/tensor/test_var.py @@ -6,6 +6,7 @@ import tests.unittest_tools as utt from aesara.graph.basic import Constant, equal_computations from aesara.tensor import get_vector_length +from aesara.tensor.basic import constant from aesara.tensor.elemwise import DimShuffle from aesara.tensor.math import dot from aesara.tensor.subtensor import AdvancedSubtensor, Subtensor @@ -21,7 +22,12 @@ tensor3, ) from aesara.tensor.type_other import MakeSlice -from aesara.tensor.var import TensorConstant, TensorVariable +from aesara.tensor.var import ( + DenseTensorConstant, + DenseTensorVariable, + TensorConstant, + TensorVariable, +) @pytest.mark.parametrize( @@ -247,3 +253,14 @@ def test_get_vector_length(): x = TensorVariable(TensorType("int64", (None,))) with pytest.raises(ValueError): get_vector_length(x) + + +def test_dense_types(): + + x = matrix() + assert isinstance(x, DenseTensorVariable) + assert not isinstance(x, DenseTensorConstant) + + x = constant(1) + assert not isinstance(x, DenseTensorVariable) + assert isinstance(x, DenseTensorConstant) diff --git a/tests/test_ifelse.py b/tests/test_ifelse.py index 5751e1ea56..218409ed18 100644 --- a/tests/test_ifelse.py +++ b/tests/test_ifelse.py @@ -335,13 +335,13 @@ def test_sparse_tensor_error(self): z = aesara.sparse.matrix("csr", dtype=self.dtype, name="z") cond = iscalar("cond") - with pytest.raises(TypeError): + with pytest.raises(NotImplementedError): ifelse(cond, x, y) - with pytest.raises(TypeError): + with pytest.raises(NotImplementedError): ifelse(cond, y, x) - with pytest.raises(TypeError): + with pytest.raises(NotImplementedError): ifelse(cond, x, z) - with pytest.raises(TypeError): + with pytest.raises(NotImplementedError): ifelse(cond, z, x) with pytest.raises(TypeError): ifelse(cond, y, z) From 8eaeee3dd84775abda45976c61ad8fcd527ea3a3 Mon Sep 17 00:00:00 2001 From: Anatoly Date: Fri, 28 Jan 2022 14:53:49 +0300 Subject: [PATCH 2/3] Warn when dense conversion is used by sparse tensor methods --- aesara/sparse/basic.py | 89 +++++++++++++++ aesara/sparse/type.py | 12 +- tests/sparse/test_var.py | 234 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 333 insertions(+), 2 deletions(-) create mode 100644 tests/sparse/test_var.py diff --git a/aesara/sparse/basic.py b/aesara/sparse/basic.py index f0df5c1cb0..b27ae5bef8 100644 --- a/aesara/sparse/basic.py +++ b/aesara/sparse/basic.py @@ -250,6 +250,95 @@ 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, SparseType) + 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", + "conj", + "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)" diff --git a/aesara/sparse/type.py b/aesara/sparse/type.py index 6a86bb2cfd..771a445450 100644 --- a/aesara/sparse/type.py +++ b/aesara/sparse/type.py @@ -148,8 +148,16 @@ def may_share_memory(a, b): return True return False - def make_variable(self, name=None): - return self.variable_type(self, name=name) + def convert_variable(self, var): + res = super().convert_variable(var) + + if res and not isinstance(res.type, SparseType): + # TODO: Convert to this sparse format + raise NotImplementedError() + + # TODO: Convert sparse `var`s with different formats to this format? + + return res def __hash__(self): return super().__hash__() ^ hash(self.format) diff --git a/tests/sparse/test_var.py b/tests/sparse/test_var.py new file mode 100644 index 0000000000..04f571eb74 --- /dev/null +++ b/tests/sparse/test_var.py @@ -0,0 +1,234 @@ +from contextlib import ExitStack + +import numpy as np +import pytest +from scipy.sparse.csr import csr_matrix + +import aesara +import aesara.sparse as sparse +import aesara.tensor as at +from aesara.sparse.type import SparseType +from aesara.tensor.type import DenseTensorType + + +class TestSparseVariable: + @pytest.mark.parametrize( + "method, exp_type, cm", + [ + ("__abs__", DenseTensorType, None), + ("__neg__", SparseType, ExitStack()), + ("__ceil__", DenseTensorType, None), + ("__floor__", DenseTensorType, None), + ("__trunc__", DenseTensorType, None), + ("transpose", DenseTensorType, None), + ("any", DenseTensorType, None), + ("all", DenseTensorType, None), + ("flatten", DenseTensorType, None), + ("ravel", DenseTensorType, None), + ("arccos", DenseTensorType, None), + ("arcsin", DenseTensorType, None), + ("arctan", DenseTensorType, None), + ("arccosh", DenseTensorType, None), + ("arcsinh", DenseTensorType, None), + ("arctanh", DenseTensorType, None), + ("ceil", DenseTensorType, None), + ("cos", DenseTensorType, None), + ("cosh", DenseTensorType, None), + ("deg2rad", DenseTensorType, None), + ("exp", DenseTensorType, None), + ("exp2", DenseTensorType, None), + ("expm1", DenseTensorType, None), + ("floor", DenseTensorType, None), + ("log", DenseTensorType, None), + ("log10", DenseTensorType, None), + ("log1p", DenseTensorType, None), + ("log2", DenseTensorType, None), + ("rad2deg", DenseTensorType, None), + ("sin", DenseTensorType, None), + ("sinh", DenseTensorType, None), + ("sqrt", DenseTensorType, None), + ("tan", DenseTensorType, None), + ("tanh", DenseTensorType, None), + ("copy", DenseTensorType, None), + ("sum", DenseTensorType, ExitStack()), + ("prod", DenseTensorType, None), + ("mean", DenseTensorType, None), + ("var", DenseTensorType, None), + ("std", DenseTensorType, None), + ("min", DenseTensorType, None), + ("max", DenseTensorType, None), + ("argmin", DenseTensorType, None), + ("argmax", DenseTensorType, None), + ("nonzero", DenseTensorType, ExitStack()), + ("nonzero_values", DenseTensorType, None), + ("argsort", DenseTensorType, ExitStack()), + ("conj", DenseTensorType, None), + ("round", DenseTensorType, None), + ("trace", DenseTensorType, None), + ("zeros_like", SparseType, ExitStack()), + ("ones_like", DenseTensorType, ExitStack()), + ("cumsum", DenseTensorType, None), + ("cumprod", DenseTensorType, None), + ("ptp", DenseTensorType, None), + ("squeeze", DenseTensorType, None), + ("diagonal", DenseTensorType, None), + ], + ) + def test_unary(self, method, exp_type, cm): + x = at.dmatrix("x") + x = sparse.csr_from_dense(x) + + method_to_call = getattr(x, method) + + if cm is None: + cm = pytest.warns(UserWarning, match=".*converted to dense.*") + + if exp_type == SparseType: + exp_res_type = csr_matrix + else: + exp_res_type = np.ndarray + + with cm: + z = method_to_call() + + if not isinstance(z, tuple): + z_outs = (z,) + else: + z_outs = z + + assert all(isinstance(out.type, exp_type) for out in z_outs) + + f = aesara.function([x], z, on_unused_input="ignore") + + res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) + + if not isinstance(res, list): + res_outs = [res] + else: + res_outs = res + + assert all(isinstance(out, exp_res_type) for out in res_outs) + + @pytest.mark.parametrize( + "method, exp_type", + [ + ("__lt__", SparseType), + ("__le__", SparseType), + ("__gt__", SparseType), + ("__ge__", SparseType), + ("__and__", DenseTensorType), + ("__or__", DenseTensorType), + ("__xor__", DenseTensorType), + ("__add__", SparseType), + ("__sub__", SparseType), + ("__mul__", SparseType), + ("__pow__", DenseTensorType), + ("__mod__", DenseTensorType), + ("__divmod__", DenseTensorType), + ("__truediv__", DenseTensorType), + ("__floordiv__", DenseTensorType), + ], + ) + def test_binary(self, method, exp_type): + x = at.lmatrix("x") + y = at.lmatrix("y") + x = sparse.csr_from_dense(x) + y = sparse.csr_from_dense(y) + + method_to_call = getattr(x, method) + + if exp_type == SparseType: + exp_res_type = csr_matrix + cm = ExitStack() + else: + exp_res_type = np.ndarray + cm = pytest.warns(UserWarning, match=".*converted to dense.*") + + with cm: + z = method_to_call(y) + + if not isinstance(z, tuple): + z_outs = (z,) + else: + z_outs = z + + assert all(isinstance(out.type, exp_type) for out in z_outs) + + f = aesara.function([x, y], z) + res = f( + [[1, 0, 2], [-1, 0, 0]], + [[1, 1, 2], [1, 4, 1]], + ) + + if not isinstance(res, list): + res_outs = [res] + else: + res_outs = res + + assert all(isinstance(out, exp_res_type) for out in res_outs) + + def test_reshape(self): + x = at.dmatrix("x") + x = sparse.csr_from_dense(x) + + with pytest.warns(UserWarning, match=".*converted to dense.*"): + z = x.reshape((3, 2)) + + assert isinstance(z.type, DenseTensorType) + + f = aesara.function([x], z) + exp_res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) + assert isinstance(exp_res, np.ndarray) + + def test_dimshuffle(self): + x = at.dmatrix("x") + x = sparse.csr_from_dense(x) + + with pytest.warns(UserWarning, match=".*converted to dense.*"): + z = x.dimshuffle((1, 0)) + + assert isinstance(z.type, DenseTensorType) + + f = aesara.function([x], z) + exp_res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) + assert isinstance(exp_res, np.ndarray) + + def test_getitem(self): + x = at.dmatrix("x") + x = sparse.csr_from_dense(x) + + z = x[:, :2] + assert isinstance(z.type, SparseType) + + f = aesara.function([x], z) + exp_res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) + assert isinstance(exp_res, csr_matrix) + + def test_dot(self): + x = at.lmatrix("x") + y = at.lmatrix("y") + x = sparse.csr_from_dense(x) + y = sparse.csr_from_dense(y) + + z = x.__dot__(y) + assert isinstance(z.type, SparseType) + + f = aesara.function([x, y], z) + exp_res = f( + [[1, 0, 2], [-1, 0, 0]], + [[-1], [2], [1]], + ) + assert isinstance(exp_res, csr_matrix) + + def test_repeat(self): + x = at.dmatrix("x") + x = sparse.csr_from_dense(x) + + with pytest.warns(UserWarning, match=".*converted to dense.*"): + z = x.repeat(2, axis=1) + + assert isinstance(z.type, DenseTensorType) + + f = aesara.function([x], z) + exp_res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) + assert isinstance(exp_res, np.ndarray) From d7179fc5e028cb4daf5cf1f1babb568687fc0010 Mon Sep 17 00:00:00 2001 From: "Brandon T. Willard" Date: Tue, 15 Mar 2022 17:36:59 -0500 Subject: [PATCH 3/3] Rename SparseType to SparseTensorType --- aesara/__init__.py | 2 +- aesara/link/c/type.py | 2 +- aesara/misc/may_share_memory.py | 4 +- aesara/sparse/__init__.py | 2 +- aesara/sparse/basic.py | 96 ++++++++++++++++------------ aesara/sparse/sandbox/sp2.py | 6 +- aesara/sparse/sharedvar.py | 6 +- aesara/sparse/type.py | 28 ++++---- aesara/tensor/basic.py | 4 +- doc/extending/other_ops.rst | 6 +- tests/compile/function/test_pfunc.py | 4 +- tests/sparse/test_basic.py | 80 +++++++++++------------ tests/sparse/test_type.py | 4 +- tests/sparse/test_var.py | 28 ++++---- tests/typed_list/test_basic.py | 4 +- 15 files changed, 145 insertions(+), 131 deletions(-) diff --git a/aesara/__init__.py b/aesara/__init__.py index b0c9c2c805..2e50c79392 100644 --- a/aesara/__init__.py +++ b/aesara/__init__.py @@ -167,7 +167,7 @@ def get_scalar_constant_value(v): """ # Is it necessary to test for presence of aesara.sparse at runtime? sparse = globals().get("sparse") - if sparse and isinstance(v.type, sparse.SparseType): + if sparse and isinstance(v.type, sparse.SparseTensorType): if v.owner is not None and isinstance(v.owner.op, sparse.CSM): data = v.owner.inputs[0] return tensor.get_scalar_constant_value(data) diff --git a/aesara/link/c/type.py b/aesara/link/c/type.py index 1920f36e61..4696b3333b 100644 --- a/aesara/link/c/type.py +++ b/aesara/link/c/type.py @@ -17,7 +17,7 @@ class CType(Type, CLinkerType): - `TensorType`: for numpy.ndarray - - `SparseType`: for scipy.sparse + - `SparseTensorType`: for scipy.sparse But you are encouraged to write your own, as described in WRITEME. diff --git a/aesara/misc/may_share_memory.py b/aesara/misc/may_share_memory.py index b522834e33..89545d8e44 100644 --- a/aesara/misc/may_share_memory.py +++ b/aesara/misc/may_share_memory.py @@ -12,7 +12,7 @@ try: import scipy.sparse - from aesara.sparse.basic import SparseType + from aesara.sparse.basic import SparseTensorType def _is_sparse(a): return scipy.sparse.issparse(a) @@ -64,4 +64,4 @@ def may_share_memory(a, b, raise_other_type=True): if a_gpua or b_gpua: return False - return SparseType.may_share_memory(a, b) + return SparseTensorType.may_share_memory(a, b) diff --git a/aesara/sparse/__init__.py b/aesara/sparse/__init__.py index 7b6873105f..376da5dcce 100644 --- a/aesara/sparse/__init__.py +++ b/aesara/sparse/__init__.py @@ -9,7 +9,7 @@ enable_sparse = False warn("SciPy can't be imported. Sparse matrix support is disabled.") -from aesara.sparse.type import SparseType, _is_sparse +from aesara.sparse.type import SparseTensorType, _is_sparse if enable_sparse: diff --git a/aesara/sparse/basic.py b/aesara/sparse/basic.py index b27ae5bef8..ba80e38db7 100644 --- a/aesara/sparse/basic.py +++ b/aesara/sparse/basic.py @@ -22,7 +22,7 @@ from aesara.link.c.op import COp from aesara.link.c.type import generic from aesara.misc.safe_asarray import _asarray -from aesara.sparse.type import SparseType, _is_sparse +from aesara.sparse.type import SparseTensorType, _is_sparse from aesara.sparse.utils import hash_from_sparse from aesara.tensor import basic as at from aesara.tensor.basic import Split @@ -80,11 +80,11 @@ def _is_sparse_variable(x): if not isinstance(x, Variable): raise NotImplementedError( "this function should only be called on " - "*variables* (of type sparse.SparseType " + "*variables* (of type sparse.SparseTensorType " "or TensorType, for instance), not ", x, ) - return isinstance(x.type, SparseType) + return isinstance(x.type, SparseTensorType) def _is_dense_variable(x): @@ -100,7 +100,7 @@ def _is_dense_variable(x): if not isinstance(x, Variable): raise NotImplementedError( "this function should only be called on " - "*variables* (of type sparse.SparseType or " + "*variables* (of type sparse.SparseTensorType or " "TensorType, for instance), not ", x, ) @@ -159,13 +159,15 @@ def as_sparse_variable(x, name=None, ndim=None, **kwargs): else: x = x.outputs[0] if isinstance(x, Variable): - if not isinstance(x.type, SparseType): - raise TypeError("Variable type field must be a SparseType.", x, x.type) + if not isinstance(x.type, SparseTensorType): + raise TypeError( + "Variable type field must be a SparseTensorType.", x, x.type + ) return x try: return constant(x, name=name) except TypeError: - raise TypeError(f"Cannot convert {x} to SparseType", type(x)) + raise TypeError(f"Cannot convert {x} to SparseTensorType", type(x)) as_sparse = as_sparse_variable @@ -198,10 +200,10 @@ def constant(x, name=None): raise TypeError("sparse.constant must be called on a " "scipy.sparse.spmatrix") try: return SparseConstant( - SparseType(format=x.format, dtype=x.dtype), x.copy(), name=name + SparseTensorType(format=x.format, dtype=x.dtype), x.copy(), name=name ) except TypeError: - raise TypeError(f"Could not convert {x} to SparseType", type(x)) + raise TypeError(f"Could not convert {x} to SparseTensorType", type(x)) def sp_ones_like(x): @@ -259,7 +261,7 @@ def to_dense(self, *args, **kwargs): self = self.toarray() new_args = [ arg.toarray() - if hasattr(arg, "type") and isinstance(arg.type, SparseType) + if hasattr(arg, "type") and isinstance(arg.type, SparseTensorType) else arg for arg in args ] @@ -503,15 +505,15 @@ def __repr__(self): return str(self) -SparseType.variable_type = SparseVariable -SparseType.constant_type = SparseConstant +SparseTensorType.variable_type = SparseVariable +SparseTensorType.constant_type = SparseConstant -# for more dtypes, call SparseType(format, dtype) +# for more dtypes, call SparseTensorType(format, dtype) def matrix(format, name=None, dtype=None): if dtype is None: dtype = config.floatX - type = SparseType(format=format, dtype=dtype) + type = SparseTensorType(format=format, dtype=dtype) return type(name) @@ -527,15 +529,15 @@ def bsr_matrix(name=None, dtype=None): return matrix("bsr", name, dtype) -# for more dtypes, call SparseType(format, dtype) -csc_dmatrix = SparseType(format="csc", dtype="float64") -csr_dmatrix = SparseType(format="csr", dtype="float64") -bsr_dmatrix = SparseType(format="bsr", dtype="float64") -csc_fmatrix = SparseType(format="csc", dtype="float32") -csr_fmatrix = SparseType(format="csr", dtype="float32") -bsr_fmatrix = SparseType(format="bsr", dtype="float32") +# for more dtypes, call SparseTensorType(format, dtype) +csc_dmatrix = SparseTensorType(format="csc", dtype="float64") +csr_dmatrix = SparseTensorType(format="csr", dtype="float64") +bsr_dmatrix = SparseTensorType(format="bsr", dtype="float64") +csc_fmatrix = SparseTensorType(format="csc", dtype="float32") +csr_fmatrix = SparseTensorType(format="csr", dtype="float32") +bsr_fmatrix = SparseTensorType(format="bsr", dtype="float32") -all_dtypes = list(SparseType.dtype_specs_map.keys()) +all_dtypes = list(SparseTensorType.dtype_specs_map.keys()) complex_dtypes = [t for t in all_dtypes if t[:7] == "complex"] float_dtypes = [t for t in all_dtypes if t[:5] == "float"] int_dtypes = [t for t in all_dtypes if t[:3] == "int"] @@ -725,7 +727,7 @@ def make_node(self, data, indices, indptr, shape): return Apply( self, [data, indices, indptr, shape], - [SparseType(dtype=data.type.dtype, format=self.format)()], + [SparseTensorType(dtype=data.type.dtype, format=self.format)()], ) def perform(self, node, inputs, outputs): @@ -931,7 +933,9 @@ def __init__(self, out_type): def make_node(self, x): x = as_sparse_variable(x) assert x.format in ("csr", "csc") - return Apply(self, [x], [SparseType(dtype=self.out_type, format=x.format)()]) + return Apply( + self, [x], [SparseTensorType(dtype=self.out_type, format=x.format)()] + ) def perform(self, node, inputs, outputs): (x,) = inputs @@ -1014,7 +1018,7 @@ def __str__(self): return f"{self.__class__.__name__}{{structured_grad={self.sparse_grad}}}" def __call__(self, x): - if not isinstance(x.type, SparseType): + if not isinstance(x.type, SparseTensorType): return x return super().__call__(x) @@ -1097,7 +1101,7 @@ def __str__(self): return f"{self.__class__.__name__}{{{self.format}}}" def __call__(self, x): - if isinstance(x.type, SparseType): + if isinstance(x.type, SparseTensorType): return x return super().__call__(x) @@ -1116,12 +1120,14 @@ def make_node(self, x): else: assert x.ndim == 2 - return Apply(self, [x], [SparseType(dtype=x.type.dtype, format=self.format)()]) + return Apply( + self, [x], [SparseTensorType(dtype=x.type.dtype, format=self.format)()] + ) def perform(self, node, inputs, outputs): (x,) = inputs (out,) = outputs - out[0] = SparseType.format_cls[self.format](x) + out[0] = SparseTensorType.format_cls[self.format](x) def grad(self, inputs, gout): (x,) = inputs @@ -1585,7 +1591,11 @@ def make_node(self, x): return Apply( self, [x], - [SparseType(dtype=x.type.dtype, format=self.format_map[x.type.format])()], + [ + SparseTensorType( + dtype=x.type.dtype, format=self.format_map[x.type.format] + )() + ], ) def perform(self, node, inputs, outputs): @@ -2002,7 +2012,7 @@ def make_node(self, diag): if diag.type.ndim != 1: raise TypeError("data argument must be a vector", diag.type) - return Apply(self, [diag], [SparseType(dtype=diag.dtype, format="csc")()]) + return Apply(self, [diag], [SparseTensorType(dtype=diag.dtype, format="csc")()]) def perform(self, node, inputs, outputs): (z,) = outputs @@ -2146,7 +2156,7 @@ def make_node(self, x, y): assert y.format in ("csr", "csc") out_dtype = aes.upcast(x.type.dtype, y.type.dtype) return Apply( - self, [x, y], [SparseType(dtype=out_dtype, format=x.type.format)()] + self, [x, y], [SparseTensorType(dtype=out_dtype, format=x.type.format)()] ) def perform(self, node, inputs, outputs): @@ -2183,7 +2193,7 @@ def make_node(self, x, y): if x.type.format != y.type.format: raise NotImplementedError() return Apply( - self, [x, y], [SparseType(dtype=x.type.dtype, format=x.type.format)()] + self, [x, y], [SparseTensorType(dtype=x.type.dtype, format=x.type.format)()] ) def perform(self, node, inputs, outputs): @@ -2286,7 +2296,7 @@ def make_node(self, x, y): if x.type.dtype != y.type.dtype: raise NotImplementedError() return Apply( - self, [x, y], [SparseType(dtype=x.type.dtype, format=x.type.format)()] + self, [x, y], [SparseTensorType(dtype=x.type.dtype, format=x.type.format)()] ) def perform(self, node, inputs, outputs): @@ -2426,7 +2436,7 @@ def make_node(self, x, y): assert y.format in ("csr", "csc") out_dtype = aes.upcast(x.type.dtype, y.type.dtype) return Apply( - self, [x, y], [SparseType(dtype=out_dtype, format=x.type.format)()] + self, [x, y], [SparseTensorType(dtype=out_dtype, format=x.type.format)()] ) def perform(self, node, inputs, outputs): @@ -2469,7 +2479,7 @@ def make_node(self, x, y): # 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 = SparseType(dtype=dtype, format=x.type.format)() + out = SparseTensorType(dtype=dtype, format=x.type.format)() return Apply(self, [x, y], [out]) def perform(self, node, inputs, outputs): @@ -2559,7 +2569,7 @@ def make_node(self, x, y): f"Got {x.type.dtype} and {y.type.dtype}." ) return Apply( - self, [x, y], [SparseType(dtype=x.type.dtype, format=x.type.format)()] + self, [x, y], [SparseTensorType(dtype=x.type.dtype, format=x.type.format)()] ) def perform(self, node, inputs, outputs): @@ -2694,7 +2704,9 @@ def make_node(self, x, y): if x.type.format != y.type.format: raise NotImplementedError() - return Apply(self, [x, y], [SparseType(dtype="uint8", format=x.type.format)()]) + return Apply( + self, [x, y], [SparseTensorType(dtype="uint8", format=x.type.format)()] + ) def perform(self, node, inputs, outputs): (x, y) = inputs @@ -3050,7 +3062,9 @@ def make_node(self, *mat): for x in var: assert x.format in ("csr", "csc") - return Apply(self, var, [SparseType(dtype=self.dtype, format=self.format)()]) + return Apply( + self, var, [SparseTensorType(dtype=self.dtype, format=self.format)()] + ) def perform(self, node, block, outputs): (out,) = outputs @@ -3578,7 +3592,7 @@ def make_node(self, x, y): raise NotImplementedError() inputs = [x, y] # Need to convert? e.g. assparse - outputs = [SparseType(dtype=x.type.dtype, format=myformat)()] + outputs = [SparseTensorType(dtype=x.type.dtype, format=myformat)()] return Apply(self, inputs, outputs) def perform(self, node, inp, out_): @@ -3702,7 +3716,7 @@ def make_node(self, a, b): raise NotImplementedError("non-matrix b") if _is_sparse_variable(b): - return Apply(self, [a, b], [SparseType(a.type.format, dtype_out)()]) + return Apply(self, [a, b], [SparseTensorType(a.type.format, dtype_out)()]) else: return Apply( self, @@ -3719,7 +3733,7 @@ def perform(self, node, inputs, outputs): ) variable = a * b - if isinstance(node.outputs[0].type, SparseType): + if isinstance(node.outputs[0].type, SparseTensorType): assert _is_sparse(variable) out[0] = variable return diff --git a/aesara/sparse/sandbox/sp2.py b/aesara/sparse/sandbox/sp2.py index 95851887c4..e86db84ae4 100644 --- a/aesara/sparse/sandbox/sp2.py +++ b/aesara/sparse/sandbox/sp2.py @@ -7,7 +7,7 @@ from aesara.graph.op import Op from aesara.sparse.basic import ( Remove0, - SparseType, + SparseTensorType, _is_sparse, as_sparse_variable, remove0, @@ -108,7 +108,9 @@ def make_node(self, n, p, shape): assert shape.dtype in discrete_dtypes return Apply( - self, [n, p, shape], [SparseType(dtype=self.dtype, format=self.format)()] + self, + [n, p, shape], + [SparseTensorType(dtype=self.dtype, format=self.format)()], ) def perform(self, node, inputs, outputs): diff --git a/aesara/sparse/sharedvar.py b/aesara/sparse/sharedvar.py index eb512bd678..d0a681ec96 100644 --- a/aesara/sparse/sharedvar.py +++ b/aesara/sparse/sharedvar.py @@ -3,7 +3,7 @@ import scipy.sparse from aesara.compile import SharedVariable, shared_constructor -from aesara.sparse.basic import SparseType, _sparse_py_operators +from aesara.sparse.basic import SparseTensorType, _sparse_py_operators class SparseTensorSharedVariable(_sparse_py_operators, SharedVariable): @@ -16,7 +16,7 @@ def sparse_constructor( value, name=None, strict=False, allow_downcast=None, borrow=False, format=None ): """ - SharedVariable Constructor for SparseType. + SharedVariable Constructor for SparseTensorType. writeme @@ -29,7 +29,7 @@ def sparse_constructor( if format is None: format = value.format - type = SparseType(format=format, dtype=value.dtype) + type = SparseTensorType(format=format, dtype=value.dtype) if not borrow: value = copy.deepcopy(value) return SparseTensorSharedVariable( diff --git a/aesara/sparse/type.py b/aesara/sparse/type.py index 771a445450..8812b76c98 100644 --- a/aesara/sparse/type.py +++ b/aesara/sparse/type.py @@ -25,9 +25,8 @@ def _is_sparse(x): return isinstance(x, scipy.sparse.spmatrix) -class SparseType(TensorType, HasDataType): - """ - Fundamental way to create a sparse node. +class SparseTensorType(TensorType, HasDataType): + """A `Type` for sparse tensors. Parameters ---------- @@ -42,8 +41,7 @@ class SparseType(TensorType, HasDataType): Notes ----- - As far as I can tell, L{scipy.sparse} objects must be matrices, i.e. - have dimension 2. + Currently, sparse tensors can only be matrices (i.e. have two dimensions). """ @@ -126,15 +124,13 @@ def filter(self, value, strict=False, allow_downcast=None): raise NotImplementedError() return sp - @staticmethod - def may_share_memory(a, b): - # This is Fred suggestion for a quick and dirty way of checking - # aliasing .. this can potentially be further refined (ticket #374) + @classmethod + def may_share_memory(cls, a, b): if _is_sparse(a) and _is_sparse(b): return ( - SparseType.may_share_memory(a, b.data) - or SparseType.may_share_memory(a, b.indices) - or SparseType.may_share_memory(a, b.indptr) + cls.may_share_memory(a, b.data) + or cls.may_share_memory(a, b.indices) + or cls.may_share_memory(a, b.indptr) ) if _is_sparse(b) and isinstance(a, np.ndarray): a, b = b, a @@ -151,7 +147,7 @@ def may_share_memory(a, b): def convert_variable(self, var): res = super().convert_variable(var) - if res and not isinstance(res.type, SparseType): + if res and not isinstance(res.type, type(self)): # TODO: Convert to this sparse format raise NotImplementedError() @@ -232,9 +228,8 @@ def is_super(self, otype): return False -# Register SparseType's C code for ViewOp. aesara.compile.register_view_op_c_code( - SparseType, + SparseTensorType, """ Py_XDECREF(%(oname)s); %(oname)s = %(iname)s; @@ -242,3 +237,6 @@ def is_super(self, otype): """, 1, ) + +# This is a deprecated alias used for (temporary) backward-compatibility +SparseType = SparseTensorType diff --git a/aesara/tensor/basic.py b/aesara/tensor/basic.py index a227a02446..468c332849 100644 --- a/aesara/tensor/basic.py +++ b/aesara/tensor/basic.py @@ -314,9 +314,9 @@ def get_scalar_constant_value( except ValueError: raise NotScalarConstantError() - from aesara.sparse.type import SparseType + from aesara.sparse.type import SparseTensorType - if isinstance(v.type, SparseType): + if isinstance(v.type, SparseTensorType): raise NotScalarConstantError() return data diff --git a/doc/extending/other_ops.rst b/doc/extending/other_ops.rst index 85c36be915..8965912ab1 100644 --- a/doc/extending/other_ops.rst +++ b/doc/extending/other_ops.rst @@ -44,7 +44,7 @@ usual dense tensors. In particular, in the instead of ``as_tensor_variable(x)``. Another difference is that you need to use ``SparseVariable`` and -``SparseType`` instead of ``TensorVariable`` and ``TensorType``. +``SparseTensorType`` instead of ``TensorVariable`` and ``TensorType``. Do not forget that we support only sparse matrices (so only 2 dimensions) and (like in SciPy) they do not support broadcasting operations by default @@ -55,7 +55,7 @@ you can create output variables like this: .. code-block:: python out_format = inputs[0].format # or 'csr' or 'csc' if the output format is fixed - SparseType(dtype=inputs[0].dtype, format=out_format).make_variable() + SparseTensorType(dtype=inputs[0].dtype, format=out_format).make_variable() See the sparse :class:`Aesara.sparse.basic.Cast` `Op` code for a good example of a sparse `Op` with Python code. @@ -226,7 +226,7 @@ along with pointers to the relevant documentation. primitive type. The C type associated with this Aesara type is the represented C primitive itself. -* :ref:`SparseType ` : Aesara `Type` used to represent sparse +* :ref:`SparseTensorType ` : Aesara `Type` used to represent sparse tensors. There is no equivalent C type for this Aesara `Type` but you can split a sparse variable into its parts as TensorVariables. Those can then be used as inputs to an op with C code. diff --git a/tests/compile/function/test_pfunc.py b/tests/compile/function/test_pfunc.py index 2aeb7477bd..6cb0e40e76 100644 --- a/tests/compile/function/test_pfunc.py +++ b/tests/compile/function/test_pfunc.py @@ -751,8 +751,8 @@ def test_sparse_input_aliasing_affecting_inplace_operations(self): # operations are used) and to break the elemwise composition # with some non-elemwise op (here dot) - x = sparse.SparseType("csc", dtype="float64")() - y = sparse.SparseType("csc", dtype="float64")() + x = sparse.SparseTensorType("csc", dtype="float64")() + y = sparse.SparseTensorType("csc", dtype="float64")() f = function([In(x, mutable=True), In(y, mutable=True)], (x + y) + (x + y)) # Test 1. If the same variable is given twice diff --git a/tests/sparse/test_basic.py b/tests/sparse/test_basic.py index 2a1d6dbe5c..e0f967ea09 100644 --- a/tests/sparse/test_basic.py +++ b/tests/sparse/test_basic.py @@ -38,7 +38,7 @@ Remove0, SamplingDot, SparseFromDense, - SparseType, + SparseTensorType, SquareDiagonal, StructuredDot, StructuredDotGradCSC, @@ -413,7 +413,7 @@ def test_getitem_2d(self): pass def test_getitem_scalar(self): - x = SparseType("csr", dtype=config.floatX)() + x = SparseTensorType("csr", dtype=config.floatX)() self._compile_and_check( [x], [x[2, 2]], @@ -451,7 +451,7 @@ def test_csm_grad(self): ) def test_transpose(self): - x = SparseType("csr", dtype=config.floatX)() + x = SparseTensorType("csr", dtype=config.floatX)() self._compile_and_check( [x], [x.T], @@ -460,7 +460,7 @@ def test_transpose(self): ) def test_neg(self): - x = SparseType("csr", dtype=config.floatX)() + x = SparseTensorType("csr", dtype=config.floatX)() self._compile_and_check( [x], [-x], @@ -469,8 +469,8 @@ def test_neg(self): ) def test_add_ss(self): - x = SparseType("csr", dtype=config.floatX)() - y = SparseType("csr", dtype=config.floatX)() + x = SparseTensorType("csr", dtype=config.floatX)() + y = SparseTensorType("csr", dtype=config.floatX)() self._compile_and_check( [x, y], [x + y], @@ -482,7 +482,7 @@ def test_add_ss(self): ) def test_add_sd(self): - x = SparseType("csr", dtype=config.floatX)() + x = SparseTensorType("csr", dtype=config.floatX)() y = matrix() self._compile_and_check( [x, y], @@ -495,8 +495,8 @@ def test_add_sd(self): ) def test_mul_ss(self): - x = SparseType("csr", dtype=config.floatX)() - y = SparseType("csr", dtype=config.floatX)() + x = SparseTensorType("csr", dtype=config.floatX)() + y = SparseTensorType("csr", dtype=config.floatX)() self._compile_and_check( [x, y], [x * y], @@ -508,7 +508,7 @@ def test_mul_ss(self): ) def test_mul_sd(self): - x = SparseType("csr", dtype=config.floatX)() + x = SparseTensorType("csr", dtype=config.floatX)() y = matrix() self._compile_and_check( [x, y], @@ -522,7 +522,7 @@ def test_mul_sd(self): ) def test_remove0(self): - x = SparseType("csr", dtype=config.floatX)() + x = SparseTensorType("csr", dtype=config.floatX)() self._compile_and_check( [x], [Remove0()(x)], @@ -531,8 +531,8 @@ def test_remove0(self): ) def test_dot(self): - x = SparseType("csc", dtype=config.floatX)() - y = SparseType("csc", dtype=config.floatX)() + x = SparseTensorType("csc", dtype=config.floatX)() + y = SparseTensorType("csc", dtype=config.floatX)() self._compile_and_check( [x, y], [Dot()(x, y)], @@ -545,12 +545,12 @@ def test_dot(self): def test_dot_broadcast(self): for x, y in [ - (SparseType("csr", "float32")(), vector()[:, None]), - (SparseType("csr", "float32")(), vector()[None, :]), - (SparseType("csr", "float32")(), matrix()), - (vector()[:, None], SparseType("csr", "float32")()), - (vector()[None, :], SparseType("csr", "float32")()), - (matrix(), SparseType("csr", "float32")()), + (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 = at.dot(x, y) @@ -562,8 +562,8 @@ def test_dot_broadcast(self): assert dense_out.broadcastable == sparse_out.broadcastable def test_structured_dot(self): - x = SparseType("csc", dtype=config.floatX)() - y = SparseType("csc", dtype=config.floatX)() + x = SparseTensorType("csc", dtype=config.floatX)() + y = SparseTensorType("csc", dtype=config.floatX)() self._compile_and_check( [x, y], [structured_dot(x, y)], @@ -583,8 +583,8 @@ def test_structured_dot_grad(self): ("csc", StructuredDotGradCSC), ("csr", StructuredDotGradCSR), ]: - x = SparseType(format, dtype=config.floatX)() - y = SparseType(format, dtype=config.floatX)() + x = SparseTensorType(format, dtype=config.floatX)() + y = SparseTensorType(format, dtype=config.floatX)() grads = aesara.grad(dense_from_sparse(structured_dot(x, y)).sum(), [x, y]) self._compile_and_check( [x, y], @@ -606,7 +606,7 @@ def test_structured_dot_grad(self): ) def test_dense_from_sparse(self): - x = SparseType("csr", dtype=config.floatX)() + x = SparseTensorType("csr", dtype=config.floatX)() self._compile_and_check( [x], [dense_from_sparse(x)], @@ -1130,7 +1130,7 @@ def test_csm_properties(self): for format in ("csc", "csr"): for dtype in ("float32", "float64"): - x = SparseType(format, dtype=dtype)() + x = SparseTensorType(format, dtype=dtype)() f = aesara.function([x], csm_properties(x)) spmat = sp_types[format](random_lil((4, 3), dtype, 3)) @@ -1288,7 +1288,7 @@ def test_upcast(self): for dense_dtype in typenames: for sparse_dtype in typenames: correct_dtype = aesara.scalar.upcast(sparse_dtype, dense_dtype) - a = SparseType("csc", dtype=sparse_dtype)() + a = SparseTensorType("csc", dtype=sparse_dtype)() b = matrix(dtype=dense_dtype) d = structured_dot(a, b) assert d.type.dtype == correct_dtype @@ -1375,8 +1375,8 @@ def test_dot_sparse_sparse(self): for sparse_format_a in ["csc", "csr", "bsr"]: for sparse_format_b in ["csc", "csr", "bsr"]: - a = SparseType(sparse_format_a, dtype=sparse_dtype)() - b = SparseType(sparse_format_b, dtype=sparse_dtype)() + a = SparseTensorType(sparse_format_a, dtype=sparse_dtype)() + b = SparseTensorType(sparse_format_b, dtype=sparse_dtype)() d = at.dot(a, b) f = aesara.function([a, b], Out(d, borrow=True)) for M, N, K, nnz in [ @@ -1397,7 +1397,7 @@ def test_csc_correct_output_faster_than_scipy(self): sparse_dtype = "float64" dense_dtype = "float64" - a = SparseType("csc", dtype=sparse_dtype)() + a = SparseTensorType("csc", dtype=sparse_dtype)() b = matrix(dtype=dense_dtype) d = at.dot(a, b) f = aesara.function([a, b], Out(d, borrow=True)) @@ -1445,7 +1445,7 @@ def test_csr_correct_output_faster_than_scipy(self): sparse_dtype = "float32" dense_dtype = "float32" - a = SparseType("csr", dtype=sparse_dtype)() + a = SparseTensorType("csr", dtype=sparse_dtype)() b = matrix(dtype=dense_dtype) d = at.dot(a, b) f = aesara.function([a, b], d) @@ -1567,8 +1567,8 @@ def test_sparse_sparse(self): ("csr", "csc"), ("csr", "csr"), ]: - x = sparse.SparseType(format=x_f, dtype=d1)("x") - y = sparse.SparseType(format=x_f, dtype=d2)("x") + 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 @@ -1886,7 +1886,7 @@ def test(self): def test_shape_i(): sparse_dtype = "float32" - a = SparseType("csr", dtype=sparse_dtype)() + a = SparseTensorType("csr", dtype=sparse_dtype)() f = aesara.function([a], a.shape[1]) assert f(sp.sparse.csr_matrix(random_lil((100, 10), sparse_dtype, 3))) == 10 @@ -1896,7 +1896,7 @@ def test_shape(): # does not actually create a dense tensor in the process. sparse_dtype = "float32" - a = SparseType("csr", dtype=sparse_dtype)() + a = SparseTensorType("csr", dtype=sparse_dtype)() f = aesara.function([a], a.shape) assert np.all( f(sp.sparse.csr_matrix(random_lil((100, 10), sparse_dtype, 3))) == (100, 10) @@ -1946,7 +1946,7 @@ def as_ar(a): (b.transpose(), a, False), ]: - assert SparseType.may_share_memory(a_, b_) == rep + assert SparseTensorType.may_share_memory(a_, b_) == rep def test_sparse_shared_memory(): @@ -1955,8 +1955,8 @@ def test_sparse_shared_memory(): a = random_lil((3, 4), "float32", 3).tocsr() m1 = random_lil((4, 4), "float32", 3).tocsr() m2 = random_lil((4, 4), "float32", 3).tocsr() - x = SparseType("csr", dtype="float32")() - y = SparseType("csr", dtype="float32")() + x = SparseTensorType("csr", dtype="float32")() + y = SparseTensorType("csr", dtype="float32")() sdot = sparse.structured_dot z = sdot(x * 3, m1) + sdot(y * 2, m2) @@ -1966,7 +1966,7 @@ def test_sparse_shared_memory(): def f_(x, y, m1=m1, m2=m2): return ((x * 3) * m1) + ((y * 2) * m2) - assert SparseType.may_share_memory(a, a) # This is trivial + assert SparseTensorType.may_share_memory(a, a) # This is trivial result = f(a, a) result_ = f_(a, a) assert (result_.todense() == result.todense()).all() @@ -3192,7 +3192,7 @@ def test_mul_s_v(self): for format in ("csr", "csc"): for dtype in ("float32", "float64"): - x = sparse.SparseType(format, dtype=dtype)() + x = sparse.SparseTensorType(format, dtype=dtype)() y = vector(dtype=dtype) f = aesara.function([x, y], mul_s_v(x, y)) @@ -3220,7 +3220,7 @@ def test_structured_add_s_v(self): for format in ("csr", "csc"): for dtype in ("float32", "float64"): - x = sparse.SparseType(format, dtype=dtype)() + x = sparse.SparseTensorType(format, dtype=dtype)() y = vector(dtype=dtype) f = aesara.function([x, y], structured_add_s_v(x, y)) diff --git a/tests/sparse/test_type.py b/tests/sparse/test_type.py index 28ec9a758d..eb6ce331ce 100644 --- a/tests/sparse/test_type.py +++ b/tests/sparse/test_type.py @@ -1,12 +1,12 @@ import pytest from aesara.sparse import matrix as sp_matrix -from aesara.sparse.type import SparseType +from aesara.sparse.type import SparseTensorType from aesara.tensor import dmatrix def test_clone(): - st = SparseType("csr", "float64") + st = SparseTensorType("csr", "float64") assert st == st.clone() diff --git a/tests/sparse/test_var.py b/tests/sparse/test_var.py index 04f571eb74..df2e04cbf8 100644 --- a/tests/sparse/test_var.py +++ b/tests/sparse/test_var.py @@ -7,7 +7,7 @@ import aesara import aesara.sparse as sparse import aesara.tensor as at -from aesara.sparse.type import SparseType +from aesara.sparse.type import SparseTensorType from aesara.tensor.type import DenseTensorType @@ -16,7 +16,7 @@ class TestSparseVariable: "method, exp_type, cm", [ ("__abs__", DenseTensorType, None), - ("__neg__", SparseType, ExitStack()), + ("__neg__", SparseTensorType, ExitStack()), ("__ceil__", DenseTensorType, None), ("__floor__", DenseTensorType, None), ("__trunc__", DenseTensorType, None), @@ -65,7 +65,7 @@ class TestSparseVariable: ("conj", DenseTensorType, None), ("round", DenseTensorType, None), ("trace", DenseTensorType, None), - ("zeros_like", SparseType, ExitStack()), + ("zeros_like", SparseTensorType, ExitStack()), ("ones_like", DenseTensorType, ExitStack()), ("cumsum", DenseTensorType, None), ("cumprod", DenseTensorType, None), @@ -83,7 +83,7 @@ def test_unary(self, method, exp_type, cm): if cm is None: cm = pytest.warns(UserWarning, match=".*converted to dense.*") - if exp_type == SparseType: + if exp_type == SparseTensorType: exp_res_type = csr_matrix else: exp_res_type = np.ndarray @@ -112,16 +112,16 @@ def test_unary(self, method, exp_type, cm): @pytest.mark.parametrize( "method, exp_type", [ - ("__lt__", SparseType), - ("__le__", SparseType), - ("__gt__", SparseType), - ("__ge__", SparseType), + ("__lt__", SparseTensorType), + ("__le__", SparseTensorType), + ("__gt__", SparseTensorType), + ("__ge__", SparseTensorType), ("__and__", DenseTensorType), ("__or__", DenseTensorType), ("__xor__", DenseTensorType), - ("__add__", SparseType), - ("__sub__", SparseType), - ("__mul__", SparseType), + ("__add__", SparseTensorType), + ("__sub__", SparseTensorType), + ("__mul__", SparseTensorType), ("__pow__", DenseTensorType), ("__mod__", DenseTensorType), ("__divmod__", DenseTensorType), @@ -137,7 +137,7 @@ def test_binary(self, method, exp_type): method_to_call = getattr(x, method) - if exp_type == SparseType: + if exp_type == SparseTensorType: exp_res_type = csr_matrix cm = ExitStack() else: @@ -198,7 +198,7 @@ def test_getitem(self): x = sparse.csr_from_dense(x) z = x[:, :2] - assert isinstance(z.type, SparseType) + assert isinstance(z.type, SparseTensorType) f = aesara.function([x], z) exp_res = f([[1.1, 0.0, 2.0], [-1.0, 0.0, 0.0]]) @@ -211,7 +211,7 @@ def test_dot(self): y = sparse.csr_from_dense(y) z = x.__dot__(y) - assert isinstance(z.type, SparseType) + assert isinstance(z.type, SparseTensorType) f = aesara.function([x, y], z) exp_res = f( diff --git a/tests/typed_list/test_basic.py b/tests/typed_list/test_basic.py index 52f7b81047..3192719ea8 100644 --- a/tests/typed_list/test_basic.py +++ b/tests/typed_list/test_basic.py @@ -451,7 +451,7 @@ def test_non_tensor_type(self): def test_sparse(self): sp = pytest.importorskip("scipy") mySymbolicSparseList = TypedListType( - sparse.SparseType("csr", aesara.config.floatX) + sparse.SparseTensorType("csr", aesara.config.floatX) )() mySymbolicSparse = sparse.csr_matrix() @@ -519,7 +519,7 @@ def test_non_tensor_type(self): def test_sparse(self): sp = pytest.importorskip("scipy") mySymbolicSparseList = TypedListType( - sparse.SparseType("csr", aesara.config.floatX) + sparse.SparseTensorType("csr", aesara.config.floatX) )() mySymbolicSparse = sparse.csr_matrix()