diff --git a/noxfile.py b/noxfile.py index 8a76f9f..4ff5526 100644 --- a/noxfile.py +++ b/noxfile.py @@ -1,4 +1,4 @@ -from nox import options, Session, session +from nox import Session, options, parametrize, session options.sessions = ["test", "test_numpy", "coverage", "lint"] @@ -25,5 +25,12 @@ def coverage(s: Session): @session(venv_backend="none") -def lint(s: Session) -> None: - s.run("ruff", "check", ".") +@parametrize("command", [["ruff", "check", "."], ["ruff", "format", "--check", "."]]) +def lint(s: Session, command: list[str]): + s.run(*command) + + +@session(venv_backend="none") +def format(s: Session) -> None: + s.run("ruff", "check", ".", "--select", "I", "--fix") + s.run("ruff", "format", ".") diff --git a/pyproject.toml b/pyproject.toml index 99e5cb5..d5e8b1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,11 +19,12 @@ source = [ [tool.ruff] src = ["src"] -line-length = 120 +line-length = 99 extend-exclude = ["src/taco/**"] extend-select = [ + "I", # isort "N", # pep8-naming ] diff --git a/setup.py b/setup.py index c8eced4..cd18d52 100644 --- a/setup.py +++ b/setup.py @@ -1,15 +1,15 @@ import platform import subprocess +from distutils.command.build import build from pathlib import Path -from distutils.command.build import build -from setuptools import setup, find_packages +from setuptools import find_packages, setup from wheel.bdist_wheel import bdist_wheel project_dir = Path(__file__).parent.resolve() -taco_source_dir = project_dir.joinpath('src/taco') -taco_build_dir = project_dir.joinpath('build/taco/') -taco_install_dir = project_dir.joinpath('src/tensora/taco/') +taco_source_dir = project_dir.joinpath("src/taco") +taco_build_dir = project_dir.joinpath("build/taco/") +taco_install_dir = project_dir.joinpath("src/tensora/taco/") class TensoraBuild(build): @@ -17,20 +17,25 @@ def run(self): # Build taco os = platform.system() if os == "Linux": - install_path = r'-DCMAKE_INSTALL_RPATH=\$ORIGIN/../lib' + install_path = r"-DCMAKE_INSTALL_RPATH=\$ORIGIN/../lib" elif os == "Darwin": - install_path = r'-DCMAKE_INSTALL_RPATH=@loader_path/../lib' + install_path = r"-DCMAKE_INSTALL_RPATH=@loader_path/../lib" else: - raise NotImplementedError(f'Tensora cannot be installed on {os}') + raise NotImplementedError(f"Tensora cannot be installed on {os}") taco_build_dir.mkdir(parents=True, exist_ok=True) - subprocess.check_call(['cmake', str(taco_source_dir), - '-DCMAKE_BUILD_TYPE=Release', - f'-DCMAKE_INSTALL_PREFIX={taco_install_dir}', - install_path], - cwd=taco_build_dir) - subprocess.check_call(['make', '-j8'], cwd=taco_build_dir) - subprocess.check_call(['make', 'install'], cwd=taco_build_dir) + subprocess.check_call( + [ + "cmake", + str(taco_source_dir), + "-DCMAKE_BUILD_TYPE=Release", + f"-DCMAKE_INSTALL_PREFIX={taco_install_dir}", + install_path, + ], + cwd=taco_build_dir, + ) + subprocess.check_call(["make", "-j8"], cwd=taco_build_dir) + subprocess.check_call(["make", "install"], cwd=taco_build_dir) super().run() @@ -44,44 +49,38 @@ def finalize_options(self): setup( - name='tensora', - version='0.0.8', - - description='Library for dense and sparse tensors built on the tensor algebra compiler.', - long_description=Path('README.md').read_text(encoding='utf-8'), - long_description_content_type='text/markdown', - keywords='tensor sparse matrix array', - - author='David Hagen', - author_email='david@drhagen.com', - url='https://github.com/drhagen/tensora', - license='MIT', - - package_dir={'': 'src'}, - packages=find_packages('src'), - package_data={'tensora': ['taco/bin/taco', 'taco/lib/libtaco.*']}, - - install_requires=Path('requirements.txt').read_text(encoding='utf-8').splitlines(), - extras_require={'numpy': ['numpy'], 'scipy': ['scipy']}, - + name="tensora", + version="0.0.8", + description="Library for dense and sparse tensors built on the tensor algebra compiler.", + long_description=Path("README.md").read_text(encoding="utf-8"), + long_description_content_type="text/markdown", + keywords="tensor sparse matrix array", + author="David Hagen", + author_email="david@drhagen.com", + url="https://github.com/drhagen/tensora", + license="MIT", + package_dir={"": "src"}, + packages=find_packages("src"), + package_data={"tensora": ["taco/bin/taco", "taco/lib/libtaco.*"]}, + install_requires=Path("requirements.txt").read_text(encoding="utf-8").splitlines(), + extras_require={"numpy": ["numpy"], "scipy": ["scipy"]}, classifiers=[ - 'Development Status :: 3 - Alpha', - 'Intended Audience :: Science/Research', - 'Intended Audience :: Developers', - 'Topic :: Software Development :: Libraries', - 'License :: OSI Approved :: MIT License', - 'Operating System :: POSIX :: Linux', - 'Operating System :: MacOS', - 'Programming Language :: Python', - 'Programming Language :: Python :: 3', - 'Programming Language :: Python :: 3 :: Only', - 'Programming Language :: Python :: 3.10', - 'Programming Language :: Python :: 3.11', + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "Intended Audience :: Developers", + "Topic :: Software Development :: Libraries", + "License :: OSI Approved :: MIT License", + "Operating System :: POSIX :: Linux", + "Operating System :: MacOS", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", ], - cmdclass={ - 'build': TensoraBuild, - 'bdist_wheel': TensoraBdistWheel, + "build": TensoraBuild, + "bdist_wheel": TensoraBdistWheel, }, zip_safe=False, ) diff --git a/src/tensora/__init__.py b/src/tensora/__init__.py index 1bc4f58..42d1e5a 100644 --- a/src/tensora/__init__.py +++ b/src/tensora/__init__.py @@ -1,3 +1,3 @@ -from .format import Mode, Format # noqa: F401 +from .format import Format, Mode # noqa: F401 +from .function import evaluate, tensor_method # noqa: F401 from .tensor import Tensor # noqa: F401 -from .function import tensor_method, evaluate # noqa: F401 diff --git a/src/tensora/compile.py b/src/tensora/compile.py index f9a3762..8651c4a 100644 --- a/src/tensora/compile.py +++ b/src/tensora/compile.py @@ -1,20 +1,25 @@ -__all__ = ['taco_kernel', 'allocate_taco_structure', 'taco_structure_to_cffi', 'take_ownership_of_arrays', - 'tensor_cdefs', 'take_ownership_of_tensor_members'] +__all__ = [ + "taco_kernel", + "allocate_taco_structure", + "taco_structure_to_cffi", + "take_ownership_of_arrays", + "tensor_cdefs", + "take_ownership_of_tensor_members", +] import re import subprocess import tempfile +import threading from pathlib import Path -from typing import List, Tuple, FrozenSet, Any +from typing import Any, FrozenSet, List, Tuple from weakref import WeakKeyDictionary from cffi import FFI -import threading - lock = threading.Lock() -taco_binary = Path(__file__).parent.joinpath('taco/bin/taco') +taco_binary = Path(__file__).parent.joinpath("taco/bin/taco") global_weakkeydict = WeakKeyDictionary() @@ -81,7 +86,7 @@ tensor_cdefs.cdef(taco_type_header) # This library only has definitions, in order to call `dlopen`, `set_source` must be called with empty `source` first -tensor_cdefs.set_source('_main', '') +tensor_cdefs.set_source("_main", "") tensor_lib = tensor_cdefs.dlopen(None) @@ -111,9 +116,12 @@ def taco_kernel(expression: str, formats: FrozenSet[Tuple[str, str]]) -> Tuple[L cffi pointers to taco_tensor_t instances in order specified by the list of variable names. """ # Call taco to write the kernels to standard out - result = subprocess.run([taco_binary, expression, '-print-evaluate', '-print-nocolor'] - + [f'-f={name}:{format}' for name, format in formats], - capture_output=True, text=True) + result = subprocess.run( + [taco_binary, expression, "-print-evaluate", "-print-nocolor"] + + [f"-f={name}:{format}" for name, format in formats], + capture_output=True, + text=True, + ) if result.returncode != 0: raise RuntimeError(result.stderr) @@ -123,17 +131,20 @@ def taco_kernel(expression: str, formats: FrozenSet[Tuple[str, str]]) -> Tuple[L # Determine signature # 1) Find function by name and capture its parameter list # 2) Find each parameter by `*` and capture its name - signature_match = re.search(r'int evaluate\(([^)]*)\)', source) + signature_match = re.search(r"int evaluate\(([^)]*)\)", source) signature = signature_match.group(0) - parameter_list_matches = re.finditer(r'\*([^,]*)', signature_match.group(1)) + parameter_list_matches = re.finditer(r"\*([^,]*)", signature_match.group(1)) parameter_names = [match.group(1) for match in parameter_list_matches] # Use cffi to compile the kernels ffibuilder = FFI() ffibuilder.include(tensor_cdefs) - ffibuilder.cdef(signature + ';') - ffibuilder.set_source('taco_kernel', taco_define_header + taco_type_header + source, - extra_compile_args=['-Wno-unused-variable', '-Wno-unknown-pragmas']) + ffibuilder.cdef(signature + ";") + ffibuilder.set_source( + "taco_kernel", + taco_define_header + taco_type_header + source, + extra_compile_args=["-Wno-unused-variable", "-Wno-unknown-pragmas"], + ) with tempfile.TemporaryDirectory() as temp_dir: # Lock because FFI.compile is not thread safe: https://foss.heptapod.net/pypy/cffi/-/issues/490 @@ -152,7 +163,9 @@ def taco_kernel(expression: str, formats: FrozenSet[Tuple[str, str]]) -> Tuple[L return parameter_names, lib -def allocate_taco_structure(mode_types: Tuple[int, ...], dimensions: Tuple[int, ...], mode_ordering: Tuple[int, ...]): +def allocate_taco_structure( + mode_types: Tuple[int, ...], dimensions: Tuple[int, ...], mode_ordering: Tuple[int, ...] +): """Allocate all parts of a taco tensor except growable arrays. All int32_t[] tensor.indices[*][*] and double[] tensor.vals are NULL pointers. All other properties are immutable. @@ -171,40 +184,44 @@ def allocate_taco_structure(mode_types: Tuple[int, ...], dimensions: Tuple[int, """ # Validate inputs if not (len(mode_types) == len(dimensions) == len(mode_ordering)): - raise ValueError(f'Must all be the same length: mode_types = {mode_types}, dimensions = {dimensions}, ' - f'mode_ordering = {mode_ordering}') + raise ValueError( + f"Must all be the same length: mode_types = {mode_types}, dimensions = {dimensions}, " + f"mode_ordering = {mode_ordering}" + ) for mode_type in mode_types: if mode_type not in (0, 1): - raise ValueError(f'mode_types must only contain elements 0 or 1: {mode_types}') + raise ValueError(f"mode_types must only contain elements 0 or 1: {mode_types}") for dimension in dimensions: if dimension < 0: # cffi will reject integers too big to fit into an int32_t - raise ValueError(f'All values in dimensions must be positive: {dimensions}') + raise ValueError(f"All values in dimensions must be positive: {dimensions}") if set(mode_ordering) != set(range(len(mode_types))): - raise ValueError(f'mode_ordering must contain each number in the set {{0, 1, ..., order - 1}} exactly once: ' - f'{mode_ordering}') + raise ValueError( + f"mode_ordering must contain each number in the set {{0, 1, ..., order - 1}} exactly once: " + f"{mode_ordering}" + ) # This structure mimics the taco structure and holds the objects owning the memory of the arrays, ensuring that the # pointers stay valid as long as the cffi taco structure has not been garbage collected. memory_holder = {} - cffi_tensor = tensor_cdefs.new('taco_tensor_t*') + cffi_tensor = tensor_cdefs.new("taco_tensor_t*") cffi_tensor.order = len(mode_types) - cffi_dimensions = tensor_cdefs.new('int32_t[]', dimensions) - memory_holder['dimensions'] = cffi_dimensions + cffi_dimensions = tensor_cdefs.new("int32_t[]", dimensions) + memory_holder["dimensions"] = cffi_dimensions cffi_tensor.dimensions = cffi_dimensions - cffi_mode_ordering = tensor_cdefs.new('int32_t[]', mode_ordering) - memory_holder['mode_ordering'] = cffi_mode_ordering + cffi_mode_ordering = tensor_cdefs.new("int32_t[]", mode_ordering) + memory_holder["mode_ordering"] = cffi_mode_ordering cffi_tensor.mode_ordering = cffi_mode_ordering - cffi_mode_types = tensor_cdefs.new('taco_mode_t[]', mode_types) - memory_holder['mode_types'] = cffi_mode_types + cffi_mode_types = tensor_cdefs.new("taco_mode_t[]", mode_types) + memory_holder["mode_types"] = cffi_mode_types cffi_tensor.mode_types = cffi_mode_types converted_levels = [] @@ -215,15 +232,15 @@ def allocate_taco_structure(mode_types: Tuple[int, ...], dimensions: Tuple[int, converted_arrays = [] elif mode == tensor_lib.taco_mode_sparse: converted_arrays = [tensor_cdefs.NULL, tensor_cdefs.NULL] - cffi_level = tensor_cdefs.new('int32_t*[]', converted_arrays) + cffi_level = tensor_cdefs.new("int32_t*[]", converted_arrays) memory_holder_levels.append(cffi_level) memory_holder_levels_arrays.append(converted_arrays) converted_levels.append(cffi_level) - cffi_levels = tensor_cdefs.new('int32_t**[]', converted_levels) - memory_holder['indices'] = cffi_levels - memory_holder['*indices'] = memory_holder_levels - memory_holder['**indices'] = memory_holder_levels_arrays - cffi_tensor.indices = tensor_cdefs.cast('uint8_t***', cffi_levels) + cffi_levels = tensor_cdefs.new("int32_t**[]", converted_levels) + memory_holder["indices"] = cffi_levels + memory_holder["*indices"] = memory_holder_levels + memory_holder["**indices"] = memory_holder_levels_arrays + cffi_tensor.indices = tensor_cdefs.cast("uint8_t***", cffi_levels) cffi_tensor.vals = tensor_cdefs.NULL @@ -234,8 +251,14 @@ def allocate_taco_structure(mode_types: Tuple[int, ...], dimensions: Tuple[int, return cffi_tensor -def taco_structure_to_cffi(indices: List[List[List[int]]], vals: List[float], *, - mode_types: Tuple[int, ...], dimensions: Tuple[int, ...], mode_ordering: Tuple[int, ...]): +def taco_structure_to_cffi( + indices: List[List[List[int]]], + vals: List[float], + *, + mode_types: Tuple[int, ...], + dimensions: Tuple[int, ...], + mode_ordering: Tuple[int, ...], +): """Build a cffi taco tensor from Python data. This takes Python data with a one-to-one mapping to taco tensor attributes and builds a cffi taco tensor from it. @@ -255,58 +278,76 @@ def taco_structure_to_cffi(indices: List[List[List[int]]], vals: List[float], *, # Validate inputs if len(indices) != len(mode_types): - raise ValueError(f'Length of indices ({len(indices)}) must be equal to the length of mode_types, dimensions, ' - f'and mode_ordering ({len(mode_types)})') + raise ValueError( + f"Length of indices ({len(indices)}) must be equal to the length of mode_types, dimensions, " + f"and mode_ordering ({len(mode_types)})" + ) nnz = 1 for i_level in range(cffi_tensor.order): if mode_types[i_level] == 0: if len(indices[i_level]) != 0: - raise ValueError(f'Level {i_level} is a dense mode and therefore expects indices[{i_level}] to be ' - f'empty: {indices[i_level]}') + raise ValueError( + f"Level {i_level} is a dense mode and therefore expects indices[{i_level}] to be " + f"empty: {indices[i_level]}" + ) nnz *= dimensions[mode_ordering[i_level]] elif mode_types[i_level] == 1: if len(indices[i_level]) != 2: - raise ValueError(f'Level {i_level} is a compressed mode and therefore expects indices[{i_level}] to be ' - f'length 2 not length {len(indices[i_level])}: {indices[i_level]}') + raise ValueError( + f"Level {i_level} is a compressed mode and therefore expects indices[{i_level}] to be " + f"length 2 not length {len(indices[i_level])}: {indices[i_level]}" + ) pos = indices[i_level][0] crd = indices[i_level][1] if len(pos) != nnz + 1: - raise ValueError(f'The pos array of level {i_level} (indices[{i_level}][0]) must have length {nnz}, ' - f'the number of explicit indexes so far, not length {len(pos)}: {pos}') + raise ValueError( + f"The pos array of level {i_level} (indices[{i_level}][0]) must have length {nnz}, " + f"the number of explicit indexes so far, not length {len(pos)}: {pos}" + ) if pos[0] != 0: - raise ValueError(f'The first element of the pos array of level {i_level} must be 0: {pos}') + raise ValueError( + f"The first element of the pos array of level {i_level} must be 0: {pos}" + ) if not weakly_increasing(pos): - raise ValueError(f'The pos array of level {i_level} (indices[{i_level}][0]) must be weakly ' - f'monotonically increasing: {pos}') + raise ValueError( + f"The pos array of level {i_level} (indices[{i_level}][0]) must be weakly " + f"monotonically increasing: {pos}" + ) if len(crd) != pos[-1]: - raise ValueError(f'The crd array of level {i_level} (indices[{i_level}][1]) must have length ' - f"{pos[-1]}, the last element of this level's pos array, not length {len(crd)}: {crd}") + raise ValueError( + f"The crd array of level {i_level} (indices[{i_level}][1]) must have length " + f"{pos[-1]}, the last element of this level's pos array, not length {len(crd)}: {crd}" + ) if not all(0 <= x < dimensions[mode_ordering[i_level]] for x in crd): - raise ValueError(f'All values in the crd array of level {i_level} (indices[{i_level}][1]) must be ' - f'nonnegative and less than the size of this dimension: {crd}') + raise ValueError( + f"All values in the crd array of level {i_level} (indices[{i_level}][1]) must be " + f"nonnegative and less than the size of this dimension: {crd}" + ) nnz = len(crd) if len(vals) != nnz: - raise ValueError(f'Length of vals must be equal to the number of indexes implicitly defined by indices {nnz} ' - f'not {len(vals)}: {vals}') + raise ValueError( + f"Length of vals must be equal to the number of indexes implicitly defined by indices {nnz} " + f"not {len(vals)}: {vals}" + ) # Get the partial constructed memory holder stored by allocate_taco_structure memory_holder = global_weakkeydict[cffi_tensor] - cffi_indices = tensor_cdefs.cast('int32_t***', cffi_tensor.indices) + cffi_indices = tensor_cdefs.cast("int32_t***", cffi_tensor.indices) for i_level, (mode, level) in enumerate(zip(mode_types, indices)): if mode == tensor_lib.taco_mode_dense: pass elif mode == tensor_lib.taco_mode_sparse: for i_array, array in enumerate(level): - cffi_array = tensor_cdefs.new('int32_t[]', array) - memory_holder['**indices'][i_level][i_array] = cffi_array + cffi_array = tensor_cdefs.new("int32_t[]", array) + memory_holder["**indices"][i_level][i_array] = cffi_array cffi_indices[i_level][i_array] = cffi_array - cffi_vals = tensor_cdefs.new('double[]', vals) - memory_holder['vals'] = cffi_vals - cffi_tensor.vals = tensor_cdefs.cast('uint8_t*', cffi_vals) + cffi_vals = tensor_cdefs.new("double[]", vals) + memory_holder["vals"] = cffi_vals + cffi_tensor.vals = tensor_cdefs.cast("uint8_t*", cffi_vals) cffi_tensor.vals_size = len(vals) @@ -333,15 +374,19 @@ def take_ownership_of_arrays(cffi_tensor) -> None: modes = cffi_tensor.mode_types[0:order] - cffi_levels = tensor_cdefs.cast('int32_t***', cffi_tensor.indices) + cffi_levels = tensor_cdefs.cast("int32_t***", cffi_tensor.indices) for i_dimension, mode in enumerate(modes): if mode == tensor_lib.taco_mode_dense: pass if mode == tensor_lib.taco_mode_sparse: - memory_holder['**indices'][i_dimension][0] = tensor_cdefs.gc(cffi_levels[i_dimension][0], tensor_lib.free) - memory_holder['**indices'][i_dimension][1] = tensor_cdefs.gc(cffi_levels[i_dimension][1], tensor_lib.free) + memory_holder["**indices"][i_dimension][0] = tensor_cdefs.gc( + cffi_levels[i_dimension][0], tensor_lib.free + ) + memory_holder["**indices"][i_dimension][1] = tensor_cdefs.gc( + cffi_levels[i_dimension][1], tensor_lib.free + ) - memory_holder['vals'] = tensor_cdefs.gc(cffi_tensor.vals, tensor_lib.free) + memory_holder["vals"] = tensor_cdefs.gc(cffi_tensor.vals, tensor_lib.free) def take_ownership_of_tensor_members(cffi_tensor) -> None: @@ -363,24 +408,26 @@ def take_ownership_of_tensor_members(cffi_tensor) -> None: # First, take ownership of everything that is owned after taco_structure_to_cffi order = cffi_tensor.order - memory_holder['dimensions'] = tensor_cdefs.gc(cffi_tensor.dimensions, tensor_lib.free) + memory_holder["dimensions"] = tensor_cdefs.gc(cffi_tensor.dimensions, tensor_lib.free) - memory_holder['mode_ordering'] = tensor_cdefs.gc(cffi_tensor.mode_ordering, tensor_lib.free) + memory_holder["mode_ordering"] = tensor_cdefs.gc(cffi_tensor.mode_ordering, tensor_lib.free) - memory_holder['mode_types'] = tensor_cdefs.gc(cffi_tensor.mode_types, tensor_lib.free) + memory_holder["mode_types"] = tensor_cdefs.gc(cffi_tensor.mode_types, tensor_lib.free) memory_holder_levels = [] memory_holder_levels_arrays = [] for i_dimension, mode in enumerate(cffi_tensor.mode_types[0:order]): - memory_holder_levels.append(tensor_cdefs.gc(cffi_tensor.indices[i_dimension], tensor_lib.free)) + memory_holder_levels.append( + tensor_cdefs.gc(cffi_tensor.indices[i_dimension], tensor_lib.free) + ) if mode == tensor_lib.taco_mode_dense: memory_holder_levels_arrays.append([]) elif mode == tensor_lib.taco_mode_sparse: # It does not matter what the values are here. They will be overwritten in take_ownership_of_arrays. memory_holder_levels_arrays.append([tensor_cdefs.NULL, tensor_cdefs.NULL]) - memory_holder['indices'] = tensor_cdefs.gc(cffi_tensor.indices, tensor_lib.free) - memory_holder['*indices'] = memory_holder_levels - memory_holder['**indices'] = memory_holder_levels_arrays + memory_holder["indices"] = tensor_cdefs.gc(cffi_tensor.indices, tensor_lib.free) + memory_holder["*indices"] = memory_holder_levels + memory_holder["**indices"] = memory_holder_levels_arrays global_weakkeydict[cffi_tensor] = memory_holder @@ -401,7 +448,7 @@ def take_ownership_of_tensor(cffi_tensor) -> None: Args: cffi_tensor: A cffi taco_tensor_t*. """ - global_weakkeydict[cffi_tensor] = {'tensor': tensor_cdefs.gc(cffi_tensor, tensor_lib.free)} + global_weakkeydict[cffi_tensor] = {"tensor": tensor_cdefs.gc(cffi_tensor, tensor_lib.free)} take_ownership_of_tensor_members(cffi_tensor) diff --git a/src/tensora/expression/ast.py b/src/tensora/expression/ast.py index a676716..da36036 100644 --- a/src/tensora/expression/ast.py +++ b/src/tensora/expression/ast.py @@ -1,9 +1,21 @@ -__all__ = ['Node', 'Expression', 'Literal', 'Integer', 'Float', 'Variable', 'Scalar', 'Tensor', 'Add', 'Subtract', - 'Multiply', 'Assignment'] +__all__ = [ + "Node", + "Expression", + "Literal", + "Integer", + "Float", + "Variable", + "Scalar", + "Tensor", + "Add", + "Subtract", + "Multiply", + "Assignment", +] from abc import abstractmethod from dataclasses import dataclass -from typing import List, Dict, Tuple, Set +from typing import Dict, List, Set, Tuple class Node: @@ -100,14 +112,16 @@ class Tensor(Variable): indexes: List[str] def deparse(self): - return self.name + '(' + ','.join(self.indexes) + ')' + return self.name + "(" + ",".join(self.indexes) + ")" def merge_index_participants(left: Expression, right: Expression): left_indexes = left.index_participants() right_indexes = right.index_participants() - return {index_name: left_indexes.get(index_name, set()) | right_indexes.get(index_name, set()) - for index_name in {*left_indexes.keys(), *right_indexes.keys()}} + return { + index_name: left_indexes.get(index_name, set()) | right_indexes.get(index_name, set()) + for index_name in {*left_indexes.keys(), *right_indexes.keys()} + } @dataclass(frozen=True) @@ -116,7 +130,7 @@ class Add(Expression): right: Expression def deparse(self): - return self.left.deparse() + ' + ' + self.right.deparse() + return self.left.deparse() + " + " + self.right.deparse() def variable_orders(self) -> Dict[str, int]: return {**self.left.variable_orders(), **self.right.variable_orders()} @@ -131,7 +145,7 @@ class Subtract(Expression): right: Expression def deparse(self): - return self.left.deparse() + ' - ' + self.right.deparse() + return self.left.deparse() + " - " + self.right.deparse() def variable_orders(self) -> Dict[str, int]: return {**self.left.variable_orders(), **self.right.variable_orders()} @@ -147,16 +161,16 @@ class Multiply(Expression): def deparse(self): if isinstance(self.left, (Add, Subtract)): - left_string = '(' + self.left.deparse() + ')' + left_string = "(" + self.left.deparse() + ")" else: left_string = self.left.deparse() if isinstance(self.right, (Add, Subtract)): - right_string = '(' + self.right.deparse() + ')' + right_string = "(" + self.right.deparse() + ")" else: right_string = self.right.deparse() - return left_string + ' * ' + right_string + return left_string + " * " + right_string def variable_orders(self) -> Dict[str, int]: return {**self.left.variable_orders(), **self.right.variable_orders()} @@ -171,7 +185,7 @@ class Assignment(Node): expression: Expression def deparse(self) -> str: - return self.target.deparse() + ' = ' + self.expression.deparse() + return self.target.deparse() + " = " + self.expression.deparse() def variable_orders(self) -> Dict[str, int]: return {**self.target.variable_orders(), **self.expression.variable_orders()} diff --git a/src/tensora/expression/parser.py b/src/tensora/expression/parser.py index 48e1ef8..a4b72c2 100644 --- a/src/tensora/expression/parser.py +++ b/src/tensora/expression/parser.py @@ -1,45 +1,45 @@ -__all__ = ['parse_assignment'] +__all__ = ["parse_assignment"] from functools import reduce -from parsita import ParserContext, lit, reg, rep, rep1sep, Result +from parsita import ParserContext, Result, lit, reg, rep, rep1sep from parsita.util import splat -from .ast import Assignment, Add, Subtract, Multiply, Tensor, Scalar, Integer, Float +from .ast import Add, Assignment, Float, Integer, Multiply, Scalar, Subtract, Tensor def make_expression(first, rest): value = first for op, term in rest: - if op == '+': + if op == "+": value = Add(value, term) else: value = Subtract(value, term) return value -class TensorExpressionParsers(ParserContext, whitespace=r'[ ]*'): - name = reg(r'[A-Za-z][A-Za-z0-9]*') +class TensorExpressionParsers(ParserContext, whitespace=r"[ ]*"): + name = reg(r"[A-Za-z][A-Za-z0-9]*") # taco does not support negatives or exponents - floating_point = reg(r'[0-9]+\.[0-9]+') > (lambda x: Float(float(x))) - integer = reg(r'[0-9]+') > (lambda x: Integer(int(x))) + floating_point = reg(r"[0-9]+\.[0-9]+") > (lambda x: Float(float(x))) + integer = reg(r"[0-9]+") > (lambda x: Integer(int(x))) number = floating_point | integer # taco requires at least one index; scalar tensors are not parsed as `a()` # taco also allows for `y_{i}` and `y_i` to mean `y(i)`, but that is not supported here - tensor = name & '(' >> rep1sep(name, ',') << ')' > splat(Tensor) + tensor = name & "(" >> rep1sep(name, ",") << ")" > splat(Tensor) scalar = name > Scalar variable = tensor | scalar - parentheses = '(' >> expression << ')' # noqa: F821 + parentheses = "(" >> expression << ")" # noqa: F821 factor = variable | number | parentheses - term = rep1sep(factor, '*') > (lambda x: reduce(Multiply, x)) - expression = term & rep(lit('+', '-') & term) > splat(make_expression) + term = rep1sep(factor, "*") > (lambda x: reduce(Multiply, x)) + expression = term & rep(lit("+", "-") & term) > splat(make_expression) - simple_assignment = variable & '=' >> expression > splat(Assignment) - add_assignment = variable & '+=' >> expression > splat(lambda v, e: Assignment(v, Add(v, e))) + simple_assignment = variable & "=" >> expression > splat(Assignment) + add_assignment = variable & "+=" >> expression > splat(lambda v, e: Assignment(v, Add(v, e))) assignment = simple_assignment | add_assignment diff --git a/src/tensora/format/__init__.py b/src/tensora/format/__init__.py index b6ad69d..40e6417 100644 --- a/src/tensora/format/__init__.py +++ b/src/tensora/format/__init__.py @@ -1,2 +1,2 @@ -from .format import Mode, Format # noqa: F401 +from .format import Format, Mode # noqa: F401 from .parser import parse_format # noqa: F401 diff --git a/src/tensora/format/format.py b/src/tensora/format/format.py index 95f8e04..885f742 100644 --- a/src/tensora/format/format.py +++ b/src/tensora/format/format.py @@ -1,4 +1,4 @@ -__all__ = ['Mode', 'Format'] +__all__ = ["Mode", "Format"] from dataclasses import dataclass from enum import Enum @@ -7,19 +7,19 @@ class Mode(Enum): # Manually map these to the entries in .taco_compile.taco_type_header.taco_mode_t - dense = (0, 'd') - compressed = (1, 's') + dense = (0, "d") + compressed = (1, "s") - def __init__(self, c_int: int, character: 'str'): + def __init__(self, c_int: int, character: "str"): self.c_int = c_int self.character = character @staticmethod - def from_c_int(value: int) -> 'Mode': + def from_c_int(value: int) -> "Mode": for member in Mode: if member.value[0] == value: return member - raise ValueError(f'No member of DimensionalMode has the integer value {value}') + raise ValueError(f"No member of DimensionalMode has the integer value {value}") @dataclass(frozen=True) @@ -29,8 +29,10 @@ class Format: def __post_init__(self): if len(self.modes) != len(self.ordering): - raise ValueError(f'Length of modes ({len(self.modes)}) must be equal to length of ordering ' - f'({len(self.ordering)})') + raise ValueError( + f"Length of modes ({len(self.modes)}) must be equal to length of ordering " + f"({len(self.ordering)})" + ) @property def order(self): @@ -38,6 +40,8 @@ def order(self): def deparse(self): if self.ordering == tuple(range(self.order)): - return ''.join(mode.character for mode in self.modes) + return "".join(mode.character for mode in self.modes) else: - return ''.join(mode.character + str(ordering) for mode, ordering in zip(self.modes, self.ordering)) + return "".join( + mode.character + str(ordering) for mode, ordering in zip(self.modes, self.ordering) + ) diff --git a/src/tensora/format/parser.py b/src/tensora/format/parser.py index 21a8350..b6bab65 100644 --- a/src/tensora/format/parser.py +++ b/src/tensora/format/parser.py @@ -1,9 +1,20 @@ -__all__ = ['parse_format'] - -from parsita import ParserContext, reg, lit, rep, eof, Result, Success, Failure, ParseError, StringReader +__all__ = ["parse_format"] + +from parsita import ( + Failure, + ParseError, + ParserContext, + Result, + StringReader, + Success, + eof, + lit, + reg, + rep, +) from parsita.util import constant -from .format import Mode, Format +from .format import Format, Mode def make_format_with_orderings(dims): @@ -16,13 +27,15 @@ def make_format_with_orderings(dims): class FormatTextParsers(ParserContext): - integer = reg(r'[0-9]+') > int - dense = lit('d') > constant(Mode.dense) - compressed = lit('s') > constant(Mode.compressed) + integer = reg(r"[0-9]+") > int + dense = lit("d") > constant(Mode.dense) + compressed = lit("s") > constant(Mode.compressed) mode = dense | compressed # Use eof to ensure each parser goes to end - format_without_orderings = rep(mode) << eof > (lambda modes: Format(tuple(modes), tuple(range(len(modes))))) + format_without_orderings = rep(mode) << eof > ( + lambda modes: Format(tuple(modes), tuple(range(len(modes)))) + ) format_with_orderings = rep(mode & integer) << eof > make_format_with_orderings format = format_without_orderings | format_with_orderings @@ -35,9 +48,11 @@ def parse_format(format: str) -> Result[Format]: elif isinstance(parse_result, Success): parse_value = parse_result.unwrap() if set(range(parse_value.order)) != set(parse_value.ordering): - return Failure(ParseError( - StringReader(format), - f'format ordering as some order of the set {set(range(parse_value.order))}' - )) + return Failure( + ParseError( + StringReader(format), + f"format ordering as some order of the set {set(range(parse_value.order))}", + ) + ) else: return parse_result diff --git a/src/tensora/function.py b/src/tensora/function.py index 773bae9..c1098d0 100644 --- a/src/tensora/function.py +++ b/src/tensora/function.py @@ -1,44 +1,51 @@ -__all__ = ['tensor_method', 'evaluate', 'PureTensorMethod'] +__all__ = ["tensor_method", "evaluate", "PureTensorMethod"] from functools import lru_cache -from inspect import Signature, Parameter -from typing import Tuple, Dict +from inspect import Parameter, Signature +from typing import Dict, Tuple +from .compile import allocate_taco_structure, taco_kernel, take_ownership_of_arrays from .expression import Assignment from .format import Format, parse_format -from .compile import taco_kernel, allocate_taco_structure, take_ownership_of_arrays from .tensor import Tensor class PureTensorMethod: - """A function taking specific tensor arguments. - """ + """A function taking specific tensor arguments.""" - def __init__(self, assignment: Assignment, input_formats: Dict[str, Format], output_format: Format): + def __init__( + self, assignment: Assignment, input_formats: Dict[str, Format], output_format: Format + ): if assignment.is_mutating(): - raise ValueError(f'{assignment} mutates its target and is so is not a pure function') + raise ValueError(f"{assignment} mutates its target and is so is not a pure function") variable_orders = assignment.expression.variable_orders() # Ensure that all parameters are defined for variable_name in variable_orders.keys(): if variable_name not in input_formats: - raise ValueError(f'Variable {variable_name} in {assignment} not listed in parameters') + raise ValueError( + f"Variable {variable_name} in {assignment} not listed in parameters" + ) # Ensure that no extraneous parameters are defined for parameter_name in input_formats.keys(): if parameter_name not in variable_orders: - raise ValueError(f'Parameter {parameter_name} not in {assignment} variables') + raise ValueError(f"Parameter {parameter_name} not in {assignment} variables") # Verify that parameters have the correct order for parameter_name, format in input_formats.items(): if format.order != variable_orders[parameter_name]: - raise ValueError(f'Parameter {parameter_name} has order {format.order}, but this variable in the ' - f'assignment has order {variable_orders[parameter_name]}') + raise ValueError( + f"Parameter {parameter_name} has order {format.order}, but this variable in the " + f"assignment has order {variable_orders[parameter_name]}" + ) if output_format.order != assignment.target.order: - raise ValueError(f'Output parameter has order {output_format.order}, but the output variable in the ' - f'assignment has order {assignment.target.order}') + raise ValueError( + f"Output parameter has order {output_format.order}, but the output variable in the " + f"assignment has order {assignment.target.order}" + ) # Store validated attributes self.assignment = assignment @@ -46,14 +53,20 @@ def __init__(self, assignment: Assignment, input_formats: Dict[str, Format], out self.output_format = output_format # Create Python signature of the function - self.signature = Signature([Parameter(parameter_name, Parameter.POSITIONAL_OR_KEYWORD) - for parameter_name in input_formats.keys()]) + self.signature = Signature( + [ + Parameter(parameter_name, Parameter.POSITIONAL_OR_KEYWORD) + for parameter_name in input_formats.keys() + ] + ) # Compile taco function all_formats = {self.assignment.target.name: output_format, **input_formats} - format_strings = frozenset((parameter_name, format_to_taco_format(format)) - for parameter_name, format in all_formats.items() - if format.order != 0) # Taco does not like formats for scalars + format_strings = frozenset( + (parameter_name, format_to_taco_format(format)) + for parameter_name, format in all_formats.items() + if format.order != 0 + ) # Taco does not like formats for scalars self.parameter_order, self.cffi_lib = taco_kernel(assignment.deparse(), format_strings) def __call__(self, *args, **kwargs): @@ -61,17 +74,24 @@ def __call__(self, *args, **kwargs): bound_arguments = self.signature.bind(*args, **kwargs).arguments # Validate tensor arguments - for name, argument, format in zip(bound_arguments.keys(), bound_arguments.values(), - self.input_formats.values()): + for name, argument, format in zip( + bound_arguments.keys(), bound_arguments.values(), self.input_formats.values() + ): if argument.order != format.order: - raise ValueError(f'Argument {name} must have order {format.order} not {argument.order}') + raise ValueError( + f"Argument {name} must have order {format.order} not {argument.order}" + ) if tuple(argument.modes) != tuple(format.modes): - raise ValueError(f'Argument {name} must have modes ' - f'{tuple(mode.name for mode in format.modes)} not ' - f'{tuple(mode.name for mode in argument.modes)}') + raise ValueError( + f"Argument {name} must have modes " + f"{tuple(mode.name for mode in format.modes)} not " + f"{tuple(mode.name for mode in argument.modes)}" + ) if tuple(argument.mode_ordering) != tuple(format.ordering): - raise ValueError(f'Argument {name} must have mode ordering ' - f'{format.ordering} not {argument.mode_ordering}') + raise ValueError( + f"Argument {name} must have mode ordering " + f"{format.ordering} not {argument.mode_ordering}" + ) # Validate dimensions index_participants = self.assignment.expression.index_participants() @@ -79,24 +99,33 @@ def __call__(self, *args, **kwargs): for index, participants in index_participants.items(): # Extract the size of dimension referenced by this index on each tensor that uses it; record the variable # name and dimension for a better error - actual_sizes = [(variable, dimension, bound_arguments[variable].dimensions[dimension]) - for variable, dimension in participants] + actual_sizes = [ + (variable, dimension, bound_arguments[variable].dimensions[dimension]) + for variable, dimension in participants + ] reference_size = actual_sizes[0][2] index_sizes[index] = reference_size for variable, dimension, size in actual_sizes[1:]: if size != reference_size: - expected = ', '.join(f'{variable}.dimensions[{dimension}] == {size}' - for variable, dimension, size in actual_sizes) - raise ValueError(f'{self.assignment} expected all these dimensions of these tensors to be the same ' - f'because they share the index {index}: {expected}') + expected = ", ".join( + f"{variable}.dimensions[{dimension}] == {size}" + for variable, dimension, size in actual_sizes + ) + raise ValueError( + f"{self.assignment} expected all these dimensions of these tensors to be the same " + f"because they share the index {index}: {expected}" + ) # Determine output dimensions output_dimensions = tuple(index_sizes[index] for index in self.assignment.target.indexes) - cffi_output = allocate_taco_structure(tuple(mode.c_int for mode in self.output_format.modes), - output_dimensions, self.output_format.ordering) + cffi_output = allocate_taco_structure( + tuple(mode.c_int for mode in self.output_format.modes), + output_dimensions, + self.output_format.ordering, + ) output = Tensor(cffi_output) @@ -109,19 +138,23 @@ def __call__(self, *args, **kwargs): take_ownership_of_arrays(cffi_output) if return_value != 0: - raise RuntimeError(f'Taco function failed with error code {return_value}') + raise RuntimeError(f"Taco function failed with error code {return_value}") return output -def tensor_method(assignment: str, input_formats: Dict[str, str], output_format: str) -> PureTensorMethod: +def tensor_method( + assignment: str, input_formats: Dict[str, str], output_format: str +) -> PureTensorMethod: return cachable_tensor_method(assignment, tuple(input_formats.items()), output_format) @lru_cache() -def cachable_tensor_method(assignment: str, input_formats: Tuple[Tuple[str, str], ...], output_format: str - ) -> PureTensorMethod: +def cachable_tensor_method( + assignment: str, input_formats: Tuple[Tuple[str, str], ...], output_format: str +) -> PureTensorMethod: from .expression.parser import parse_assignment + parsed_assignment = parse_assignment(assignment).unwrap() parsed_input_formats = {name: parse_format(format).unwrap() for name, format in input_formats} @@ -129,7 +162,9 @@ def cachable_tensor_method(assignment: str, input_formats: Tuple[Tuple[str, str] parsed_output = parse_format(output_format).unwrap() if parsed_assignment.is_mutating(): - raise NotImplementedError(f'Mutating tensor assignments like {assignment} not implemented yet.') + raise NotImplementedError( + f"Mutating tensor assignments like {assignment} not implemented yet." + ) else: return PureTensorMethod(parsed_assignment, parsed_input_formats, parsed_output) @@ -143,4 +178,8 @@ def evaluate(assignment: str, output_format: str, **inputs: Tensor) -> Tensor: def format_to_taco_format(format: Format): - return ''.join(mode.character for mode in format.modes) + ':' + ','.join(map(str, format.ordering)) + return ( + "".join(mode.character for mode in format.modes) + + ":" + + ",".join(map(str, format.ordering)) + ) diff --git a/src/tensora/tensor.py b/src/tensora/tensor.py index 356399f..170417f 100644 --- a/src/tensora/tensor.py +++ b/src/tensora/tensor.py @@ -1,11 +1,11 @@ -__all__ = ['Tensor'] +__all__ = ["Tensor"] import itertools from numbers import Real -from typing import List, Tuple, Dict, Iterable, Union, Any, Iterator, Optional +from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union -from .format import Mode, Format, parse_format from .compile import taco_structure_to_cffi +from .format import Format, Mode, parse_format class Tensor: @@ -19,8 +19,12 @@ def __init__(self, cffi_tensor): self.cffi_tensor = cffi_tensor @staticmethod - def from_lol(lol, *, - dimensions: Optional[Tuple[int, ...]] = None, format: Union[Format, str, None] = None) -> 'Tensor': + def from_lol( + lol, + *, + dimensions: Optional[Tuple[int, ...]] = None, + format: Union[Format, str, None] = None, + ) -> "Tensor": if dimensions is None: dimensions = default_lol_dimensions(lol) @@ -34,13 +38,24 @@ def from_lol(lol, *, return Tensor.from_aos(coordinates, values, dimensions=dimensions, format=format) @staticmethod - def from_dok(dictionary: Dict[Tuple[int, ...], float], *, - dimensions: Optional[Tuple[int, ...]] = None, format: Union[Format, str, None] = None) -> 'Tensor': - return Tensor.from_aos(dictionary.keys(), dictionary.values(), dimensions=dimensions, format=format) + def from_dok( + dictionary: Dict[Tuple[int, ...], float], + *, + dimensions: Optional[Tuple[int, ...]] = None, + format: Union[Format, str, None] = None, + ) -> "Tensor": + return Tensor.from_aos( + dictionary.keys(), dictionary.values(), dimensions=dimensions, format=format + ) @staticmethod - def from_aos(coordinates: Iterable[Tuple[int, ...]], values: Iterable[float], *, - dimensions: Optional[Tuple[int, ...]] = None, format: Union[Format, str, None] = None) -> 'Tensor': + def from_aos( + coordinates: Iterable[Tuple[int, ...]], + values: Iterable[float], + *, + dimensions: Optional[Tuple[int, ...]] = None, + format: Union[Format, str, None] = None, + ) -> "Tensor": # Lengths of modes, dimensions, and elements in coordinates must be equal. Lengths of coordinates and values # must be equal if dimensions is None: @@ -54,7 +69,9 @@ def from_aos(coordinates: Iterable[Tuple[int, ...]], values: Iterable[float], *, # Reorder with first level first, etc. level_dimensions = tuple(dimensions[i] for i in format.ordering) - level_coordinates = [tuple(coordinate[i] for i in format.ordering) for coordinate in coordinates] + level_coordinates = [ + tuple(coordinate[i] for i in format.ordering) for coordinate in coordinates + ] tree = coordinates_to_tree(level_coordinates, values) @@ -62,20 +79,32 @@ def from_aos(coordinates: Iterable[Tuple[int, ...]], values: Iterable[float], *, cffi_modes = tuple(x.c_int for x in format.modes) - cffi_tensor = taco_structure_to_cffi(indexes, vals, mode_types=cffi_modes, dimensions=dimensions, - mode_ordering=format.ordering) + cffi_tensor = taco_structure_to_cffi( + indexes, + vals, + mode_types=cffi_modes, + dimensions=dimensions, + mode_ordering=format.ordering, + ) return Tensor(cffi_tensor) @staticmethod - def from_soa(coordinates: Tuple[Iterable[int], ...], values: Iterable[float], *, - dimensions: Optional[Tuple[int, ...]] = None, format: Union[Format, str, None] = None) -> 'Tensor': + def from_soa( + coordinates: Tuple[Iterable[int], ...], + values: Iterable[float], + *, + dimensions: Optional[Tuple[int, ...]] = None, + format: Union[Format, str, None] = None, + ) -> "Tensor": # Lengths of coordinates, modes, and dimensions must be equal. Lengths of elements of coordinates and values # must be equal transposed_coordinates = [*zip(*coordinates)] - return Tensor.from_aos(transposed_coordinates, values, dimensions=dimensions, format=format) + return Tensor.from_aos( + transposed_coordinates, values, dimensions=dimensions, format=format + ) @staticmethod def from_numpy(array, *, format: Union[Format, str, None] = None): @@ -107,12 +136,18 @@ def from_scipy_sparse(matrix, *, format: Union[Format, str, None] = None): soa_matrix = matrix.tocoo() - return Tensor.from_soa((soa_matrix.row, soa_matrix.col), soa_matrix.data, - dimensions=matrix.shape, format=format) + return Tensor.from_soa( + (soa_matrix.row, soa_matrix.col), + soa_matrix.data, + dimensions=matrix.shape, + format=format, + ) @staticmethod - def from_scalar(scalar: float) -> 'Tensor': - return Tensor(taco_structure_to_cffi([], [scalar], mode_types=(), dimensions=(), mode_ordering=())) + def from_scalar(scalar: float) -> "Tensor": + return Tensor( + taco_structure_to_cffi([], [scalar], mode_types=(), dimensions=(), mode_ordering=()) + ) def to_format(self, format: Union[Format, str]): return Tensor.from_dok(self.to_dok(), dimensions=self.dimensions, format=format) @@ -123,15 +158,17 @@ def order(self) -> int: @property def dimensions(self) -> Tuple[int, ...]: - return tuple(self.cffi_tensor.dimensions[0:self.order]) + return tuple(self.cffi_tensor.dimensions[0 : self.order]) @property def modes(self) -> Tuple[Mode, ...]: - return tuple(Mode.from_c_int(value) for value in self.cffi_tensor.mode_types[0:self.order]) + return tuple( + Mode.from_c_int(value) for value in self.cffi_tensor.mode_types[0 : self.order] + ) @property def mode_ordering(self) -> Tuple[int, ...]: - return tuple(self.cffi_tensor.mode_ordering[0:self.order]) + return tuple(self.cffi_tensor.mode_ordering[0 : self.order]) @property def taco_indices(self) -> List[List[List[int]]]: @@ -141,7 +178,7 @@ def taco_indices(self) -> List[List[List[int]]]: dimensions = self.dimensions modes = self.modes mode_ordering = self.mode_ordering - cffi_indexes = tensor_cdefs.cast('int32_t***', self.cffi_tensor.indices) + cffi_indexes = tensor_cdefs.cast("int32_t***", self.cffi_tensor.indices) indices = [] nnz = 1 @@ -150,8 +187,8 @@ def taco_indices(self) -> List[List[List[int]]]: indices.append([]) nnz *= dimensions[mode_ordering[i_dimension]] elif modes[i_dimension] == Mode.compressed: - pos = list(cffi_indexes[i_dimension][0][0:nnz + 1]) - crd = list(cffi_indexes[i_dimension][1][0:pos[-1]]) + pos = list(cffi_indexes[i_dimension][0][0 : nnz + 1]) + crd = list(cffi_indexes[i_dimension][1][0 : pos[-1]]) indices.append([pos, crd]) nnz = len(crd) @@ -165,7 +202,7 @@ def taco_vals(self) -> List[float]: dimensions = self.dimensions modes = self.modes mode_ordering = self.mode_ordering - cffi_indexes = tensor_cdefs.cast('int32_t***', self.cffi_tensor.indices) + cffi_indexes = tensor_cdefs.cast("int32_t***", self.cffi_tensor.indices) nnz = 1 for i_dimension in range(order): @@ -174,7 +211,7 @@ def taco_vals(self) -> List[float]: elif modes[i_dimension] == Mode.compressed: nnz = cffi_indexes[i_dimension][0][nnz] - cffi_vals = tensor_cdefs.cast('double*', self.cffi_tensor.vals) + cffi_vals = tensor_cdefs.cast("double*", self.cffi_tensor.vals) return list(cffi_vals[0:nnz]) @property @@ -188,8 +225,8 @@ def items(self) -> Iterator[Tuple[Tuple[int, ...], float]]: modes = self.modes dimensions = self.dimensions mode_ordering = self.mode_ordering - cffi_indexes = tensor_cdefs.cast('int32_t***', self.cffi_tensor.indices) - cffi_values = tensor_cdefs.cast('double*', self.cffi_tensor.vals) + cffi_indexes = tensor_cdefs.cast("int32_t***", self.cffi_tensor.indices) + cffi_values = tensor_cdefs.cast("double*", self.cffi_tensor.vals) level_dimensions = [dimensions[i] for i in mode_ordering] def recurse(i_level, prefix, position): @@ -227,23 +264,23 @@ def to_numpy(self): return array - def __add__(self, other) -> 'Tensor': - return evaluate_binary_operator(self, other, '+') + def __add__(self, other) -> "Tensor": + return evaluate_binary_operator(self, other, "+") - def __radd__(self, other) -> 'Tensor': - return evaluate_binary_operator(other, self, '+') + def __radd__(self, other) -> "Tensor": + return evaluate_binary_operator(other, self, "+") - def __sub__(self, other) -> 'Tensor': - return evaluate_binary_operator(self, other, '-') + def __sub__(self, other) -> "Tensor": + return evaluate_binary_operator(self, other, "-") - def __rsub__(self, other) -> 'Tensor': - return evaluate_binary_operator(other, self, '-') + def __rsub__(self, other) -> "Tensor": + return evaluate_binary_operator(other, self, "-") - def __mul__(self, other) -> 'Tensor': - return evaluate_binary_operator(self, other, '*') + def __mul__(self, other) -> "Tensor": + return evaluate_binary_operator(self, other, "*") - def __rmul__(self, other) -> 'Tensor': - return evaluate_binary_operator(other, self, '*') + def __rmul__(self, other) -> "Tensor": + return evaluate_binary_operator(other, self, "*") def __matmul__(self, other): return evaluate_matrix_multiplication_operator(self, other) @@ -255,27 +292,29 @@ def __float__(self): from .compile import tensor_cdefs if self.order != 0: - raise ValueError(f'Can only convert Tensor of order 0 to float, not order {self.order}') + raise ValueError( + f"Can only convert Tensor of order 0 to float, not order {self.order}" + ) - cffi_vals = tensor_cdefs.cast('double*', self.cffi_tensor.vals) + cffi_vals = tensor_cdefs.cast("double*", self.cffi_tensor.vals) return cffi_vals[0] def __getstate__(self): return { - 'dimensions': self.dimensions, - 'mode_types': tuple(mode.c_int for mode in self.format.modes), - 'mode_ordering': self.format.ordering, - 'indices': self.taco_indices, - 'vals': self.taco_vals, + "dimensions": self.dimensions, + "mode_types": tuple(mode.c_int for mode in self.format.modes), + "mode_ordering": self.format.ordering, + "indices": self.taco_indices, + "vals": self.taco_vals, } def __setstate__(self, state): self.cffi_tensor = taco_structure_to_cffi( - indices=state['indices'], - vals=state['vals'], - mode_types=state['mode_types'], - dimensions=state['dimensions'], - mode_ordering=state['mode_ordering'], + indices=state["indices"], + vals=state["vals"], + mode_types=state["mode_types"], + dimensions=state["dimensions"], + mode_ordering=state["mode_ordering"], ) def __eq__(self, other): @@ -285,11 +324,12 @@ def __eq__(self, other): return NotImplemented def __repr__(self): - return f'Tensor.from_dok({str(self.to_dok())}, dimensions={self.dimensions}, format={self.format.deparse()!r})' + return f"Tensor.from_dok({str(self.to_dok())}, dimensions={self.dimensions}, format={self.format.deparse()!r})" -def lol_to_coordinates_and_values(data: Any, keep_zero: bool = False - ) -> Tuple[Iterable[Tuple[int, ...]], Iterable[float]]: +def lol_to_coordinates_and_values( + data: Any, keep_zero: bool = False +) -> Tuple[Iterable[Tuple[int, ...]], Iterable[float]]: coordinates = [] values = [] @@ -334,8 +374,9 @@ def recurse(node: Dict[int, Any], remaining_coordinates: Tuple[int, ...], payloa return tree -def tree_to_indices_and_values(tree: Any, modes: Tuple[Mode, ...], dimensions: Tuple[int, ...] - ) -> Tuple[List[List[List[int]]], List[float]]: +def tree_to_indices_and_values( + tree: Any, modes: Tuple[Mode, ...], dimensions: Tuple[int, ...] +) -> Tuple[List[List[List[int]]], List[float]]: order = len(modes) # Initialize indexes structure @@ -383,59 +424,79 @@ def recurse(node, i_level): return indexes, values -def evaluate_binary_operator(left: Union[Tensor, Real], right: Union[Tensor, Real], operator: str) -> Tensor: +def evaluate_binary_operator( + left: Union[Tensor, Real], right: Union[Tensor, Real], operator: str +) -> Tensor: from .function import evaluate def indexes_string(tensor): - return ','.join(f'i{i}' for i in range(tensor.order)) + return ",".join(f"i{i}" for i in range(tensor.order)) if isinstance(left, Tensor) and isinstance(right, Tensor): if left.dimensions != right.dimensions: - raise ValueError(f'Cannot apply operator {operator} between tensor with dimensions {left.dimensions} and ' - f'tensor with dimensions {right.dimensions}') + raise ValueError( + f"Cannot apply operator {operator} between tensor with dimensions {left.dimensions} and " + f"tensor with dimensions {right.dimensions}" + ) - if operator == '*': + if operator == "*": # Output has density of least dense tensor - output_format = ''.join('d' if mode1 == Mode.dense and mode2 == Mode.dense else 's' - for mode1, mode2 in zip(left.format.modes, right.format.modes)) - elif operator in ('+', '-'): + output_format = "".join( + "d" if mode1 == Mode.dense and mode2 == Mode.dense else "s" + for mode1, mode2 in zip(left.format.modes, right.format.modes) + ) + elif operator in ("+", "-"): # Output has density of most dense tensor - output_format = ''.join('d' if mode1 == Mode.dense or mode2 == Mode.dense else 's' - for mode1, mode2 in zip(left.format.modes, right.format.modes)) + output_format = "".join( + "d" if mode1 == Mode.dense or mode2 == Mode.dense else "s" + for mode1, mode2 in zip(left.format.modes, right.format.modes) + ) else: raise NotImplementedError() indexes = indexes_string(left) - return evaluate(f'output({indexes}) = left({indexes}) {operator} right({indexes})', - output_format, left=left, right=right) + return evaluate( + f"output({indexes}) = left({indexes}) {operator} right({indexes})", + output_format, + left=left, + right=right, + ) elif isinstance(left, Tensor) and isinstance(right, Real): - if operator == '*': + if operator == "*": # Output has density of tensor output_format = left.format.deparse() - elif operator in ('+', '-'): + elif operator in ("+", "-"): # Output is full dense - output_format = 'd' * left.order + output_format = "d" * left.order else: raise NotImplementedError() indexes = indexes_string(left) - return evaluate(f'output({indexes}) = left({indexes}) {operator} right', - output_format, left=left, right=Tensor.from_scalar(float(right))) + return evaluate( + f"output({indexes}) = left({indexes}) {operator} right", + output_format, + left=left, + right=Tensor.from_scalar(float(right)), + ) elif isinstance(left, Real) and isinstance(right, Tensor): - if operator == '*': + if operator == "*": # Output has density of tensor output_format = right.format.deparse() - elif operator in ('+', '-'): + elif operator in ("+", "-"): # Output is full dense - output_format = 'd' * right.order + output_format = "d" * right.order else: raise NotImplementedError() indexes = indexes_string(right) - return evaluate(f'output({indexes}) = left {operator} right({indexes})', - output_format, left=Tensor.from_scalar(float(left)), right=right) + return evaluate( + f"output({indexes}) = left {operator} right({indexes})", + output_format, + left=Tensor.from_scalar(float(left)), + right=right, + ) else: return NotImplemented @@ -446,25 +507,33 @@ def evaluate_matrix_multiplication_operator(left: Tensor, right: Tensor): if isinstance(left, Tensor) and isinstance(right, Tensor): if left.order == 1 and right.order == 1: - scalar_tensor = evaluate('output = left(i) * right(i)', '', left=left, right=right) + scalar_tensor = evaluate("output = left(i) * right(i)", "", left=left, right=right) return float(scalar_tensor) elif left.order == 2 and right.order == 1: # Output format is the uncontracted dimension of the matrix output_format = left.format.modes[left.format.ordering[0]].character - return evaluate('output(i) = left(i,j) * right(j)', output_format, left=left, right=right) + return evaluate( + "output(i) = left(i,j) * right(j)", output_format, left=left, right=right + ) elif left.order == 1 and right.order == 2: # Output format is the uncontracted dimension of the matrix output_format = right.format.modes[right.format.ordering[1]].character - return evaluate('output(j) = left(i) * right(i,j)', output_format, left=left, right=right) + return evaluate( + "output(j) = left(i) * right(i,j)", output_format, left=left, right=right + ) elif left.order == 2 and right.order == 2: # Output format are the uncontracted dimensions of the matrices left_output_format = left.format.modes[left.format.ordering[0]].character right_output_format = right.format.modes[right.format.ordering[1]].character output_format = left_output_format + right_output_format - return evaluate('output(i,k) = left(i,j) * right(j,k)', output_format, left=left, right=right) + return evaluate( + "output(i,k) = left(i,j) * right(j,k)", output_format, left=left, right=right + ) else: - raise ValueError(f'Matrix multiply is only defined between tensors of orders 1 and 2, not orders ' - f'{left.order} and {right.order}') + raise ValueError( + f"Matrix multiply is only defined between tensors of orders 1 and 2, not orders " + f"{left.order} and {right.order}" + ) else: return NotImplemented @@ -506,8 +575,10 @@ def default_aos_dimensions(coordinates: Iterable[Tuple[int, ...]]) -> Tuple[int, maximums = list(coordinate) else: if len(coordinate) != order: - raise ValueError(f'All coordinates must be the same length; the first coordinate has length' - f'{order}, but this coordinate is not that length: {coordinate}') + raise ValueError( + f"All coordinates must be the same length; the first coordinate has length" + f"{order}, but this coordinate is not that length: {coordinate}" + ) for i, (dimension, index) in enumerate(zip(maximums, coordinate)): if index > dimension: @@ -527,12 +598,19 @@ def default_format_given_nnz(dimensions: Tuple[int, ...], nnz: int) -> Format: if nnz < required_threshold: break - return Format((Mode.dense,) * needed_dense + (Mode.compressed,) * (len(dimensions) - needed_dense), - tuple(range(len(dimensions)))) + return Format( + (Mode.dense,) * needed_dense + (Mode.compressed,) * (len(dimensions) - needed_dense), + tuple(range(len(dimensions))), + ) -def taco_indexes_from_aos_coordinates(coordinates: Iterable[Tuple[int, ...]], values: Iterable[float], *, - modes: Tuple[Mode, ...], dimensions=Tuple[int, ...]): # pragma: no cover +def taco_indexes_from_aos_coordinates( + coordinates: Iterable[Tuple[int, ...]], + values: Iterable[float], + *, + modes: Tuple[Mode, ...], + dimensions=Tuple[int, ...], +): # pragma: no cover # This is an experimental alternative to coordinates_to_tree and tree_to_indices_and_values. It is not currently # used anywhere. @@ -583,7 +661,7 @@ def __iter__(self): idx = [] unique_coordinates = [] previous_coordinate = None - for coordinate in zip(*soa_coordinates[0:i_level + 1]): + for coordinate in zip(*soa_coordinates[0 : i_level + 1]): if coordinate != previous_coordinate: while coordinate[0:-1] != current_prefix: # The prefix has changed. Mark in pos the position of this prefix. Some prefixes may be @@ -601,6 +679,6 @@ def __iter__(self): levels.append([pos, idx]) previous_prefixes = unique_coordinates else: - raise RuntimeError(f'Unknown mode: {mode}') + raise RuntimeError(f"Unknown mode: {mode}") return levels diff --git a/tests/test_combinatorically.py b/tests/test_combinatorically.py index 362b3f3..0262022 100644 --- a/tests/test_combinatorically.py +++ b/tests/test_combinatorically.py @@ -4,87 +4,105 @@ def assert_same_as_dense(expression, format_out, **tensor_pairs): - tensors_in_format = {name: Tensor.from_lol(data, format=format) for name, (data, format) in tensor_pairs.items()} + tensors_in_format = { + name: Tensor.from_lol(data, format=format) for name, (data, format) in tensor_pairs.items() + } tensors_as_dense = {name: Tensor.from_lol(data) for name, (data, _) in tensor_pairs.items()} actual = evaluate(expression, format_out, **tensors_in_format) - expected = evaluate(expression, ''.join('d' for c in format_out if c in ('d', 's')), **tensors_as_dense) + expected = evaluate( + expression, "".join("d" for c in format_out if c in ("d", "s")), **tensors_as_dense + ) assert actual == expected -@pytest.mark.parametrize('dense', [[3, 2, 4], [0, 0, 0]]) -@pytest.mark.parametrize('format_in', ['s', 'd']) -@pytest.mark.parametrize('format_out', ['s', 'd']) +@pytest.mark.parametrize("dense", [[3, 2, 4], [0, 0, 0]]) +@pytest.mark.parametrize("format_in", ["s", "d"]) +@pytest.mark.parametrize("format_out", ["s", "d"]) def test_copy_1(dense, format_in, format_out): a = Tensor.from_lol(dense, format=format_in) - actual = evaluate('b(i) = a(i)', format_out, a=a) + actual = evaluate("b(i) = a(i)", format_out, a=a) assert actual == a -@pytest.mark.skip('taco fails to compile most of these') -@pytest.mark.parametrize('dense', [[[0, 2, 4], [0, -1, 0]], [[0, 0, 0], [0, 0, 0]]]) -@pytest.mark.parametrize('format_in', ['ss', 'dd', 'sd', 'ds', 's1s0', 'd1d0', 's1d0', 'd1s0']) -@pytest.mark.parametrize('format_out', ['ss', 'dd', 'sd', 'ds', 's1s0', 'd1d0', 's1d0', 'd1s0']) +@pytest.mark.skip("taco fails to compile most of these") +@pytest.mark.parametrize("dense", [[[0, 2, 4], [0, -1, 0]], [[0, 0, 0], [0, 0, 0]]]) +@pytest.mark.parametrize("format_in", ["ss", "dd", "sd", "ds", "s1s0", "d1d0", "s1d0", "d1s0"]) +@pytest.mark.parametrize("format_out", ["ss", "dd", "sd", "ds", "s1s0", "d1d0", "s1d0", "d1s0"]) def test_copy_2(dense, format_in, format_out): a = Tensor.from_lol(dense, format=format_in) - actual = evaluate('b(i,j) = a(i,j)', format_out, a=a) + actual = evaluate("b(i,j) = a(i,j)", format_out, a=a) assert actual == a -@pytest.mark.skip('taco fails on all of these') -@pytest.mark.parametrize('dense', [[[0, 2, 4], [0, -1, 0], [2, 0, 3]], [[0, 0, 0], [0, 0, 0], [0, 0, 0]]]) -@pytest.mark.parametrize('format_in', ['ss', 'dd', 'sd', 'ds', 's1s0', 'd1d0', 's1d0', 'd1s0']) -@pytest.mark.parametrize('format_out', ['s', 'd']) +@pytest.mark.skip("taco fails on all of these") +@pytest.mark.parametrize( + "dense", [[[0, 2, 4], [0, -1, 0], [2, 0, 3]], [[0, 0, 0], [0, 0, 0], [0, 0, 0]]] +) +@pytest.mark.parametrize("format_in", ["ss", "dd", "sd", "ds", "s1s0", "d1d0", "s1d0", "d1s0"]) +@pytest.mark.parametrize("format_out", ["s", "d"]) def test_diag(dense, format_in, format_out): - assert_same_as_dense('diagA(i) = A(i,i)', format_out, A=(dense, format_in)) + assert_same_as_dense("diagA(i) = A(i,i)", format_out, A=(dense, format_in)) -@pytest.mark.parametrize('dense1', [[0, 2, 4, 0], [0, 0, 0, 0]]) -@pytest.mark.parametrize('dense2', [[-1, 3.5, 0, 0], [0, 0, 0, 0]]) -@pytest.mark.parametrize('format1', ['s', 'd']) -@pytest.mark.parametrize('format2', ['s', 'd']) +@pytest.mark.parametrize("dense1", [[0, 2, 4, 0], [0, 0, 0, 0]]) +@pytest.mark.parametrize("dense2", [[-1, 3.5, 0, 0], [0, 0, 0, 0]]) +@pytest.mark.parametrize("format1", ["s", "d"]) +@pytest.mark.parametrize("format2", ["s", "d"]) def test_vector_dot(dense1, dense2, format1, format2): - assert_same_as_dense('out = in1(i) * in2(i)', '', in1=(dense1, format1), in2=(dense2, format2)) + assert_same_as_dense("out = in1(i) * in2(i)", "", in1=(dense1, format1), in2=(dense2, format2)) -@pytest.mark.parametrize('dense1', [[0, 2, 4, 0], [0, 0, 0, 0]]) -@pytest.mark.parametrize('dense2', [[-1, 3.5, 0, 0], [0, 0, 0, 0]]) -@pytest.mark.parametrize('format1', ['s', 'd']) -@pytest.mark.parametrize('format2', ['s', 'd']) -@pytest.mark.parametrize('format_out', ['s', 'd']) -@pytest.mark.parametrize('operator', ['+', '-', '*']) +@pytest.mark.parametrize("dense1", [[0, 2, 4, 0], [0, 0, 0, 0]]) +@pytest.mark.parametrize("dense2", [[-1, 3.5, 0, 0], [0, 0, 0, 0]]) +@pytest.mark.parametrize("format1", ["s", "d"]) +@pytest.mark.parametrize("format2", ["s", "d"]) +@pytest.mark.parametrize("format_out", ["s", "d"]) +@pytest.mark.parametrize("operator", ["+", "-", "*"]) def test_vector_binary(operator, dense1, dense2, format1, format2, format_out): - assert_same_as_dense(f'out(i) = in1(i) {operator} in2(i)', format_out, in1=(dense1, format1), in2=(dense2, format2)) - - -@pytest.mark.skip('taco fails to compile most of these') -@pytest.mark.parametrize('dense1', [[[0, 2, 4], [0, -1, 0]], [[0, 0, 0], [0, 0, 0]]]) -@pytest.mark.parametrize('dense2', [[[-1, 3.5], [0, 0], [4, 0]], [[0, 0], [0, 0], [0, 0]]]) -@pytest.mark.parametrize('format1', ['ss', 'dd', 'sd', 'ds', 's1s0', 'd1d0', 's1d0', 'd1s0']) -@pytest.mark.parametrize('format2', ['ss', 'dd', 'sd', 'ds', 's1s0', 'd1d0', 's1d0', 'd1s0']) -@pytest.mark.parametrize('format_out', ['ss', 'dd', 'sd', 'ds', 's1s0', 'd1d0', 's1d0', 'd1s0']) + assert_same_as_dense( + f"out(i) = in1(i) {operator} in2(i)", + format_out, + in1=(dense1, format1), + in2=(dense2, format2), + ) + + +@pytest.mark.skip("taco fails to compile most of these") +@pytest.mark.parametrize("dense1", [[[0, 2, 4], [0, -1, 0]], [[0, 0, 0], [0, 0, 0]]]) +@pytest.mark.parametrize("dense2", [[[-1, 3.5], [0, 0], [4, 0]], [[0, 0], [0, 0], [0, 0]]]) +@pytest.mark.parametrize("format1", ["ss", "dd", "sd", "ds", "s1s0", "d1d0", "s1d0", "d1s0"]) +@pytest.mark.parametrize("format2", ["ss", "dd", "sd", "ds", "s1s0", "d1d0", "s1d0", "d1s0"]) +@pytest.mark.parametrize("format_out", ["ss", "dd", "sd", "ds", "s1s0", "d1d0", "s1d0", "d1s0"]) def test_matrix_dot(dense1, dense2, format1, format2, format_out): - assert_same_as_dense('out(i,k) = in1(i,j) * in2(j,k)', format_out, in1=(dense1, format1), in2=(dense2, format2)) + assert_same_as_dense( + "out(i,k) = in1(i,j) * in2(j,k)", format_out, in1=(dense1, format1), in2=(dense2, format2) + ) -@pytest.mark.parametrize('dense1', [[[0, 2, 4], [0, -1, 0]], [[0, 0, 0], [0, 0, 0]]]) -@pytest.mark.parametrize('dense2', [[-1, 3.5, 0], [0, 0, 0]]) -@pytest.mark.parametrize('format1', ['ss', 'dd', 'sd', 'ds', 's1s0', 'd1d0', 's1d0', 'd1s0']) -@pytest.mark.parametrize('format2', ['s', 'd']) -@pytest.mark.parametrize('format_out', ['d']) +@pytest.mark.parametrize("dense1", [[[0, 2, 4], [0, -1, 0]], [[0, 0, 0], [0, 0, 0]]]) +@pytest.mark.parametrize("dense2", [[-1, 3.5, 0], [0, 0, 0]]) +@pytest.mark.parametrize("format1", ["ss", "dd", "sd", "ds", "s1s0", "d1d0", "s1d0", "d1s0"]) +@pytest.mark.parametrize("format2", ["s", "d"]) +@pytest.mark.parametrize("format_out", ["d"]) def test_matrix_vector_product(dense1, dense2, format1, format2, format_out): - assert_same_as_dense('out(i) = in1(i,j) * in2(j)', format_out, in1=(dense1, format1), in2=(dense2, format2)) - - -@pytest.mark.parametrize('dense1', [[[0, 2, 4], [0, -1, 0]], [[0, 0, 0], [0, 0, 0]]]) -@pytest.mark.parametrize('dense2', [[[-1, 3.5], [0, 0], [4, 0]], [[0, 0], [0, 0], [0, 0]]]) -@pytest.mark.parametrize('dense3', [[[-3, 0], [7, 0]], [[0, 0], [0, 0]]]) -@pytest.mark.parametrize('format1', ['dd', 'ds']) -@pytest.mark.parametrize('format2', ['dd']) -@pytest.mark.parametrize('format3', ['dd', 'ds']) -@pytest.mark.parametrize('format_out', ['dd']) + assert_same_as_dense( + "out(i) = in1(i,j) * in2(j)", format_out, in1=(dense1, format1), in2=(dense2, format2) + ) + + +@pytest.mark.parametrize("dense1", [[[0, 2, 4], [0, -1, 0]], [[0, 0, 0], [0, 0, 0]]]) +@pytest.mark.parametrize("dense2", [[[-1, 3.5], [0, 0], [4, 0]], [[0, 0], [0, 0], [0, 0]]]) +@pytest.mark.parametrize("dense3", [[[-3, 0], [7, 0]], [[0, 0], [0, 0]]]) +@pytest.mark.parametrize("format1", ["dd", "ds"]) +@pytest.mark.parametrize("format2", ["dd"]) +@pytest.mark.parametrize("format3", ["dd", "ds"]) +@pytest.mark.parametrize("format_out", ["dd"]) def test_matrix_multiply_add(dense1, dense2, dense3, format1, format2, format3, format_out): - assert_same_as_dense('out(i,k) = in1(i,j) * in2(j,k) + in3(i,k)', format_out, - in1=(dense1, format1), - in2=(dense2, format2), - in3=(dense3, format3)) + assert_same_as_dense( + "out(i,k) = in1(i,j) * in2(j,k) + in3(i,k)", + format_out, + in1=(dense1, format1), + in2=(dense2, format2), + in3=(dense3, format3), + ) diff --git a/tests/test_compile.py b/tests/test_compile.py index 30181c2..0712ab5 100644 --- a/tests/test_compile.py +++ b/tests/test_compile.py @@ -3,8 +3,14 @@ from cffi import FFI from tensora import Tensor -from tensora.compile import tensor_cdefs, taco_define_header, taco_type_header, lock, \ - take_ownership_of_tensor_members, take_ownership_of_tensor +from tensora.compile import ( + lock, + taco_define_header, + taco_type_header, + take_ownership_of_tensor, + take_ownership_of_tensor_members, + tensor_cdefs, +) source = """ taco_tensor_t create_tensor() { @@ -71,14 +77,19 @@ ffi = FFI() ffi.include(tensor_cdefs) -ffi.cdef(""" +ffi.cdef( + """ taco_tensor_t create_tensor(); taco_tensor_t* create_pointer_to_tensor(); -""") -ffi.set_source('taco_kernel', taco_define_header + taco_type_header + source, - extra_compile_args=['-Wno-unused-variable', '-Wno-unknown-pragmas']) - -expected_tensor = Tensor.from_lol([[6, 0, 9, 8], [0, 0, 0, 0], [5, 0, 0, 7]], format='ds') +""" +) +ffi.set_source( + "taco_kernel", + taco_define_header + taco_type_header + source, + extra_compile_args=["-Wno-unused-variable", "-Wno-unknown-pragmas"], +) + +expected_tensor = Tensor.from_lol([[6, 0, 9, 8], [0, 0, 0, 0], [5, 0, 0, 7]], format="ds") with tempfile.TemporaryDirectory() as temp_dir: # Lock because FFI.compile is not thread safe: https://foss.heptapod.net/pypy/cffi/-/issues/490 diff --git a/tests/test_expression.py b/tests/test_expression.py index 82ccaee..86e205b 100644 --- a/tests/test_expression.py +++ b/tests/test_expression.py @@ -4,30 +4,50 @@ from tensora.expression.parser import parse_assignment assignment_strings = [ - ('A(i) = B(i,j) * C(j)', - Assignment(Tensor('A', ['i']), Multiply(Tensor('B', ['i', 'j']), Tensor('C', ['j'])))), - ('ab(i) = a(i) + b(i)', - Assignment(Tensor('ab', ['i']), Add(Tensor('a', ['i']), Tensor('b', ['i'])))), - ('D(i) = A(i) - B(i)', - Assignment(Tensor('D', ['i']), Subtract(Tensor('A', ['i']), Tensor('B', ['i'])))), - ('A2(i) = 2 * a * A(i)', - Assignment(Tensor('A2', ['i']), Multiply(Multiply(Integer(2), Scalar('a')), Tensor('A', ['i'])))), - ('B2(i) = 2.0 * B(i)', - Assignment(Tensor('B2', ['i']), Multiply(Float(2.0), Tensor('B', ['i'])))), - ('ab2(i) = 2.0 * (a(i) + b(i))', - Assignment(Tensor('ab2', ['i']), Multiply(Float(2.0), Add(Tensor('a', ['i']), Tensor('b', ['i']))))), - ('ab2(i) = (a(i) + b(i)) * 2.0', - Assignment(Tensor('ab2', ['i']), Multiply(Add(Tensor('a', ['i']), Tensor('b', ['i'])), Float(2.0)))), + ( + "A(i) = B(i,j) * C(j)", + Assignment(Tensor("A", ["i"]), Multiply(Tensor("B", ["i", "j"]), Tensor("C", ["j"]))), + ), + ( + "ab(i) = a(i) + b(i)", + Assignment(Tensor("ab", ["i"]), Add(Tensor("a", ["i"]), Tensor("b", ["i"]))), + ), + ( + "D(i) = A(i) - B(i)", + Assignment(Tensor("D", ["i"]), Subtract(Tensor("A", ["i"]), Tensor("B", ["i"]))), + ), + ( + "A2(i) = 2 * a * A(i)", + Assignment( + Tensor("A2", ["i"]), Multiply(Multiply(Integer(2), Scalar("a")), Tensor("A", ["i"])) + ), + ), + ( + "B2(i) = 2.0 * B(i)", + Assignment(Tensor("B2", ["i"]), Multiply(Float(2.0), Tensor("B", ["i"]))), + ), + ( + "ab2(i) = 2.0 * (a(i) + b(i))", + Assignment( + Tensor("ab2", ["i"]), Multiply(Float(2.0), Add(Tensor("a", ["i"]), Tensor("b", ["i"]))) + ), + ), + ( + "ab2(i) = (a(i) + b(i)) * 2.0", + Assignment( + Tensor("ab2", ["i"]), Multiply(Add(Tensor("a", ["i"]), Tensor("b", ["i"])), Float(2.0)) + ), + ), ] -@pytest.mark.parametrize('string,assignment', assignment_strings) +@pytest.mark.parametrize("string,assignment", assignment_strings) def test_assignment_parsing(string, assignment): actual = parse_assignment(string).unwrap() assert actual == assignment -@pytest.mark.parametrize('string,assignment', assignment_strings) +@pytest.mark.parametrize("string,assignment", assignment_strings) def test_assignment_deparsing(string, assignment): deparsed = assignment.deparse() assert deparsed == string @@ -38,32 +58,49 @@ def parse(string): def test_assignment_to_string(): - string = 'A(i) = 2 * B(i,j) * (C(j) + D(j))' + string = "A(i) = 2 * B(i,j) * (C(j) + D(j))" assert str(parse(string)) == string -@pytest.mark.parametrize('string,output', [ - ('A(i) = B(i,j) * C(j)', False), - ('y(i) = X(i,j) * y(i)', True), - ('z(i) += X(i,j) * y(i)', True), -]) +@pytest.mark.parametrize( + "string,output", + [ + ("A(i) = B(i,j) * C(j)", False), + ("y(i) = X(i,j) * y(i)", True), + ("z(i) += X(i,j) * y(i)", True), + ], +) def test_is_mutating(string, output): assert parse(string).is_mutating() == output -@pytest.mark.parametrize('string,output', [ - ('y(i) = 0.5 * (b - a) * (x1(i,j) + x2(i,j)) * z(j)', {'y': 1, 'b': 0, 'a': 0, 'x1': 2, 'x2': 2, 'z': 1}), - ('B2(i,k) = B(i,j) * B(j,k)', {'B2': 2, 'B': 2}), -]) +@pytest.mark.parametrize( + "string,output", + [ + ( + "y(i) = 0.5 * (b - a) * (x1(i,j) + x2(i,j)) * z(j)", + {"y": 1, "b": 0, "a": 0, "x1": 2, "x2": 2, "z": 1}, + ), + ("B2(i,k) = B(i,j) * B(j,k)", {"B2": 2, "B": 2}), + ], +) def test_variable_order(string, output): assert parse(string).variable_orders() == output -@pytest.mark.parametrize('string,output', [ - ('y(i) = 0.5 * (b - a) * (x1(i,j) + x2(i,j)) * z(j)', {'i': {('y', 0), ('x1', 0), ('x2', 0)}, - 'j': {('x1', 1), ('x2', 1), ('z', 0)}}), - ('B2(i,k) = B(i,j) * B(j,k)', {'i': {('B2', 0), ('B', 0)}, 'j': {('B', 1), ('B', 0)}, 'k': {('B2', 1), ('B', 1)}}), - ('diagA2(i) = A(i,i) + A(i,i)', {'i': {('diagA2', 0), ('A', 0), ('A', 1)}}), -]) +@pytest.mark.parametrize( + "string,output", + [ + ( + "y(i) = 0.5 * (b - a) * (x1(i,j) + x2(i,j)) * z(j)", + {"i": {("y", 0), ("x1", 0), ("x2", 0)}, "j": {("x1", 1), ("x2", 1), ("z", 0)}}, + ), + ( + "B2(i,k) = B(i,j) * B(j,k)", + {"i": {("B2", 0), ("B", 0)}, "j": {("B", 1), ("B", 0)}, "k": {("B2", 1), ("B", 1)}}, + ), + ("diagA2(i) = A(i,i) + A(i,i)", {"i": {("diagA2", 0), ("A", 0), ("A", 1)}}), + ], +) def test_index_participants(string, output): assert parse(string).index_participants() == output diff --git a/tests/test_format.py b/tests/test_format.py index 919e0cf..199b0fe 100644 --- a/tests/test_format.py +++ b/tests/test_format.py @@ -5,29 +5,29 @@ from tensora.format import parse_format format_strings = [ - ('', Format((), ())), - ('d', Format((Mode.dense,), (0,))), - ('s', Format((Mode.compressed,), (0,))), - ('ds', Format((Mode.dense, Mode.compressed), (0, 1))), - ('sd', Format((Mode.compressed, Mode.dense), (0, 1))), - ('d1s0', Format((Mode.dense, Mode.compressed), (1, 0))), - ('d1s0s2', Format((Mode.dense, Mode.compressed, Mode.compressed), (1, 0, 2))), + ("", Format((), ())), + ("d", Format((Mode.dense,), (0,))), + ("s", Format((Mode.compressed,), (0,))), + ("ds", Format((Mode.dense, Mode.compressed), (0, 1))), + ("sd", Format((Mode.compressed, Mode.dense), (0, 1))), + ("d1s0", Format((Mode.dense, Mode.compressed), (1, 0))), + ("d1s0s2", Format((Mode.dense, Mode.compressed, Mode.compressed), (1, 0, 2))), ] -@pytest.mark.parametrize('string,format', format_strings) +@pytest.mark.parametrize("string,format", format_strings) def test_parse_format(string, format): actual = parse_format(string).unwrap() assert actual == format -@pytest.mark.parametrize('string,format', format_strings) +@pytest.mark.parametrize("string,format", format_strings) def test_deparse_format(string, format): actual = format.deparse() assert actual == string -@pytest.mark.parametrize('string', ['df', '1d0s', 'd0s', 'd0s1s1', 'd1s2s3']) +@pytest.mark.parametrize("string", ["df", "1d0s", "d0s", "d0s1s1", "d1s2s3"]) def test_parse_bad_format(string): actual = parse_format(string) assert isinstance(actual, Failure) @@ -43,13 +43,13 @@ def test_format_attributes(): def test_mode_dense_attributes(): mode_dense = Mode.from_c_int(0) assert mode_dense.c_int == 0 - assert mode_dense.character == 'd' + assert mode_dense.character == "d" def test_mode_sparse_attributes(): mode_dense = Mode.from_c_int(1) assert mode_dense.c_int == 1 - assert mode_dense.character == 's' + assert mode_dense.character == "s" def test_mode_from_illegal_int(): diff --git a/tests/test_from_data.py b/tests/test_from_data.py index a8dbe48..7b5b72b 100644 --- a/tests/test_from_data.py +++ b/tests/test_from_data.py @@ -1,6 +1,6 @@ import pytest -from tensora import Mode, Format, Tensor +from tensora import Format, Mode, Tensor try: # numpy is an optional dependency of tensora. It is only used to convert from numpy. @@ -23,7 +23,7 @@ def test_from_dok(): (1, 2, 0): 4.5, } - actual = Tensor.from_dok(dok, dimensions=(3, 3, 2), format='sss') + actual = Tensor.from_dok(dok, dimensions=(3, 3, 2), format="sss") assert actual.to_dok() == dok assert actual.format == Format((Mode.compressed, Mode.compressed, Mode.compressed), (0, 1, 2)) @@ -62,7 +62,7 @@ def test_from_soa(): 4.5, ] - actual = Tensor.from_soa(soa, values, dimensions=(3, 3, 2), format='sss') + actual = Tensor.from_soa(soa, values, dimensions=(3, 3, 2), format="sss") assert actual.to_dok() == dok assert actual.format == Format((Mode.compressed, Mode.compressed, Mode.compressed), (0, 1, 2)) @@ -115,7 +115,7 @@ def test_from_aos(): 4.5, ] - actual = Tensor.from_aos(aos, values, dimensions=(3, 3, 2), format='sss') + actual = Tensor.from_aos(aos, values, dimensions=(3, 3, 2), format="sss") assert actual.to_dok() == dok assert actual.format == Format((Mode.compressed, Mode.compressed, Mode.compressed), (0, 1, 2)) @@ -161,7 +161,7 @@ def test_from_lol(): lol = [[[0, 0], [0, 0], [0, 2]], [[-2, 0], [0, 0], [4.5, 4.5]], [[0, 0], [0, 0], [0, 0]]] - actual = Tensor.from_lol(lol, dimensions=(3, 3, 2), format='sss') + actual = Tensor.from_lol(lol, dimensions=(3, 3, 2), format="sss") assert actual.to_dok() == dok assert actual.format == Format((Mode.compressed, Mode.compressed, Mode.compressed), (0, 1, 2)) @@ -207,8 +207,9 @@ def recurse(tree, indexes): tensor_from_dok = Tensor.from_dok(dok, dimensions=dimensions, format=format) tensor_from_dok_no_zeros = Tensor.from_dok(dok_no_zeros, dimensions=dimensions, format=format) tensor_from_aos = Tensor.from_aos(coordinates, values, dimensions=dimensions, format=format) - tensor_from_aos_no_zeros = Tensor.from_aos(coordinates_no_zeros, values_no_zeros, dimensions=dimensions, - format=format) + tensor_from_aos_no_zeros = Tensor.from_aos( + coordinates_no_zeros, values_no_zeros, dimensions=dimensions, format=format + ) assert tensor_from_dok == tensor_from_lol assert tensor_from_dok_no_zeros == tensor_from_lol @@ -238,70 +239,101 @@ def recurse(tree, indexes): assert Tensor.from_scipy_sparse(scipy_coo_matrix) == tensor_from_lol -@pytest.mark.parametrize('lol', [3.5, 0.0, 1]) +@pytest.mark.parametrize("lol", [3.5, 0.0, 1]) def test_convert_0(lol): - assert_all_methods_same(lol, (), '') - - -@pytest.mark.parametrize('lol,dimensions', [ - ([], (0,)), - ([3], (1,)), - ([3, -2, 0], (3,)), -]) -@pytest.mark.parametrize('format', ['d', 's']) + assert_all_methods_same(lol, (), "") + + +@pytest.mark.parametrize( + "lol,dimensions", + [ + ([], (0,)), + ([3], (1,)), + ([3, -2, 0], (3,)), + ], +) +@pytest.mark.parametrize("format", ["d", "s"]) def test_convert_1(lol, dimensions, format): assert_all_methods_same(lol, dimensions, format) -@pytest.mark.parametrize('lol,dimensions', [ - ([], (0, 0)), - ([], (0, 3)), - ([[], []], (2, 0)), - ([[0, -2, 3]], (1, 3)), - ([[0], [3]], (2, 1)), - ([[0, 2, 3], [0, -2, 3]], (2, 3)), - ([[0, 0, 0], [0, 0, 0]], (2, 3)), -]) -@pytest.mark.parametrize('format', ['ss', 'dd', 'sd', 'ds', 's1s0', 'd1d0', 's1d0', 'd1s0']) +@pytest.mark.parametrize( + "lol,dimensions", + [ + ([], (0, 0)), + ([], (0, 3)), + ([[], []], (2, 0)), + ([[0, -2, 3]], (1, 3)), + ([[0], [3]], (2, 1)), + ([[0, 2, 3], [0, -2, 3]], (2, 3)), + ([[0, 0, 0], [0, 0, 0]], (2, 3)), + ], +) +@pytest.mark.parametrize("format", ["ss", "dd", "sd", "ds", "s1s0", "d1d0", "s1d0", "d1s0"]) def test_convert_2(lol, dimensions, format): assert_all_methods_same(lol, dimensions, format) -@pytest.mark.parametrize('lol,dimensions', [ - ([], (0, 0, 0)), - ([], (0, 2, 3)), - ([[], [], [], []], (4, 0, 3)), - ([[[], []], [[], []], [[], []], [[], []]], (4, 2, 0)), - ([[[0, 0, 0], [1, 2, 3]]], (1, 2, 3)), - ([[[0, 0, 0]], [[0, 0, 0]], [[0, 0, 1]], [[1, 0, 0]]], (4, 1, 3)), - ([[[0], [1]], [[0], [-2]], [[0], [0]], [[0], [1]]], (4, 2, 1)), - ([[[0, 0, 0], [1, 2, 3]], [[0, 0, 0], [-2, 3, 0]], [[0, 0, 0], [0, 0, 1]], [[0, 0, 0], [1, 0, 0]]], (4, 2, 3)), - ([[[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]]], (4, 2, 3)), -]) -@pytest.mark.parametrize('modes', ['sss', 'ssd', 'sds', 'sdd', 'dss', 'dsd', 'dds', 'ddd']) -@pytest.mark.parametrize('ordering', ['012', '021', '120', '102', '201', '210']) +@pytest.mark.parametrize( + "lol,dimensions", + [ + ([], (0, 0, 0)), + ([], (0, 2, 3)), + ([[], [], [], []], (4, 0, 3)), + ([[[], []], [[], []], [[], []], [[], []]], (4, 2, 0)), + ([[[0, 0, 0], [1, 2, 3]]], (1, 2, 3)), + ([[[0, 0, 0]], [[0, 0, 0]], [[0, 0, 1]], [[1, 0, 0]]], (4, 1, 3)), + ([[[0], [1]], [[0], [-2]], [[0], [0]], [[0], [1]]], (4, 2, 1)), + ( + [ + [[0, 0, 0], [1, 2, 3]], + [[0, 0, 0], [-2, 3, 0]], + [[0, 0, 0], [0, 0, 1]], + [[0, 0, 0], [1, 0, 0]], + ], + (4, 2, 3), + ), + ( + [ + [[0, 0, 0], [0, 0, 0]], + [[0, 0, 0], [0, 0, 0]], + [[0, 0, 0], [0, 0, 0]], + [[0, 0, 0], [0, 0, 0]], + ], + (4, 2, 3), + ), + ], +) +@pytest.mark.parametrize("modes", ["sss", "ssd", "sds", "sdd", "dss", "dsd", "dds", "ddd"]) +@pytest.mark.parametrize("ordering", ["012", "021", "120", "102", "201", "210"]) def test_convert_3(lol, dimensions, modes, ordering): format = modes[0] + ordering[0] + modes[1] + ordering[1] + modes[2] + ordering[2] assert_all_methods_same(lol, dimensions, format) -@pytest.mark.parametrize('data', [ - 0, - 2.0, -]) +@pytest.mark.parametrize( + "data", + [ + 0, + 2.0, + ], +) def test_interconversion_0(data): expected = Tensor.from_lol(data) - actual = expected.to_format('') + actual = expected.to_format("") assert expected == actual -@pytest.mark.parametrize('data', [ - [0, 0], - [0, 1], -]) -@pytest.mark.parametrize('format1', ['d', 's']) -@pytest.mark.parametrize('format2', ['d', 's']) +@pytest.mark.parametrize( + "data", + [ + [0, 0], + [0, 1], + ], +) +@pytest.mark.parametrize("format1", ["d", "s"]) +@pytest.mark.parametrize("format2", ["d", "s"]) def test_interconversion_1(data, format1, format2): expected = Tensor.from_lol(data) actual = expected.to_format(format2) @@ -309,12 +341,15 @@ def test_interconversion_1(data, format1, format2): assert expected == actual -@pytest.mark.parametrize('data', [ - [[0, 0], [0, 0]], - [[0, 1], [-2, 0]], -]) -@pytest.mark.parametrize('format1', ['ss', 'dd', 'sd', 'ds', 's1s0', 'd1d0', 's1d0', 'd1s0']) -@pytest.mark.parametrize('format2', ['ss', 'dd', 'sd', 'ds', 's1s0', 'd1d0', 's1d0', 'd1s0']) +@pytest.mark.parametrize( + "data", + [ + [[0, 0], [0, 0]], + [[0, 1], [-2, 0]], + ], +) +@pytest.mark.parametrize("format1", ["ss", "dd", "sd", "ds", "s1s0", "d1d0", "s1d0", "d1s0"]) +@pytest.mark.parametrize("format2", ["ss", "dd", "sd", "ds", "s1s0", "d1d0", "s1d0", "d1s0"]) def test_interconversion_2(data, format1, format2): expected = Tensor.from_lol(data) actual = expected.to_format(format2) diff --git a/tests/test_function.py b/tests/test_function.py index 42d68de..565994d 100644 --- a/tests/test_function.py +++ b/tests/test_function.py @@ -1,95 +1,66 @@ +from multiprocessing.pool import ThreadPool from random import randrange -from multiprocessing.pool import ThreadPool -from tensora import Tensor, tensor_method, evaluate +from tensora import Tensor, evaluate, tensor_method def test_csr_matrix_vector_product(): - A = Tensor.from_aos( - [[1, 0], [0, 1], [1, 2]], - [2.0, -2.0, 4.0], - dimensions=(2, 3), format='ds' - ) + A = Tensor.from_aos([[1, 0], [0, 1], [1, 2]], [2.0, -2.0, 4.0], dimensions=(2, 3), format="ds") - x = Tensor.from_aos( - [[0], [1], [2]], - [3.0, 2.5, 2.0], - dimensions=(3,), format='d' - ) + x = Tensor.from_aos([[0], [1], [2]], [3.0, 2.5, 2.0], dimensions=(3,), format="d") - expected = Tensor.from_aos( - [[0], [1]], - [-5.0, 14.0], - dimensions=(2,), format='d' - ) + expected = Tensor.from_aos([[0], [1]], [-5.0, 14.0], dimensions=(2,), format="d") - function = tensor_method('y(i) = A(i,j) * x(j)', dict(A='ds', x='d'), 'd') + function = tensor_method("y(i) = A(i,j) * x(j)", dict(A="ds", x="d"), "d") actual = function(A, x) assert actual == expected - actual = evaluate('y(i) = A(i,j) * x(j)', 'd', A=A, x=x) + actual = evaluate("y(i) = A(i,j) * x(j)", "d", A=A, x=x) assert actual == expected def test_csc_matrix_vector_product(): A = Tensor.from_aos( - [[1, 0], [0, 1], [1, 2]], - [2.0, -2.0, 4.0], - dimensions=(2, 3), format='d1s0' + [[1, 0], [0, 1], [1, 2]], [2.0, -2.0, 4.0], dimensions=(2, 3), format="d1s0" ) - x = Tensor.from_aos( - [[0], [1], [2]], - [3.0, 2.5, 2.0], - dimensions=(3,), format='d' - ) + x = Tensor.from_aos([[0], [1], [2]], [3.0, 2.5, 2.0], dimensions=(3,), format="d") - expected = Tensor.from_aos( - [[0], [1]], - [-5.0, 14.0], - dimensions=(2,), format='d' - ) + expected = Tensor.from_aos([[0], [1]], [-5.0, 14.0], dimensions=(2,), format="d") - function = tensor_method('y(i) = A(i,j) * x(j)', dict(A='d1s0', x='d'), 'd') + function = tensor_method("y(i) = A(i,j) * x(j)", dict(A="d1s0", x="d"), "d") actual = function(A, x) assert actual == expected - actual = evaluate('y(i) = A(i,j) * x(j)', 'd', A=A, x=x) + actual = evaluate("y(i) = A(i,j) * x(j)", "d", A=A, x=x) assert actual == expected def test_csr_matrix_plus_csr_matrix(): - A = Tensor.from_aos( - [[1, 0], [0, 1], [1, 2]], - [2.0, -2.0, 4.0], - dimensions=(2, 3), format='ds' - ) + A = Tensor.from_aos([[1, 0], [0, 1], [1, 2]], [2.0, -2.0, 4.0], dimensions=(2, 3), format="ds") - B = Tensor.from_aos( - [[1, 1], [1, 2], [0, 2]], - [-3.0, 4.0, 3.5], - dimensions=(2, 3), format='ds' - ) + B = Tensor.from_aos([[1, 1], [1, 2], [0, 2]], [-3.0, 4.0, 3.5], dimensions=(2, 3), format="ds") expected = Tensor.from_aos( [[1, 0], [0, 1], [1, 2], [1, 1], [0, 2]], [2.0, -2.0, 8.0, -3.0, 3.5], - dimensions=(2, 3), format='ds' + dimensions=(2, 3), + format="ds", ) - function = tensor_method('C(i,j) = A(i,j) + B(i,j)', dict(A='ds', B='ds'), 'ds') + function = tensor_method("C(i,j) = A(i,j) + B(i,j)", dict(A="ds", B="ds"), "ds") actual = function(A, B) assert actual == expected - actual = evaluate('C(i,j) = A(i,j) + B(i,j)', 'ds', A=A, B=B) + actual = evaluate("C(i,j) = A(i,j) + B(i,j)", "ds", A=A, B=B) assert actual == expected @@ -97,21 +68,13 @@ def test_csr_matrix_plus_csr_matrix(): def test_multithread_evaluation(): # As of version 1.14.4 of cffi, the FFI.compile method is not thread safe. This tests that evaluation of different # kernels is thread safe. - A = Tensor.from_aos( - [[1, 0], [0, 1], [1, 2]], - [2.0, -2.0, 4.0], - dimensions=(2, 3), format='ds' - ) + A = Tensor.from_aos([[1, 0], [0, 1], [1, 2]], [2.0, -2.0, 4.0], dimensions=(2, 3), format="ds") - x = Tensor.from_aos( - [[0], [1], [2]], - [3.0, 2.5, 2.0], - dimensions=(3,), format='d' - ) + x = Tensor.from_aos([[0], [1], [2]], [3.0, 2.5, 2.0], dimensions=(3,), format="d") def run_eval(): # Generate a random expression so that the cache cannot be hit - return evaluate(f'y{randrange(1024)}(i) = A(i,j) * x(j)', 'd', A=A, x=x) + return evaluate(f"y{randrange(1024)}(i) = A(i,j) * x(j)", "d", A=A, x=x) n = 4 with ThreadPool(n) as p: diff --git a/tests/test_numpy.py b/tests/test_numpy.py index 99b1379..7347f9b 100644 --- a/tests/test_numpy.py +++ b/tests/test_numpy.py @@ -1,10 +1,10 @@ -from tensora import Tensor - import pytest +from tensora import Tensor + @pytest.mark.parametrize( - 'array', + "array", [ 0.0, 4.5, @@ -15,9 +15,9 @@ [[[0, 0, 3], [4, 5, 0]], [[0, 0, 0], [4, 5, 6]]], ], ) -@pytest.mark.parametrize('format', ['d', 's']) +@pytest.mark.parametrize("format", ["d", "s"]) def test_to_from_numpy(array, format): - numpy = pytest.importorskip('numpy') + numpy = pytest.importorskip("numpy") expected = numpy.array(array) diff --git a/tests/test_operators.py b/tests/test_operators.py index b00f55f..0d3b5fe 100644 --- a/tests/test_operators.py +++ b/tests/test_operators.py @@ -6,18 +6,14 @@ @pytest.fixture() def a_ds(): return Tensor.from_aos( - [[1, 0], [0, 1], [1, 2]], - [2.0, -2.0, 4.0], - dimensions=(2, 3), format='ds' + [[1, 0], [0, 1], [1, 2]], [2.0, -2.0, 4.0], dimensions=(2, 3), format="ds" ) @pytest.fixture() def b_ds(): return Tensor.from_aos( - [[1, 1], [1, 2], [0, 2]], - [-3.0, 4.0, 3.5], - dimensions=(2, 3), format='ds' + [[1, 1], [1, 2], [0, 2]], [-3.0, 4.0, 3.5], dimensions=(2, 3), format="ds" ) @@ -26,7 +22,8 @@ def c_ds_add(): return Tensor.from_aos( [[1, 0], [0, 1], [1, 2], [1, 1], [0, 2]], [2.0, -2.0, 8.0, -3.0, 3.5], - dimensions=(2, 3), format='ds' + dimensions=(2, 3), + format="ds", ) diff --git a/tests/test_tensor.py b/tests/test_tensor.py index e59ec18..4308569 100644 --- a/tests/test_tensor.py +++ b/tests/test_tensor.py @@ -2,23 +2,23 @@ import pytest -from tensora import Mode, Format, Tensor +from tensora import Format, Mode, Tensor # The example tensors from the original taco paper: # http://tensor-compiler.org/kjolstad-oopsla17-tensor-compiler.pdf figure_5_taco_data = [ - ('d0d1', [[], []], [6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7]), - ('s0d1', [[[0, 2], [0, 2]], []], [6, 0, 9, 8, 5, 0, 0, 7]), - ('d0s1', [[], [[0, 3, 3, 5], [0, 2, 3, 0, 3]]], [6, 9, 8, 5, 7]), - ('s0s1', [[[0, 2], [0, 2]], [[0, 3, 5], [0, 2, 3, 0, 3]]], [6, 9, 8, 5, 7]), - ('d1d0', [[], []], [6, 0, 5, 0, 0, 0, 9, 0, 0, 8, 0, 7]), - ('s1d0', [[[0, 3], [0, 2, 3]], []], [6, 0, 5, 9, 0, 0, 8, 0, 7]), - ('d1s0', [[], [[0, 2, 2, 3, 5], [0, 2, 0, 0, 2]]], [6, 5, 9, 8, 7]), - ('s1s0', [[[0, 3], [0, 2, 3]], [[0, 2, 3, 5], [0, 2, 0, 0, 2]]], [6, 5, 9, 8, 7]), + ("d0d1", [[], []], [6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7]), + ("s0d1", [[[0, 2], [0, 2]], []], [6, 0, 9, 8, 5, 0, 0, 7]), + ("d0s1", [[], [[0, 3, 3, 5], [0, 2, 3, 0, 3]]], [6, 9, 8, 5, 7]), + ("s0s1", [[[0, 2], [0, 2]], [[0, 3, 5], [0, 2, 3, 0, 3]]], [6, 9, 8, 5, 7]), + ("d1d0", [[], []], [6, 0, 5, 0, 0, 0, 9, 0, 0, 8, 0, 7]), + ("s1d0", [[[0, 3], [0, 2, 3]], []], [6, 0, 5, 9, 0, 0, 8, 0, 7]), + ("d1s0", [[], [[0, 2, 2, 3, 5], [0, 2, 0, 0, 2]]], [6, 5, 9, 8, 7]), + ("s1s0", [[[0, 3], [0, 2, 3]], [[0, 2, 3, 5], [0, 2, 0, 0, 2]]], [6, 5, 9, 8, 7]), ] -@pytest.mark.parametrize('format,indices,vals', figure_5_taco_data) +@pytest.mark.parametrize("format,indices,vals", figure_5_taco_data) def test_figure_5(format, indices, vals): data = [[6, 0, 9, 8], [0, 0, 0, 0], [5, 0, 0, 7]] A = Tensor.from_lol(data, dimensions=(3, 4), format=format) @@ -27,13 +27,16 @@ def test_figure_5(format, indices, vals): assert A.taco_vals == vals -@pytest.mark.parametrize('format,indices,vals', figure_5_taco_data) -@pytest.mark.parametrize('permutation', [ - [0, 1, 2, 3, 4], - [4, 3, 2, 1, 0], - [0, 4, 3, 2, 1], - [2, 1, 3, 4, 0], -]) +@pytest.mark.parametrize("format,indices,vals", figure_5_taco_data) +@pytest.mark.parametrize( + "permutation", + [ + [0, 1, 2, 3, 4], + [4, 3, 2, 1, 0], + [0, 4, 3, 2, 1], + [2, 1, 3, 4, 0], + ], +) def test_unsorted_coordinates(format, indices, vals, permutation): data = [ ((0, 0), 6), @@ -64,10 +67,10 @@ def test_unsorted_coordinates(format, indices, vals, permutation): # The example tensors from the follow-up taco paper: # https://tensor-compiler.org/chou-oopsla18-taco-formats.pdf -@pytest.mark.parametrize('format,indices,vals', [ - ('d', [[]], [5, 1, 0, 0, 2, 0, 8, 0]), - ('s', [[[0, 4], [0, 1, 4, 6]]], [5, 1, 2, 8]) -]) +@pytest.mark.parametrize( + "format,indices,vals", + [("d", [[]], [5, 1, 0, 0, 2, 0, 8, 0]), ("s", [[[0, 4], [0, 1, 4, 6]]], [5, 1, 2, 8])], +) def test_figure_2a(format, indices, vals): data = [5, 1, 0, 0, 2, 0, 8, 0] a = Tensor.from_lol(data, dimensions=(8,), format=format) @@ -76,10 +79,17 @@ def test_figure_2a(format, indices, vals): assert a.taco_vals == vals -@pytest.mark.parametrize('format,indices,vals', [ - ('ds', [[], [[0, 2, 4, 4, 7], [0, 1, 0, 1, 0, 3, 4]]], [5, 1, 7, 3, 8, 4, 9]), - ('ss', [[[0, 3], [0, 1, 3]], [[0, 2, 4, 7], [0, 1, 0, 1, 0, 3, 4]]], [5, 1, 7, 3, 8, 4, 9]) -]) +@pytest.mark.parametrize( + "format,indices,vals", + [ + ("ds", [[], [[0, 2, 4, 4, 7], [0, 1, 0, 1, 0, 3, 4]]], [5, 1, 7, 3, 8, 4, 9]), + ( + "ss", + [[[0, 3], [0, 1, 3]], [[0, 2, 4, 7], [0, 1, 0, 1, 0, 3, 4]]], + [5, 1, 7, 3, 8, 4, 9], + ), + ], +) def test_figure_2e(format, indices, vals): data = { (0, 0): 5, @@ -96,10 +106,20 @@ def test_figure_2e(format, indices, vals): assert a.taco_vals == vals -@pytest.mark.parametrize('format,indices,vals', [ - ('sss', [[[0, 2], [0, 2]], [[0, 2, 5], [0, 2, 0, 2, 3]], [[0, 2, 3, 4, 6, 8], [0, 1, 1, 1, 0, 1, 0, 1]]], - [1, 7, 5, 2, 4, 8, 3, 9]), -]) +@pytest.mark.parametrize( + "format,indices,vals", + [ + ( + "sss", + [ + [[0, 2], [0, 2]], + [[0, 2, 5], [0, 2, 0, 2, 3]], + [[0, 2, 3, 4, 6, 8], [0, 1, 1, 1, 0, 1, 0, 1]], + ], + [1, 7, 5, 2, 4, 8, 3, 9], + ), + ], +) def test_figure_2m(format, indices, vals): data = { (0, 0, 0): 1, @@ -125,7 +145,7 @@ def test_from_dok(): (2, 3): 5.0, } format = Format((Mode.compressed, Mode.compressed), (0, 1)) - x = Tensor.from_dok(data, dimensions=(4, 5), format='ss') + x = Tensor.from_dok(data, dimensions=(4, 5), format="ss") assert x.order == 2 assert x.dimensions == (4, 5) @@ -248,11 +268,14 @@ def test_nonscalar_to_float(): _ = float(x) -@pytest.mark.parametrize('a,b,c', [ - ([0, 0, 1], 3, [3, 3, 4]), - (3, [0, 0, 1], [3, 3, 4]), - ([1, 2, 3], [0, 0, 1], [1, 2, 4]), -]) +@pytest.mark.parametrize( + "a,b,c", + [ + ([0, 0, 1], 3, [3, 3, 4]), + (3, [0, 0, 1], [3, 3, 4]), + ([1, 2, 3], [0, 0, 1], [1, 2, 4]), + ], +) def test_add(a, b, c): if isinstance(a, list): a = Tensor.from_lol(a) @@ -264,11 +287,14 @@ def test_add(a, b, c): assert actual == expected -@pytest.mark.parametrize('a,b,c', [ - ([0, 0, 1], 3, [-3, -3, -2]), - (3, [0, 0, 1], [3, 3, 2]), - ([1, 2, 3], [0, 0, 1], [1, 2, 2]), -]) +@pytest.mark.parametrize( + "a,b,c", + [ + ([0, 0, 1], 3, [-3, -3, -2]), + (3, [0, 0, 1], [3, 3, 2]), + ([1, 2, 3], [0, 0, 1], [1, 2, 2]), + ], +) def test_subtract(a, b, c): if isinstance(a, list): a = Tensor.from_lol(a) @@ -280,11 +306,14 @@ def test_subtract(a, b, c): assert actual == expected -@pytest.mark.parametrize('a,b,c', [ - ([0, 0, 1], 3, [0, 0, 3]), - (3, [0, 0, 2], [0, 0, 6]), - ([1, 2, 3], [0, 0, 4], [0, 0, 12]), -]) +@pytest.mark.parametrize( + "a,b,c", + [ + ([0, 0, 1], 3, [0, 0, 3]), + (3, [0, 0, 2], [0, 0, 6]), + ([1, 2, 3], [0, 0, 4], [0, 0, 12]), + ], +) def test_multiply(a, b, c): if isinstance(a, list): a = Tensor.from_lol(a) @@ -296,12 +325,15 @@ def test_multiply(a, b, c): assert actual == expected -@pytest.mark.parametrize('a,b,c', [ - ([3, 2, 5], [1, 2, 0], 7), - ([0, 0, 1], [[0, 0], [4, -1], [0, 2]], [0, 2]), - ([[0, 0, 1], [2, 2, 3]], [0, -1, 2], [2, 4]), - ([[0, 0, 1], [2, 2, 3]], [[0, 0], [4, -1], [0, 2]], [[0, 2], [8, 4]]), -]) +@pytest.mark.parametrize( + "a,b,c", + [ + ([3, 2, 5], [1, 2, 0], 7), + ([0, 0, 1], [[0, 0], [4, -1], [0, 2]], [0, 2]), + ([[0, 0, 1], [2, 2, 3]], [0, -1, 2], [2, 4]), + ([[0, 0, 1], [2, 2, 3]], [[0, 0], [4, -1], [0, 2]], [[0, 2], [8, 4]]), + ], +) def test_matrix_multiply(a, b, c): a = Tensor.from_lol(a) b = Tensor.from_lol(b) @@ -336,28 +368,36 @@ def test_matrix_multiply_mismatched_dimensions(): def test_matrix_multiply_too_many_dimensions(): a = Tensor.from_lol([3, 2, 5]) - b = Tensor.from_dok({ - (0, 0, 0): 4.5, - (1, 0, 1): 3.2, - (1, 1, 2): -3.0, - (0, 1, 1): 5.0, - }, dimensions=(3, 3, 3)) + b = Tensor.from_dok( + { + (0, 0, 0): 4.5, + (1, 0, 1): 3.2, + (1, 1, 2): -3.0, + (0, 1, 1): 5.0, + }, + dimensions=(3, 3, 3), + ) with pytest.raises(ValueError): _ = a @ b def test_equality(): - assert Tensor.from_lol([0, 2, 0], format='d') == Tensor.from_lol([0, 2, 0], format='s') - assert Tensor.from_lol([0, 1, 0], format='d') != Tensor.from_lol([0, 2, 0], format='s') - assert Tensor.from_lol([0, 1, 0], format='d') != 1 - - -@pytest.mark.parametrize('tensor', [ - Tensor.from_dok({}, dimensions=()), - Tensor.from_dok({(2, 3): 2.0, (0, 1): 0.0, (1, 2): -1.0, (0, 3): 0.0}, dimensions=(2, 4)), - Tensor.from_dok({(0, 0, 0): 4.5, (1, 0, 1): 3.2, (1, 1, 2): -3.0, (0, 1, 1): 5.0}, dimensions=(3, 3, 3)), -]) + assert Tensor.from_lol([0, 2, 0], format="d") == Tensor.from_lol([0, 2, 0], format="s") + assert Tensor.from_lol([0, 1, 0], format="d") != Tensor.from_lol([0, 2, 0], format="s") + assert Tensor.from_lol([0, 1, 0], format="d") != 1 + + +@pytest.mark.parametrize( + "tensor", + [ + Tensor.from_dok({}, dimensions=()), + Tensor.from_dok({(2, 3): 2.0, (0, 1): 0.0, (1, 2): -1.0, (0, 3): 0.0}, dimensions=(2, 4)), + Tensor.from_dok( + {(0, 0, 0): 4.5, (1, 0, 1): 3.2, (1, 1, 2): -3.0, (0, 1, 1): 5.0}, dimensions=(3, 3, 3) + ), + ], +) def test_pickle(tensor): # Ensure that not only are the tensors equal, but that they also have the same format and explicit zeros, neither of # which affects equality